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

feat(server): converge legacy compatibility (#15426)

Browse Source 
#### PR Dependency Tree


* **PR #15426** 👈

This tree was auto-generated by
[Charcoal](https://github.com/danerwilliams/charcoal)

<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->
## Summary by CodeRabbit

* **New Features**
* Added workspace BYOK profiles with provider/model catalogs, capability
validation, connection probing, credential rotation, reordering, and
secure local leases.
* Added Copilot route options, selectable targets, managed tiers,
explicit profile/model overrides, and improved streaming with tool
callbacks and abort support.
* Added Copilot availability controls to prevent access when the feature
is disabled.
* **Changes**
* Simplified Copilot configuration and removed legacy provider-specific
settings.
* Removed obsolete model, token-cost, transcript strategy, and provider
metadata fields from public responses.
<!-- end of auto-generated comment: release notes by coderabbit.ai -->
This commit is contained in:
DarkSky
2026-08-05 00:40:13 +08:00
committed by GitHub
parent fdfb6df826
commit 965f4590ff
272 changed files with 12430 additions and 33824 deletions
Download Patch File Download Diff File Expand all files Collapse all files
packages/backend/server/migrations/20260803095500_converge_copilot_runtime/migration.sql
+58
View File
@@ -0,0 +1,58 @@
DELETE FROM "app_configs"
WHERE "id" IN (
'copilot.providers.openai',
'copilot.providers.cloudflareWorkersAi',
'copilot.providers.fal',
'copilot.providers.gemini',
'copilot.providers.geminiVertex',
'copilot.providers.anthropic',
'copilot.providers.anthropicVertex',
'copilot.providers.defaults'
);
DELETE FROM "ai_workspace_byok_configs";
DO $$
BEGIN
IF to_regclass('public.runtime_states') IS NOT NULL THEN
DELETE FROM "runtime_states"
WHERE "purpose" IN (
'copilot_byok_local_lease',
'copilot_byok_local_lease:active'
);
END IF;
END $$;
ALTER TABLE "ai_workspace_byok_configs"
DROP COLUMN "endpoint",
DROP COLUMN "disabled_reason",
DROP COLUMN "last_validated_at",
DROP COLUMN "last_validation_error",
ADD COLUMN "definition" JSONB NOT NULL,
ADD COLUMN "revision" INTEGER NOT NULL DEFAULT 1,
ADD COLUMN "credential_generation" INTEGER NOT NULL DEFAULT 1,
ADD COLUMN "validation" JSONB;
ALTER TABLE "ai_sessions_metadata"
DROP CONSTRAINT "ai_sessions_metadata_prompt_name_fkey",
DROP COLUMN "tokenCost";
UPDATE "ai_action_runs"
SET "action_id" = 'transcript.audio'
WHERE "action_id" = 'transcript.audio.gemini';
UPDATE "ai_transcript_tasks"
SET
"recipe_id" = 'transcript.audio',
"input_snapshot" = "input_snapshot"::jsonb - 'providerMeta' - 'strategy',
"public_meta" = "public_meta"::jsonb - 'providerMeta' - 'strategy',
"protected_result" = "protected_result"::jsonb - 'providerMeta' - 'strategy'
WHERE "recipe_id" = 'transcript.audio.gemini';
ALTER TABLE "ai_transcript_tasks"
DROP COLUMN "strategy";
DROP TABLE "ai_prompts_messages";
DROP TABLE "ai_prompts_metadata";
ALTER TYPE "AiPromptRole" RENAME TO "AiSessionMessageRole";
packages/backend/server/schema.prisma
+15 -56
View File
@@ -706,61 +706,23 @@ model SnapshotHistory {
@@map("snapshot_histories")
}
enum AiPromptRole {
enum AiSessionMessageRole {
system
assistant
user
}
model AiPromptMessage {
promptId Int @map("prompt_id") @db.Integer
// if a group of prompts contains multiple sentences, idx specifies the order of each sentence
idx Int @db.Integer
// system/assistant/user
role AiPromptRole
// prompt content
content String @db.Text
attachments Json? @db.Json
params Json? @db.Json
createdAt DateTime @default(now()) @map("created_at") @db.Timestamptz(3)
prompt AiPrompt @relation(fields: [promptId], references: [id], onDelete: Cascade)
@@unique([promptId, idx])
@@map("ai_prompts_messages")
}
model AiPrompt {
id Int @id @default(autoincrement()) @db.Integer
name String @unique @db.VarChar(32)
// an mark identifying which view to use to display the session
// it is only used in the frontend and does not affect the backend
action String? @db.VarChar
model String @db.VarChar
optionalModels String[] @default([]) @map("optional_models") @db.VarChar
config Json? @db.Json
createdAt DateTime @default(now()) @map("created_at") @db.Timestamptz(3)
updatedAt DateTime @default(now()) @map("updated_at") @db.Timestamptz(3)
// whether the prompt metadata is manually overridden in compat storage
modified Boolean @default(false)
messages AiPromptMessage[]
sessions AiSession[]
@@map("ai_prompts_metadata")
}
model AiSessionMessage {
id String @id @default(uuid()) @db.VarChar
sessionId String @map("session_id") @db.VarChar
compatSubmissionId String? @map("compat_submission_id") @db.VarChar
role AiPromptRole
content String @db.Text
streamObjects Json? @db.Json
attachments Json? @db.Json
params Json? @db.Json
createdAt DateTime @default(now()) @map("created_at") @db.Timestamptz(3)
updatedAt DateTime @updatedAt @map("updated_at") @db.Timestamptz(3)
id String @id @default(uuid()) @db.VarChar
sessionId String @map("session_id") @db.VarChar
compatSubmissionId String? @map("compat_submission_id") @db.VarChar
role AiSessionMessageRole
content String @db.Text
streamObjects Json? @db.Json
attachments Json? @db.Json
params Json? @db.Json
createdAt DateTime @default(now()) @map("created_at") @db.Timestamptz(3)
updatedAt DateTime @updatedAt @map("updated_at") @db.Timestamptz(3)
session AiSession @relation(fields: [sessionId], references: [id], onDelete: Cascade)
@@ -782,13 +744,11 @@ model AiSession {
// the session id of the parent session if this session is a forked session
parentSessionId String? @map("parent_session_id") @db.VarChar
messageCost Int @default(0)
tokenCost Int @default(0)
createdAt DateTime @default(now()) @map("created_at") @db.Timestamptz(3)
updatedAt DateTime @default(now()) @updatedAt @map("updated_at") @db.Timestamptz(3)
deletedAt DateTime? @map("deleted_at") @db.Timestamptz(3)
user User @relation(fields: [userId], references: [id], onDelete: Cascade)
prompt AiPrompt @relation(fields: [promptName], references: [name], onDelete: Cascade)
messages AiSessionMessage[]
context AiContext[]
actionRuns AiActionRun[]
@@ -843,7 +803,6 @@ model AiTranscriptTask {
workspaceId String @map("workspace_id") @db.VarChar
blobId String @map("blob_id") @db.VarChar
status String @db.VarChar
strategy String @db.VarChar
recipeId String @map("recipe_id") @db.VarChar
recipeVersion String @map("recipe_version") @db.VarChar
actionRunId String? @map("action_run_id") @db.VarChar
@@ -986,12 +945,12 @@ model AiWorkspaceByokConfig {
name String @db.VarChar
description String? @db.VarChar
encryptedApiKey String @map("encrypted_api_key") @db.Text
endpoint String? @db.Text
definition Json @db.JsonB
revision Int @default(1)
credentialGeneration Int @default(1) @map("credential_generation")
validation Json? @db.JsonB
sortOrder Int @default(0) @map("sort_order")
enabled Boolean @default(true)
disabledReason String? @map("disabled_reason") @db.VarChar
lastValidatedAt DateTime? @map("last_validated_at") @db.Timestamptz(3)
lastValidationError String? @map("last_validation_error") @db.Text
lastUsedAt DateTime? @map("last_used_at") @db.Timestamptz(3)
lastErrorAt DateTime? @map("last_error_at") @db.Timestamptz(3)
lastError String? @map("last_error") @db.Text
packages/backend/server/src/__tests__/copilot/__snapshots__/copilot.e2e.ts.md
-167
View File
@@ -1,167 +0,0 @@
# Snapshot report for `src/__tests__/copilot.e2e.ts`
The actual snapshot is saved in `copilot.e2e.ts.snap`.
Generated by [AVA](https://avajs.dev).
## should be able to retry with api
> should be able to list history after retry
[
{
messages: [
{
content: 'generate text to text stream',
role: 'assistant',
},
],
pinned: false,
tokens: 10,
},
]
> should be able to list history after retry
[
{
messages: [
{
content: 'generate text to text stream',
role: 'assistant',
},
],
pinned: false,
tokens: 10,
},
]
## should be able to manage context
> should list context files
[
{
blobId: 'Ip3vuwzubwJnOlzeKQ0Gc-daDcMc7EOYnIqypOyn4bs',
chunkSize: 0,
name: 'sample.pdf',
status: 'processing',
},
]
> should list context docs
[
{
id: 'docId1',
status: 'processing',
},
]
## should be able to transcript
> should submit audio transcription job
[
{
status: 'running',
},
]
> should claim audio transcription job
[
{
actions: 'generate text to text',
status: 'claimed',
summary: 'generate text to text',
title: 'generate text to text',
transcription: [
{
end: '00:00:45',
speaker: 'A',
start: '00:00:30',
transcription: 'Hello, everyone.',
},
{
end: '00:01:10',
speaker: 'B',
start: '00:00:46',
transcription: 'Hi, thank you for joining the meeting today.',
},
],
},
]
> should submit audio transcription job
[
{
status: 'running',
},
]
> should claim audio transcription job
[
{
actions: 'generate text to text',
status: 'claimed',
summary: 'generate text to text',
title: 'generate text to text',
transcription: [
{
end: '00:00:45',
speaker: 'A',
start: '00:00:30',
transcription: 'Hello, everyone.',
},
{
end: '00:01:10',
speaker: 'B',
start: '00:00:46',
transcription: 'Hi, thank you for joining the meeting today.',
},
{
end: '00:10:45',
speaker: 'A',
start: '00:10:30',
transcription: 'Hello, everyone.',
},
{
end: '00:11:10',
speaker: 'B',
start: '00:10:46',
transcription: 'Hi, thank you for joining the meeting today.',
},
],
},
]
## should create different session types and validate prompt constraints
> should create session with should create workspace session with text prompt
[
{
pinned: false,
},
]
> should create session with should create pinned session with text prompt
[
{
docId: 'pinned-doc',
pinned: true,
},
]
> should create session with should create doc session with text prompt
[
{
docId: 'normal-doc',
pinned: false,
},
]
packages/backend/server/src/__tests__/copilot/__snapshots__/copilot.spec.ts.md
-431
View File
@@ -1,431 +0,0 @@
# Snapshot report for `src/__tests__/copilot/copilot.spec.ts`
The actual snapshot is saved in `copilot.spec.ts.snap`.
Generated by [AVA](https://avajs.dev).
## should be able to manage chat session
> should generate the final message
[
{
content: 'hello world',
params: {
word: 'world',
},
role: 'system',
},
{
content: 'hello',
role: 'user',
},
]
> should generate different message with another params
[
{
content: 'hello world',
params: {
word: 'world',
},
role: 'system',
},
{
content: 'hello',
role: 'user',
},
]
## should be able to fork chat session
> should generate the final message
[
{
content: 'hello world',
params: {
word: 'world',
},
role: 'system',
},
{
content: 'hello',
role: 'user',
},
{
content: 'world',
role: 'assistant',
},
]
> should generate the final message
[
{
content: 'hello world',
params: {
word: 'world',
},
role: 'system',
},
{
content: 'hello',
role: 'user',
},
{
content: 'world',
role: 'assistant',
},
]
> should generate the final message
[
{
content: 'hello world',
params: {
word: 'world',
},
role: 'system',
},
{
content: 'hello',
role: 'user',
},
{
content: 'world',
role: 'assistant',
},
{
content: 'aaa',
role: 'user',
},
{
content: 'bbb',
role: 'assistant',
},
]
> should generate the final message
[
{
content: 'hello world',
params: {
word: 'world',
},
role: 'system',
},
{
content: 'hello',
role: 'user',
},
{
content: 'world',
role: 'assistant',
},
{
content: 'aaa',
role: 'user',
},
{
content: 'bbb',
role: 'assistant',
},
]
## should revert message correctly
> should have three messages before revert
[
{
content: 'hello world',
params: {
word: 'world',
},
role: 'system',
},
{
content: '1',
role: 'user',
},
{
content: '2',
role: 'assistant',
},
{
content: '3',
role: 'user',
},
{
content: '4',
role: 'assistant',
},
]
> should remove assistant message after revert
[
{
content: 'hello world',
params: {
word: 'world',
},
role: 'system',
},
{
content: '1',
role: 'user',
},
{
content: '2',
role: 'assistant',
},
{
content: '3',
role: 'user',
},
]
> should remove assistant message after revert
[
{
content: 'hello world',
params: {
word: 'world',
},
role: 'system',
},
{
content: '1',
role: 'user',
},
{
content: '2',
role: 'assistant',
},
]
> should have three messages before revert
[
{
content: 'hello world',
params: {
word: 'world',
},
role: 'system',
},
{
content: '1',
role: 'user',
},
{
content: '2',
role: 'assistant',
},
{
content: '3',
role: 'user',
},
{
content: '4',
role: 'assistant',
},
]
> should remove assistant message after revert
[
{
content: 'hello world',
params: {
word: 'world',
},
role: 'system',
},
{
content: '1',
role: 'user',
},
{
content: '2',
role: 'assistant',
},
{
content: '3',
role: 'user',
},
]
> should remove assistant message after revert
[
{
content: 'hello world',
params: {
word: 'world',
},
role: 'system',
},
{
content: '1',
role: 'user',
},
{
content: '2',
role: 'assistant',
},
]
## should handle generateSessionTitle correctly under various conditions
> should generate title when conditions are met
{
chatWithPromptCalled: undefined,
exists: true,
title: 'What is Machine Learning?',
}
> should not generate title when session already has title
{
chatWithPromptCalled: false,
exists: true,
title: 'Existing Title',
}
> should not generate title when no user messages exist
{
chatWithPromptCalled: false,
exists: true,
title: null,
}
> should not generate title when no assistant messages exist
{
chatWithPromptCalled: false,
exists: true,
title: null,
}
> should use correct prompt for title generation
{
content: `[user]: Explain quantum computing briefly␊
[assistant]: Quantum computing uses quantum mechanics principles.`,
promptName: 'Summary as title',
}
## should handle copilot cron jobs correctly
> daily job scheduling calls
[
{
args: [
'copilot.session.cleanupEmptySessions',
{},
{
jobId: 'daily-copilot-cleanup-empty-sessions',
},
],
},
{
args: [
'copilot.session.generateMissingTitles',
{},
{
jobId: 'daily-copilot-generate-missing-titles',
},
],
},
{
args: [
'copilot.workspace.cleanupTrashedDocEmbeddings',
{},
{
jobId: 'daily-copilot-cleanup-trashed-doc-embeddings',
},
],
},
]
> cleanup empty sessions calls
[
{
args: [
'Date',
],
},
]
> title generation calls
{
jobCalls: [
{
args: [
'copilot.session.generateTitle',
{
sessionId: 'session1',
},
],
},
{
args: [
'copilot.session.generateTitle',
{
sessionId: 'session2',
},
],
},
],
modelCalls: [
{
args: [],
},
],
}
## capability policy host should gate pro model requests by subscription status
> should honor requested pro model
'gpt-5.6-terra'
> should fallback to default model
'gpt-5.6-luna'
> should fallback to default model when requesting pro model during trialing
'gpt-5.6-luna'
> should honor requested non-pro model during trialing
'gpt-5.6-luna'
> should pick default model when no requested model during trialing
'gpt-5.6-luna'
> should pick default model when no requested model during active
'gpt-5.6-luna'
> should honor requested pro model during active
'claude-sonnet-4-6'
> should fallback to default model when requesting non-optional model during active
'gpt-5.6-luna'
packages/backend/server/src/__tests__/copilot/__snapshots__/native-provider.spec.ts.md
-692
View File
@@ -1,692 +0,0 @@
# Snapshot report for `src/__tests__/copilot/native-provider.spec.ts`
The actual snapshot is saved in `native-provider.spec.ts.snap`.
Generated by [AVA](https://avajs.dev).
## NativeProviderAdapter streamObject should map tool and text events
> Snapshot 1
[
{
args: {
doc_id: 'a1',
},
argumentParseError: undefined,
rawArgumentsText: undefined,
thought: undefined,
toolCallId: 'call_1',
toolName: 'doc_read',
type: 'tool-call',
},
{
args: {
doc_id: 'a1',
},
argumentParseError: undefined,
rawArgumentsText: undefined,
result: {
markdown: '# a1',
},
toolCallId: 'call_1',
toolName: 'doc_read',
type: 'tool-result',
},
{
textDelta: 'ok',
type: 'text-delta',
},
]
## buildCanonicalNativeRequest should only use explicit structured contract inputs
> Snapshot 1
{
additionalProperties: false,
properties: {
summary: {
type: 'string',
},
},
required: [
'summary',
],
type: 'object',
}
## buildCanonicalNativeStructuredRequest should accept schema-only explicit structured response contracts
> Snapshot 1
{
schema: {
additionalProperties: false,
properties: {
summary: {
type: 'string',
},
},
required: [
'summary',
],
type: 'object',
},
strict: true,
}
## buildCanonicalNativeStructuredRequest should honor explicit structured options contract before system responseFormat
> Snapshot 1
{
schema: {
additionalProperties: false,
properties: {
ok: {
type: 'boolean',
},
},
required: [
'ok',
],
type: 'object',
},
strict: true,
}
## buildCanonicalNativeStructuredRequest should honor explicit responseSchema for array outputs
> Snapshot 1
{
items: {
additionalProperties: false,
properties: {
speaker: {
type: 'string',
},
text: {
type: 'string',
},
},
required: [
'speaker',
'text',
],
type: 'object',
},
type: 'array',
}
## buildCanonicalNativeStructuredRequest should consume explicit structured response contract without options.schema
> Snapshot 1
{
schema: {
additionalProperties: false,
properties: {
summary: {
type: 'string',
},
},
required: [
'summary',
],
type: 'object',
},
strict: false,
}
## buildCanonicalNativeStructuredRequest should accept explicit schema contracts without schemaHash
> Snapshot 1
{
schema: {
additionalProperties: false,
properties: {
summary: {
type: 'string',
},
},
required: [
'summary',
],
type: 'object',
},
strict: true,
}
## buildNativeRequest should canonicalize Gemini attachments
> remote file url
[
{
text: 'summarize this attachment',
type: 'text',
},
{
source: {
media_type: 'application/pdf',
url: 'https://example.com/a.pdf',
},
type: 'file',
},
]
> remote image url
[
{
text: 'describe this image',
type: 'text',
},
{
source: {
media_type: 'image/png',
url: 'https://example.com/cat.png',
},
type: 'image',
},
]
> data url
[
{
text: 'read this note',
type: 'text',
},
{
source: {
data: 'aGVsbG8gd29ybGQ=',
media_type: 'text/plain',
},
type: 'file',
},
]
> remote audio url
[
{
text: 'transcribe this clip',
type: 'text',
},
{
source: {
media_type: 'audio/mpeg',
url: 'https://example.com/a.mp3',
},
type: 'audio',
},
]
> bytes and file handle
[
{
text: 'inspect these assets',
type: 'text',
},
{
source: {
data: 'aGVsbG8=',
file_name: 'hello.txt',
media_type: 'text/plain',
},
type: 'file',
},
{
source: {
file_handle: 'file_123',
file_name: 'report.pdf',
media_type: 'application/pdf',
},
type: 'file',
},
]
## buildNativeStructuredRequest should prefer explicit schema option
> Snapshot 1
{
additionalProperties: false,
properties: {
summary: {
type: 'string',
},
},
required: [
'summary',
],
type: 'object',
}
## buildNativeStructuredRequest should ignore legacy params.schema fallback when explicit schema contract exists
> Snapshot 1
{
schema: {
additionalProperties: false,
properties: {
summary: {
type: 'string',
},
},
required: [
'summary',
],
type: 'object',
},
strict: true,
}
## defineTool should precompute json schema at definition time
> Snapshot 1
{
additionalProperties: false,
properties: {
docId: {
type: 'string',
},
includeChildren: {
type: 'boolean',
},
},
required: [
'docId',
],
type: 'object',
}
## GeminiProvider should use native path for text-only requests
> Snapshot 1
{
include: [
'reasoning',
],
middleware: {
request: [
'normalize_messages',
'tool_schema_rewrite',
],
stream: [
'stream_event_normalize',
'citation_indexing',
],
},
reasoning: {
effort: 'medium',
},
remoteAttachmentRequests: [],
}
## GeminiProvider should use native path for structured requests
> Snapshot 1
{
request: {
messages: [
{
content: [
{
text: 'Return JSON only.',
type: 'text',
},
],
role: 'system',
},
{
content: [
{
text: 'Summarize AFFiNE in one short sentence.',
type: 'text',
},
],
role: 'user',
},
],
middleware: {
request: [
'normalize_messages',
'tool_schema_rewrite',
],
},
model: 'gemini-3.6-flash',
responseMimeType: 'application/json',
schema: {
additionalProperties: false,
properties: {
summary: {
type: 'string',
},
},
required: [
'summary',
],
type: 'object',
},
strict: true,
},
result: {
summary: 'AFFiNE native',
},
}
## GeminiProvider should use native structured path for audio attachments
> Snapshot 1
{
content: [
{
text: 'transcribe the audio',
type: 'text',
},
{
source: {
data: 'YXVkaW8tYnl0ZXM=',
media_type: 'audio/mpeg',
},
type: 'audio',
},
],
remoteAttachmentRequests: [
'https://example.com/a.mp3',
],
result: [
{
a: 'Speaker 1',
e: 1,
s: 0,
t: 'Hello',
},
],
}
## GeminiProvider should use native path for embeddings
> Snapshot 1
{
request: {
dimensions: 3,
inputs: [
'first',
'second',
],
model: 'gemini-embedding-001',
taskType: 'RETRIEVAL_DOCUMENT',
},
result: [
[
0.1,
0.2,
],
[
1.1,
1.2,
],
],
}
## GeminiProvider should canonicalize native text attachments
> remote file attachment
{
content: [
{
text: 'summarize this file',
type: 'text',
},
{
source: {
data: 'cGRmLWJ5dGVz',
media_type: 'application/pdf',
},
type: 'file',
},
],
remoteAttachmentRequests: [
'https://example.com/a.pdf',
],
}
> remote image attachment
{
content: [
{
text: 'describe this image',
type: 'text',
},
{
source: {
data: 'aW1hZ2UtYnl0ZXM=',
media_type: 'image/jpeg',
},
type: 'image',
},
],
remoteAttachmentRequests: [
'https://example.com/a.jpg',
],
}
> downloaded audio webm attachment
{
content: [
{
text: 'transcribe this clip',
type: 'text',
},
{
source: {
data: 'YXVkaW8tYnl0ZXM=',
media_type: 'audio/webm',
},
type: 'audio',
},
],
remoteAttachmentRequests: [
'https://example.com/a.webm',
],
}
> google file url attachment
{
content: [
{
text: 'summarize this file',
type: 'text',
},
{
source: {
media_type: 'application/pdf',
url: 'https://generativelanguage.googleapis.com/v1beta/files/file-123',
},
type: 'file',
},
],
remoteAttachmentRequests: [],
}
## GeminiVertexProvider should prefetch bearer token for native config
> Snapshot 1
{
auth_token: 'vertex-token',
base_url: 'https://vertex.example',
}
## GeminiVertexProvider should materialize remote attachments before native text path
> remote http url
{
content: [
{
text: 'transcribe the audio',
type: 'text',
},
{
source: {
data: 'YXVkaW8tYnl0ZXM=',
media_type: 'audio/mpeg',
},
type: 'audio',
},
],
remoteAttachmentRequests: [
'https://example.com/a.mp3',
],
}
> gs url
{
content: [
{
text: 'transcribe the audio',
type: 'text',
},
{
source: {
data: 'b3B1cy1ieXRlcw==',
media_type: 'audio/opus',
},
type: 'audio',
},
],
remoteAttachmentRequests: [
'gs://bucket/audio.opus',
],
}
## OpenAIProvider should use native structured dispatch
> Snapshot 1
{
request: {
messages: [
{
content: [
{
text: 'Return JSON only.',
type: 'text',
},
],
role: 'system',
},
{
content: [
{
text: 'Summarize AFFiNE in one sentence.',
type: 'text',
},
],
role: 'user',
},
],
middleware: {
request: [
'normalize_messages',
'tool_schema_rewrite',
],
},
model: 'gpt-4.1',
responseMimeType: 'application/json',
schema: {
additionalProperties: false,
properties: {
summary: {
type: 'string',
},
},
required: [
'summary',
],
type: 'object',
},
strict: true,
},
result: {
summary: 'AFFiNE structured',
},
}
## OpenAIProvider should prefer native output_json for structured dispatch
> Snapshot 1
{
summary: 'AFFiNE structured',
}
## OpenAIProvider should use native embedding dispatch
> Snapshot 1
{
request: {
dimensions: 8,
inputs: [
'alpha',
'beta',
],
model: 'text-embedding-3-small',
taskType: 'RETRIEVAL_DOCUMENT',
},
result: [
[
0.4,
0.5,
],
[
0.4,
0.5,
],
],
}
## OpenAIProvider should use native rerank dispatch
> Snapshot 1
{
request: {
candidates: [
{
id: 'react',
text: 'React is a UI library.',
},
{
id: 'weather',
text: 'The park is sunny today.',
},
],
model: 'gpt-4.1',
query: 'programming',
},
scores: [
0.8,
0.8,
],
}
packages/backend/server/src/__tests__/copilot/__snapshots__/provider-native.spec.ts.md
-505
View File
@@ -1,505 +0,0 @@
# Snapshot report for `src/__tests__/copilot/provider-native.spec.ts`
The actual snapshot is saved in `provider-native.spec.ts.snap`.
Generated by [AVA](https://avajs.dev).
## CopilotProviderFactory should return no prepared routes when native prepare returns null
> Snapshot 1
{
chat: [
length: 0,
prepared: undefined,
providerId: undefined,
],
embedding: [
length: 0,
prepared: undefined,
],
rerank: [
length: 0,
prepared: undefined,
],
structured: [
length: 0,
prepared: undefined,
],
}
## getActiveProviderMiddleware should merge defaults with profile override
> Snapshot 1
{
node: {
text: [
'citation_footnote',
'callout',
'thinking_format',
],
},
rust: {
request: [
'clamp_max_tokens',
],
stream: undefined,
},
}
## checkParams should infer remote image capability from url extension without host mime inference
> Snapshot 1
{
attachmentKinds: [
'image',
],
attachmentSourceKinds: [
'url',
],
inputTypes: [
'image',
'text',
],
}
## llmResolveRequestedModelMatch should preserve provider-prefixed optional matches
> prefixed optional hit
{
matchedOptionalModel: true,
selectedModel: 'openai-default/gpt-5.6-terra',
}
> prefixed optional miss
{
matchedOptionalModel: false,
selectedModel: 'gpt-5.6-luna',
}
## ExecutionPlan should serialize routed request state and reject host-only signal
> Snapshot 1
{
fallbackOrder: [
'openai-main',
],
transport: {
kind: 'chat',
request: {
messages: [
{
content: [
{
text: 'hello',
type: 'text',
},
],
role: 'user',
},
],
model: 'gpt-5-mini',
},
},
}
## NativeExecutionEngine should dispatch prepared text routes through native fallback
> Snapshot 1
[
{
model: 'gpt-5-mini',
providerId: 'openai-primary',
requestShape: {
candidateCount: 0,
firstContent: 'hello from primary',
inputCount: 0,
keys: [
'messages',
'model',
],
query: undefined,
schemaKeys: undefined,
toolNames: [],
},
},
{
model: 'gpt-5-mini',
providerId: 'openai-fallback',
requestShape: {
candidateCount: 0,
firstContent: 'hello from fallback',
inputCount: 0,
keys: [
'messages',
'model',
],
query: undefined,
schemaKeys: undefined,
toolNames: [],
},
},
]
## NativeExecutionEngine should prefer prepared native fallback dispatch for explicit routes
> Snapshot 1
[
{
model: 'gpt-5-mini',
providerId: 'openai-primary',
requestShape: {
candidateCount: 0,
firstContent: 'hello',
inputCount: 0,
keys: [
'messages',
'model',
],
query: undefined,
schemaKeys: undefined,
toolNames: [],
},
},
{
model: 'gpt-5-mini',
providerId: 'openai-fallback',
requestShape: {
candidateCount: 0,
firstContent: 'hello',
inputCount: 0,
keys: [
'messages',
'model',
],
query: undefined,
schemaKeys: undefined,
toolNames: [],
},
},
]
## ExecutionPlanBuilder should keep tool-loop chat routes on prepared dispatch path
> Snapshot 1
{
preparedTools: [
'answer',
],
transport: undefined,
}
## ExecutionPlanBuilder should keep single-route tool chat plans on prepared_routes path
> Snapshot 1
{
kind: 'chat',
request: {
messages: [
{
content: [
{
text: 'hello',
type: 'text',
},
],
role: 'user',
},
],
model: 'gpt-5-mini',
tools: [
{
description: 'Answer',
name: 'answer',
parameters: {
properties: {
value: {
type: 'string',
},
},
required: [
'value',
],
type: 'object',
},
},
],
},
}
## NativeExecutionEngine should route tool-loop chat prepared routes through native dispatch
> Snapshot 1
[
{
model: 'gpt-5-mini',
providerId: 'openai-primary',
requestShape: {
candidateCount: 0,
firstContent: 'hello',
inputCount: 0,
keys: [
'messages',
'model',
'tools',
],
query: undefined,
schemaKeys: undefined,
toolNames: [
'answer',
],
},
},
{
model: 'gpt-5-mini',
providerId: 'openai-fallback',
requestShape: {
candidateCount: 0,
firstContent: 'hello from fallback',
inputCount: 0,
keys: [
'messages',
'model',
'tools',
],
query: undefined,
schemaKeys: undefined,
toolNames: [
'answer',
],
},
},
]
## ExecutionPlanBuilder should build native prepared routes for structured, image, embedding and rerank
> Snapshot 1
{
embedding: {
routes: 2,
transport: undefined,
},
image: {
prepared: {
request: {
images: [],
model: 'gpt-image-1',
operation: 'generate',
prompt: 'draw a cat',
},
route: {
backendConfig: {
auth_token: 'image-key',
base_url: 'https://api.openai.com',
},
model: 'gpt-image-1',
protocol: 'openai_images',
providerId: 'openai-default',
},
},
routes: [
{
config: {
auth_token: 'image-key',
base_url: 'https://api.openai.com',
},
model: 'gpt-image-1',
protocol: 'openai_images',
provider_id: 'openai-default',
request: {
images: [],
model: 'gpt-image-1',
operation: 'generate',
prompt: 'draw a cat',
},
},
],
},
rerank: {
routes: 2,
transport: undefined,
},
structured: {
routes: 2,
transport: undefined,
},
}
## NativeExecutionEngine should dispatch structured prepared routes through native execution
> Snapshot 1
[
{
model: 'gpt-5-mini',
providerId: 'openai-primary',
requestShape: {
candidateCount: 0,
firstContent: 'hello',
inputCount: 0,
keys: [
'messages',
'model',
'schema',
],
query: undefined,
schemaKeys: [
'ok',
],
toolNames: [],
},
},
{
model: 'gpt-5-mini',
providerId: 'openai-fallback',
requestShape: {
candidateCount: 0,
firstContent: 'hello from fallback',
inputCount: 0,
keys: [
'messages',
'model',
'schema',
],
query: undefined,
schemaKeys: [
'ok',
],
toolNames: [],
},
},
]
## NativeExecutionEngine should dispatch embedding prepared routes through native execution
> Snapshot 1
{
called: true,
result: [
[
0.1,
0.2,
],
],
routes: [
{
model: 'text-embedding-3-small',
providerId: 'openai-primary',
requestShape: {
candidateCount: 0,
firstContent: null,
inputCount: 1,
keys: [
'inputs',
'model',
],
query: undefined,
schemaKeys: undefined,
toolNames: [],
},
},
{
model: 'text-embedding-3-small',
providerId: 'openai-fallback',
requestShape: {
candidateCount: 0,
firstContent: null,
inputCount: 1,
keys: [
'inputs',
'model',
],
query: undefined,
schemaKeys: undefined,
toolNames: [],
},
},
],
}
## NativeExecutionEngine should dispatch rerank prepared routes through native execution
> Snapshot 1
{
called: true,
result: [
0.9,
0.1,
],
routes: [
{
model: 'gpt-4o-mini',
providerId: 'openai-primary',
requestShape: {
candidateCount: 1,
firstContent: null,
inputCount: 0,
keys: [
'candidates',
'model',
'query',
],
query: 'programming',
schemaKeys: undefined,
toolNames: [],
},
},
{
model: 'gpt-4o-mini',
providerId: 'openai-fallback',
requestShape: {
candidateCount: 1,
firstContent: null,
inputCount: 0,
keys: [
'candidates',
'model',
'query',
],
query: 'programming fallback',
schemaKeys: undefined,
toolNames: [],
},
},
],
}
## NativeExecutionEngine should dispatch image plans through prepared native routes
> Snapshot 1
[
{
model: 'gpt-image-1',
providerId: 'openai-image',
requestShape: {
candidateCount: 0,
firstContent: null,
imageCount: 0,
inputCount: 0,
keys: [
'images',
'model',
'operation',
'prompt',
],
prompt: 'draw a cat',
query: undefined,
schemaKeys: undefined,
toolNames: [],
},
},
]
packages/backend/server/src/__tests__/copilot/byok-probe.spec.ts
-53
View File
@@ -1,53 +0,0 @@
import test from 'ava';
import Sinon from 'sinon';
import type { safeFetch } from '../../base';
import {
PROVIDER_PROBE_MAX_BYTES,
runProviderProbe,
} from '../../plugins/copilot/byok/probe';
import { ByokProvider } from '../../plugins/copilot/byok/types';
test('provider probe allows model responses and explicitly configured private targets', async t => {
const fetch = Sinon.stub<
Parameters<typeof safeFetch>,
ReturnType<typeof safeFetch>
>().resolves(new Response('{}', { status: 200 }));
await runProviderProbe(
fetch,
ByokProvider.openai,
'secret',
'http://provider.internal/v1',
true
);
t.is(fetch.firstCall.args[0], 'http://provider.internal/v1/models');
t.deepEqual(fetch.firstCall.args[2], {
timeoutMs: 10_000,
maxRedirects: 3,
maxBytes: PROVIDER_PROBE_MAX_BYTES,
allowedHeaders: ['Authorization'],
allowHttp: true,
allowPrivateTargetOrigin: true,
});
t.true(PROVIDER_PROBE_MAX_BYTES >= 64 * 1024);
});
test('provider probe keeps private targets blocked by default', async t => {
const fetch = Sinon.stub<
Parameters<typeof safeFetch>,
ReturnType<typeof safeFetch>
>().resolves(new Response('{}', { status: 200 }));
await runProviderProbe(
fetch,
ByokProvider.gemini,
'secret',
'https://provider.example/v1beta',
false
);
t.false(fetch.firstCall.args[2]?.allowHttp);
t.false(fetch.firstCall.args[2]?.allowPrivateTargetOrigin);
});
packages/backend/server/src/__tests__/copilot/byok.spec.ts
+188 -1219
View File
File diff suppressed because it is too large Load Diff
packages/backend/server/src/__tests__/copilot/capability-runtime.spec.ts
+320
View File
@@ -0,0 +1,320 @@
import ava from 'ava';
import { z } from 'zod';
import { Config, CopilotQuotaExceeded } from '../../base';
import type { BackendRuntimeProvider } from '../../core/backend-runtime';
import type { Models } from '../../models';
import type { ByokEntitlementPolicy } from '../../plugins/copilot/byok';
import type { ConversationPolicy } from '../../plugins/copilot/conversation/policy';
import { CapabilityRuntime } from '../../plugins/copilot/runtime/capability-runtime';
import { CopilotRuntimeEventConsumer } from '../../plugins/copilot/runtime/copilot-runtime-event-consumer';
import { executeToolCall } from '../../plugins/copilot/runtime/tool/bridge';
import type { ToolRuntime } from '../../plugins/copilot/runtime/tool-runtime';
const test = ava;
async function collect<T>(source: AsyncIterable<T>) {
const values: T[] = [];
for await (const value of source) values.push(value);
return values;
}
function runtimeFixture(streamError?: string, enabled = true) {
const calls: Array<{
slot: string;
request?: unknown;
targetOverride?: { profileId: string; modelId: string };
}> = [];
const backend = {
executeCopilot: async (input: {
slot: string;
request: unknown;
targetOverride?: { profileId: string; modelId: string };
}) => {
calls.push(input);
if (input.slot === 'index.embedding')
return { events: [], result: { embeddings: [[1, 2]] } };
if (input.slot === 'search.rerank')
return { events: [], result: { scores: [0.9] } };
if (
input.slot === 'image.generate' ||
input.slot === 'action.image.filter.sketch'
) {
return {
events: [],
result: {
images: [
{ url: 'https://example.com/image.png', media_type: 'image/png' },
],
},
};
}
throw new Error(`unexpected slot ${input.slot}`);
},
streamCopilot: (input: {
slot: string;
request: unknown;
targetOverride?: { profileId: string; modelId: string };
}) => {
calls.push(input);
async function* events() {
if (streamError) {
yield { type: 'error', message: streamError };
return;
}
yield { type: 'message_start', model: 'opaque/model' };
yield { type: 'text_delta', text: 'hello' };
yield { type: 'done', finish_reason: 'stop' };
}
return events();
},
assertCopilotRoute: async () => {
throw new Error('access_unavailable');
},
} as unknown as BackendRuntimeProvider;
const entitlement = {
hasServerEntitlement: async () => true,
hasLocalEntitlement: async () => true,
hasAiPlan: async () => false,
} as unknown as ByokEntitlementPolicy;
const conversation = {
hasQuota: async () => true,
} as unknown as ConversationPolicy;
const tools = { getTools: async () => ({}) } as unknown as ToolRuntime;
const consumer = {
consume: async () => {},
} as unknown as CopilotRuntimeEventConsumer;
const config = { copilot: { enabled } } as Config;
return {
calls,
runtime: new CapabilityRuntime(
backend,
entitlement,
conversation,
tools,
consumer,
config
),
};
}
test('disabled copilot rejects native execution before route access', async t => {
const { runtime, calls } = runtimeFixture(undefined, false);
t.false(await runtime.embeddingConfigured('ignored'));
await t.throwsAsync(runtime.embed('ignored', ['text']), {
message: 'Copilot is disabled.',
});
t.deepEqual(calls, []);
});
test('all operation kinds enter the native slot pipeline', async t => {
const { runtime, calls } = runtimeFixture();
t.deepEqual(await runtime.embed('ignored', ['text']), [[1, 2]]);
t.deepEqual(
await runtime.rerank('ignored', {
query: 'query',
candidates: [{ id: 'one', text: 'text' }],
}),
[0.9]
);
t.deepEqual(
await collect(
runtime.streamImageArtifacts(
{},
[{ role: 'user', content: 'draw' }],
{},
undefined,
'action.image.filter.sketch'
)
),
[{ url: 'https://example.com/image.png', media_type: 'image/png' }]
);
t.deepEqual(
calls.map(call => call.slot),
['index.embedding', 'search.rerank', 'action.image.filter.sketch']
);
});
test('image request builder receives only serializable request options', async t => {
const { runtime, calls } = runtimeFixture();
const controller = new AbortController();
await collect(
runtime.streamImageArtifacts({}, [{ role: 'user', content: 'draw' }], {
quality: 'high',
seed: 42,
signal: controller.signal,
user: 'user-1',
})
);
t.deepEqual(calls[0].request, {
model: 'route-selected',
prompt: 'draw',
operation: 'generate',
options: {
quality: 'high',
outputFormat: 'webp',
seed: 42,
},
});
});
test('text streaming consumes native generic events', async t => {
const { runtime, calls } = runtimeFixture();
const chunks = await collect(
runtime.streamText({ profileId: 'profile-1', modelId: 'vendor/model:B' }, [
{ role: 'user', content: 'hello' },
])
);
t.is(chunks.join(''), 'hello');
t.is(calls[0].slot, 'chat.default');
t.deepEqual(calls[0].targetOverride, {
profileId: 'profile-1',
modelId: 'vendor/model:B',
});
await t.throwsAsync(runtime.assertRoute('chat.default', {}, {}), {
instanceOf: CopilotQuotaExceeded,
});
const denied = runtimeFixture('access_unavailable').runtime;
await t.throwsAsync(
async () =>
await collect(
denied.streamText({}, [{ role: 'user', content: 'denied' }])
),
{ instanceOf: CopilotQuotaExceeded }
);
});
test('product event consumer attributes BYOK usage from structured identity', async t => {
const records: unknown[] = [];
const models = {
copilotUsage: { create: async (value: unknown) => records.push(value) },
copilotWorkspaceByokConfig: {
touchUsed: async () => {},
markFailure: async () => {},
},
} as unknown as Models;
const consumer = new CopilotRuntimeEventConsumer(models);
await consumer.consume(
[
{
type: 'usage',
route: {
profileId: 'profile-1',
source: 'server',
provider: 'openai',
model: 'opaque/model:B',
},
usage: { input_tokens: 3, output_tokens: 2, total_tokens: 5 },
},
],
{ workspaceId: 'workspace-1', featureKind: 'chat' }
);
t.like(records[0], {
workspaceId: 'workspace-1',
provider: 'openai',
providerSource: 'byok_server',
model: 'opaque/model:B',
promptTokens: 3,
completionTokens: 2,
totalTokens: 5,
});
});
test('tool callback validates arguments and preserves call identity', async t => {
const result = await executeToolCall(
{
echo: {
description: 'echo',
inputSchema: z.object({ value: z.string() }),
execute: async ({ value }) => ({ value }),
},
},
{ callId: 'call-1', name: 'echo', args: { value: 'ok' } },
{}
);
t.deepEqual(result, {
callId: 'call-1',
name: 'echo',
args: { value: 'ok' },
rawArgumentsText: undefined,
argumentParseError: undefined,
output: { value: 'ok' },
});
});
test('tool callback reports missing tools and invalid argument JSON', async t => {
const missing = await executeToolCall(
{},
{ callId: 'call-1', name: 'missing', args: {} },
{}
);
const invalid = await executeToolCall(
{},
{
callId: 'call-2',
name: 'missing',
args: {},
rawArgumentsText: '{',
argumentParseError: 'unexpected end',
},
{}
);
t.true(missing.isError);
t.true(invalid.isError);
t.deepEqual(invalid.output, {
message: 'Invalid tool arguments JSON',
rawArguments: '{',
error: 'unexpected end',
});
});
test('tool callback rejects invalid zod args without execution', async t => {
let executed = false;
const result = await executeToolCall(
{
echo: {
description: 'echo',
inputSchema: z.object({ value: z.string().trim() }),
execute: async () => {
executed = true;
},
},
},
{
callId: 'call-1',
name: 'echo',
args: { value: 42 },
rawArgumentsText: '{"value":42}',
},
{}
);
t.true(result.isError);
t.false(executed);
});
test('tool callback passes transformed args without prototype pollution', async t => {
const received: unknown[] = [];
const result = await executeToolCall(
{
echo: {
description: 'echo',
inputSchema: z.object({ value: z.string().trim() }).passthrough(),
execute: async args => received.push(args),
},
},
{
callId: 'call-1',
name: 'echo',
args: JSON.parse('{"value":" AFFiNE ","__proto__":{"polluted":true}}'),
},
{}
);
t.false(result.isError ?? false);
t.deepEqual(received, [{ value: 'AFFiNE' }]);
t.is((Object.prototype as Record<string, unknown>).polluted, undefined);
});
packages/backend/server/src/__tests__/copilot/utils.spec.ts → packages/backend/server/src/__tests__/copilot/citation-formatter.spec.ts
View File
packages/backend/server/src/__tests__/copilot/conversation-host.spec.ts
+193
View File
@@ -0,0 +1,193 @@
import '../../plugins/copilot/runtime/capability-runtime';
import ava from 'ava';
import { CopilotMessageNotFound, type Mutex } from '../../base';
import type { CompatSubmissionStore } from '../../plugins/copilot/compat/submission-store';
import type { ConversationPolicy } from '../../plugins/copilot/conversation/policy';
import type { Turn } from '../../plugins/copilot/core';
import { ConversationHost } from '../../plugins/copilot/runtime/hosts/conversation-host';
import {
ChatSession,
type ChatSessionService,
} from '../../plugins/copilot/session';
const test = ava;
function fixture(
options: { failFirstAcceptedWrite?: boolean; failFirstAppend?: boolean } = {}
) {
const sessionId = 'session-1';
const token = 'submission-1';
const durable = new Map<string, Turn>();
const accepted = new Map<string, { sessionId: string; turnId: string }>();
const submissions = new Map([
[
token,
{
id: token,
sessionId,
content: 'hello',
attachments: [],
params: { tone: 'brief' },
createdAt: new Date('2026-01-01T00:00:00.000Z'),
},
],
]);
let appendCount = 0;
let quota = true;
let acceptedWriteCount = 0;
const chatSession = new ChatSession(
{
sessionId,
userId: 'user-1',
workspaceId: 'workspace-1',
docId: 'doc-1',
prompt: {
name: 'Chat With AFFiNE AI',
config: {},
paramKeys: [],
params: {},
},
turns: [],
},
() => []
);
const sessions = {
get: async (id: string) => (id === sessionId ? chatSession : undefined),
findTurnByCompatSubmissionId: async (_sessionId: string, id: string) =>
durable.get(id),
appendTurn: async (input: { compatSubmissionId: string; turn: Turn }) => {
appendCount += 1;
if (options.failFirstAppend && appendCount === 1) {
throw new Error('durable append failed');
}
const stored = { ...input.turn, id: `turn-${appendCount}` };
durable.set(input.compatSubmissionId, stored);
return stored;
},
getMessage: async (_sessionId: string, turnId: string) =>
[...durable.values()].find(turn => turn.id === turnId),
revertLatestMessage: async () => {},
} as unknown as ChatSessionService;
const submissionStore = {
get: async (id: string) => submissions.get(id),
getAccepted: async (id: string) => {
const value = accepted.get(id);
return value
? { ...value, acceptedAt: new Date('2026-01-01T00:00:00.000Z') }
: undefined;
},
markAccepted: async (
id: string,
value: { sessionId: string; turnId: string }
) => {
acceptedWriteCount += 1;
if (options.failFirstAcceptedWrite && acceptedWriteCount === 1) {
throw new Error('accepted cache write failed');
}
accepted.set(id, value);
submissions.delete(id);
},
} as unknown as CompatSubmissionStore;
const mutex = {
acquire: async () => ({ async [Symbol.asyncDispose]() {} }),
} as unknown as Mutex;
const policy = {
hasQuota: async () => quota,
} as unknown as ConversationPolicy;
return {
host: new ConversationHost(sessions, submissionStore, mutex, policy),
sessionId,
token,
durable,
accepted,
submissions,
appendCount: () => appendCount,
setQuota: (value: boolean) => {
quota = value;
},
};
}
test('compat submission becomes one durable user turn and replays idempotently', async t => {
const state = fixture();
const first = await state.host.prepareTurn('user-1', state.sessionId, {
messageId: state.token,
});
t.is(first.latestTurn?.content, 'hello');
t.deepEqual(first.latestTurn?.metadata, { tone: 'brief' });
t.is(state.appendCount(), 1);
t.false(state.submissions.has(state.token));
t.truthy(state.accepted.get(state.token));
state.setQuota(false);
const replay = await state.host.prepareTurn('user-1', state.sessionId, {
messageId: state.token,
});
t.is(replay.latestTurn?.id, first.latestTurn?.id);
t.true(replay.quotaBackedRoutesAllowed);
t.is(state.appendCount(), 1);
});
test('durable compat turn recovers after accepted-cache write failure', async t => {
const state = fixture({ failFirstAcceptedWrite: true });
await t.throwsAsync(
state.host.prepareTurn('user-1', state.sessionId, {
messageId: state.token,
}),
{ message: 'accepted cache write failed' }
);
t.is(state.appendCount(), 1);
t.truthy(state.durable.get(state.token));
const recovered = await state.host.prepareTurn('user-1', state.sessionId, {
messageId: state.token,
});
t.is(recovered.latestTurn?.id, state.durable.get(state.token)?.id);
t.is(state.appendCount(), 1);
t.truthy(state.accepted.get(state.token));
});
test('compat submission remains retryable when durable append fails', async t => {
const state = fixture({ failFirstAppend: true });
await t.throwsAsync(
state.host.prepareTurn('user-1', state.sessionId, {
messageId: state.token,
}),
{ message: 'durable append failed' }
);
t.true(state.submissions.has(state.token));
t.false(state.accepted.has(state.token));
const recovered = await state.host.prepareTurn('user-1', state.sessionId, {
messageId: state.token,
});
t.is(recovered.latestTurn?.content, 'hello');
t.is(state.durable.size, 1);
});
test('compat submission cannot be consumed by another session', async t => {
const state = fixture();
const other = fixture();
other.submissions.set(state.token, {
id: state.token,
sessionId: 'session-other',
content: 'secret',
attachments: [],
params: { tone: 'brief' },
createdAt: new Date(),
});
await t.throwsAsync(
other.host.prepareTurn('user-1', other.sessionId, {
messageId: state.token,
}),
{ instanceOf: CopilotMessageNotFound }
);
t.is(other.appendCount(), 0);
});
packages/backend/server/src/__tests__/copilot/copilot-provider.spec.ts
+206 -1076
View File
File diff suppressed because it is too large Load Diff
packages/backend/server/src/__tests__/copilot/copilot.e2e.ts
+246 -2180
View File
File diff suppressed because it is too large Load Diff
packages/backend/server/src/__tests__/copilot/copilot.spec.ts
-2708
View File
File diff suppressed because it is too large Load Diff
packages/backend/server/src/__tests__/copilot/execution-metrics.spec.ts
-42
View File
@@ -1,42 +0,0 @@
import test from 'ava';
import { summarizePreparedRoutes } from '../../plugins/copilot/runtime/execution-metrics';
test('summarizePreparedRoutes should report none when no route is prepared', t => {
t.deepEqual(
summarizePreparedRoutes([{ prepared: undefined }, { prepared: undefined }]),
{
routeCount: 2,
preparedCount: 0,
preparedMode: 'none',
}
);
});
test('summarizePreparedRoutes should report partial when only some routes are prepared', t => {
t.deepEqual(
summarizePreparedRoutes([
{ prepared: { route: {} } as never },
{ prepared: undefined },
]),
{
routeCount: 2,
preparedCount: 1,
preparedMode: 'partial',
}
);
});
test('summarizePreparedRoutes should report all when every route is prepared', t => {
t.deepEqual(
summarizePreparedRoutes([
{ prepared: { route: {} } as never },
{ prepared: { route: {} } as never },
]),
{
routeCount: 2,
preparedCount: 2,
preparedMode: 'all',
}
);
});
packages/backend/server/src/__tests__/copilot/host-services.spec.ts
-1790
View File
File diff suppressed because it is too large Load Diff
packages/backend/server/src/__tests__/copilot/native-provider.spec.ts
-2000
View File
File diff suppressed because it is too large Load Diff
packages/backend/server/src/__tests__/copilot/native-runtime-contract.spec.ts
+116
View File
@@ -0,0 +1,116 @@
import { access } from 'node:fs/promises';
import path from 'node:path';
import ava from 'ava';
import {
buildLlmEmbeddingRequest,
buildLlmImageRequestFromMessages,
buildLlmRerankRequest,
llmBuildCanonicalRequest,
llmBuildCanonicalStructuredRequest,
llmGetBuiltInRouteOptions,
} from '../../native';
import { ChatQuerySchema } from '../../plugins/copilot/types';
const test = ava;
test('canonical request builders cover every execution request kind', t => {
const chat = llmBuildCanonicalRequest({
model: 'route-selected',
messages: [{ role: 'user', content: 'hello' }],
});
const structured = llmBuildCanonicalStructuredRequest({
model: 'route-selected',
messages: [{ role: 'user', content: 'hello' }],
schema: { type: 'object' },
});
const embedding = buildLlmEmbeddingRequest({
model: 'route-selected',
inputs: ['hello'],
dimensions: 4,
});
const rerank = buildLlmRerankRequest('route-selected', {
query: 'hello',
candidates: [{ id: 'one', text: 'world' }],
});
const image = buildLlmImageRequestFromMessages({
model: 'route-selected',
messages: [{ role: 'user', content: 'draw a circle' }],
});
t.is(chat.model, 'route-selected');
t.deepEqual(structured.schema, { type: 'object' });
t.is(embedding.dimensions, 4);
t.is(rerank.candidates[0].id, 'one');
t.is(image.prompt, 'draw a circle');
});
test('image requests stay provider neutral before target selection', t => {
const image = buildLlmImageRequestFromMessages({
model: 'opaque/model:id',
messages: [
{
role: 'user',
content: 'restyle',
attachments: [
{
kind: 'url',
url: 'data:image/png;base64,aW1n',
mimeType: 'image/png',
},
],
},
],
});
t.is(image.model, 'opaque/model:id');
t.is(image.images?.[0].kind, 'data');
});
test('target override is all-or-nothing and preserves opaque model ids', t => {
const parsed = ChatQuerySchema.parse({
profileId: 'profile-1',
modelId: 'vendor/model:B',
});
t.is(parsed.profileId, 'profile-1');
t.is(parsed.modelId, 'vendor/model:B');
t.throws(() => ChatQuerySchema.parse({ profileId: 'profile-1' }));
t.throws(() => ChatQuerySchema.parse({ modelId: 'vendor/model:B' }));
t.is(
ChatQuerySchema.parse({ routeTargetId: 'terra' }).routeTargetId,
'terra'
);
const route = llmGetBuiltInRouteOptions('Chat With AFFiNE AI');
t.is(route?.standardDefaultTargetId, 'luna');
t.is(route?.premiumDefaultTargetId, 'luna');
t.deepEqual(
route?.choices.map(choice => [choice.id, choice.minimumTier]),
[
['luna', 'Standard'],
['terra', 'Premium'],
['gemini', 'Premium'],
['claude', 'Premium'],
]
);
});
test('caller supplied route policy facts are rejected', t => {
for (const field of ['requirements', 'deployment', 'profiles', 'presets']) {
t.throws(() => ChatQuerySchema.parse({ [field]: 'caller-value' }));
}
});
test('Node provider registry and factory are absent', async t => {
const directory = path.join(
process.cwd(),
'packages/backend/server/src/plugins/copilot/providers'
);
for (const file of [
'factory.ts',
'provider-registry.ts',
'registry-service.ts',
]) {
await t.throwsAsync(access(path.join(directory, file)));
}
});
packages/backend/server/src/__tests__/copilot/provider-middleware.spec.ts
-43
View File
@@ -1,43 +0,0 @@
import test from 'ava';
import { resolveProviderMiddleware } from '../../plugins/copilot/providers/provider-middleware';
import { buildProviderRegistry } from '../../plugins/copilot/providers/provider-registry';
import { CopilotProviderType } from '../../plugins/copilot/providers/types';
test('resolveProviderMiddleware should include anthropic defaults', t => {
const middleware = resolveProviderMiddleware(CopilotProviderType.Anthropic);
t.is(middleware.rust, undefined);
t.deepEqual(middleware.node?.text, ['citation_footnote', 'callout']);
});
test('resolveProviderMiddleware should merge defaults and overrides', t => {
const middleware = resolveProviderMiddleware(CopilotProviderType.OpenAI, {
rust: { request: ['clamp_max_tokens'] },
node: { text: ['thinking_format'] },
});
t.deepEqual(middleware.rust?.request, ['clamp_max_tokens']);
t.deepEqual(middleware.node?.text, [
'citation_footnote',
'callout',
'thinking_format',
]);
});
test('buildProviderRegistry should normalize profile middleware defaults', t => {
const registry = buildProviderRegistry({
profiles: [
{
id: 'openai-main',
type: CopilotProviderType.OpenAI,
config: { apiKey: '1' },
},
],
});
const profile = registry.profiles.get('openai-main');
t.truthy(profile);
t.is(profile?.middleware.rust, undefined);
t.deepEqual(profile?.middleware.node?.text, ['citation_footnote', 'callout']);
});
packages/backend/server/src/__tests__/copilot/provider-native.spec.ts
-3242
View File
File diff suppressed because it is too large Load Diff
packages/backend/server/src/__tests__/copilot/provider-registry.spec.ts
-282
View File
@@ -1,282 +0,0 @@
import test from 'ava';
import { OpenAIProvider } from '../../plugins/copilot/providers';
import { CopilotProviderLifecycleService } from '../../plugins/copilot/providers/lifecycle-service';
import {
buildProviderRegistry,
resolveModel,
stripProviderPrefix,
} from '../../plugins/copilot/providers/provider-registry';
import {
CopilotProviderType,
ModelOutputType,
} from '../../plugins/copilot/providers/types';
test('buildProviderRegistry should keep explicit profile over legacy compatibility profile', t => {
const registry = buildProviderRegistry({
profiles: [
{
id: 'openai-default',
type: CopilotProviderType.OpenAI,
priority: 100,
config: { apiKey: 'new' },
},
],
openai: { apiKey: 'legacy' },
});
const profile = registry.profiles.get('openai-default');
t.truthy(profile);
t.deepEqual(profile?.config, { apiKey: 'new' });
});
test('buildProviderRegistry should reject duplicated profile ids', t => {
const error = t.throws(() =>
buildProviderRegistry({
profiles: [
{
id: 'openai-main',
type: CopilotProviderType.OpenAI,
config: { apiKey: '1' },
},
{
id: 'openai-main',
type: CopilotProviderType.OpenAI,
config: { apiKey: '2' },
},
],
})
) as Error;
t.truthy(error);
t.regex(error.message, /Duplicated copilot provider profile id/);
});
test('buildProviderRegistry should reject defaults that reference unknown providers', t => {
const error = t.throws(() =>
buildProviderRegistry({
profiles: [
{
id: 'openai-main',
type: CopilotProviderType.OpenAI,
config: { apiKey: '1' },
},
],
defaults: {
fallback: 'unknown-provider',
},
})
) as Error;
t.truthy(error);
t.regex(error.message, /defaults references unknown providerId/);
});
test('resolveModel should support explicit provider prefix and keep slash models untouched', t => {
const registry = buildProviderRegistry({
profiles: [
{
id: 'openai-main',
type: CopilotProviderType.OpenAI,
config: { apiKey: '1' },
},
{
id: 'fal-main',
type: CopilotProviderType.FAL,
config: { apiKey: '2' },
},
],
});
const prefixed = resolveModel({
registry,
modelId: 'openai-main/gpt-5-mini',
});
t.deepEqual(prefixed, {
rawModelId: 'openai-main/gpt-5-mini',
modelId: 'gpt-5-mini',
explicitProviderId: 'openai-main',
candidateProviderIds: ['openai-main'],
});
const slashModel = resolveModel({
registry,
modelId: 'lora/image-to-image',
});
t.is(slashModel.modelId, 'lora/image-to-image');
t.false(slashModel.candidateProviderIds.includes('lora'));
});
test('resolveModel should follow defaults -> fallback -> order and apply filters', t => {
const registry = buildProviderRegistry({
profiles: [
{
id: 'openai-main',
type: CopilotProviderType.OpenAI,
priority: 10,
config: { apiKey: '1' },
},
{
id: 'anthropic-main',
type: CopilotProviderType.Anthropic,
priority: 5,
config: { apiKey: '2' },
},
{
id: 'fal-main',
type: CopilotProviderType.FAL,
priority: 1,
config: { apiKey: '3' },
},
],
defaults: {
[ModelOutputType.Text]: 'anthropic-main',
fallback: 'openai-main',
},
});
const routed = resolveModel({
registry,
outputType: ModelOutputType.Text,
preferredProviderIds: ['openai-main', 'fal-main'],
});
t.deepEqual(routed.candidateProviderIds, ['openai-main', 'fal-main']);
});
test('resolveModel should resolve bare model ids by provider priority order', t => {
const registry = buildProviderRegistry({
profiles: [
{
id: 'openai-main',
type: CopilotProviderType.OpenAI,
priority: 10,
config: { apiKey: '1' },
},
{
id: 'anthropic-main',
type: CopilotProviderType.Anthropic,
priority: 5,
config: { apiKey: '2' },
},
{
id: 'fal-main',
type: CopilotProviderType.FAL,
priority: 1,
config: { apiKey: '3' },
},
],
defaults: {
[ModelOutputType.Text]: 'anthropic-main',
fallback: 'fal-main',
},
});
const routed = resolveModel({
registry,
modelId: 'shared-model',
});
t.deepEqual(routed.candidateProviderIds, [
'openai-main',
'anthropic-main',
'fal-main',
]);
});
test('stripProviderPrefix should only strip matched provider prefix', t => {
const registry = buildProviderRegistry({
profiles: [
{
id: 'openai-main',
type: CopilotProviderType.OpenAI,
config: { apiKey: '1' },
},
],
});
t.is(
stripProviderPrefix(registry, 'openai-main', 'openai-main/gpt-5-mini'),
'gpt-5-mini'
);
t.is(
stripProviderPrefix(registry, 'openai-main', 'another-main/gpt-5-mini'),
'another-main/gpt-5-mini'
);
t.is(
stripProviderPrefix(registry, 'openai-main', 'gpt-5-mini'),
'gpt-5-mini'
);
});
test('CopilotProviderLifecycleService should register current profiles and unregister stale ones', async t => {
const calls: string[] = [];
let registry = buildProviderRegistry({
profiles: [
{
id: 'openai-main',
type: CopilotProviderType.OpenAI,
config: { apiKey: '1' },
},
{
id: 'openai-backup',
type: CopilotProviderType.OpenAI,
config: { apiKey: '2' },
},
],
});
const provider = {
type: CopilotProviderType.OpenAI,
configured(execution: { providerId?: string } | undefined) {
return execution?.providerId === 'openai-main';
},
};
const service = new CopilotProviderLifecycleService(
{
get(token: unknown) {
return token === OpenAIProvider ? provider : undefined;
},
} as any,
{
register(providerId: string) {
calls.push(`register:${providerId}`);
},
unregister(providerId: string) {
calls.push(`unregister:${providerId}`);
},
} as any,
{
getRegistry() {
return registry;
},
} as any
);
await service.syncProviders();
t.deepEqual(calls.slice().sort(), [
'register:openai-main',
'unregister:openai-backup',
]);
calls.length = 0;
registry = buildProviderRegistry({
profiles: [
{
id: 'openai-backup',
type: CopilotProviderType.OpenAI,
config: { apiKey: '2' },
},
],
});
provider.configured = (execution: { providerId?: string } | undefined) =>
execution?.providerId === 'openai-backup';
await service.syncProviders();
t.deepEqual(calls.slice().sort(), [
'register:openai-backup',
'unregister:openai-main',
]);
});
packages/backend/server/src/__tests__/copilot/provider-template.spec.ts
-201
View File
@@ -1,201 +0,0 @@
import serverNativeModule from '@affine/server-native';
import test from 'ava';
import { z } from 'zod';
import type {
LlmEmbeddingRequest,
LlmRerankRequest,
LlmStructuredRequest,
} from '../../native';
import { CopilotProvider } from '../../plugins/copilot/providers/provider';
import type { ProviderDriverSpec } from '../../plugins/copilot/providers/provider-runtime-contract';
import { CopilotProviderType } from '../../plugins/copilot/providers/types';
import {
buildStructuredResponseContract,
type RequiredStructuredOutputContract,
requireStructuredOutputContract,
} from '../../plugins/copilot/runtime/contracts';
import { getProviderRuntimeHost } from '../../plugins/copilot/runtime/provider-runtime-context';
import { nativeUserText, singleUserPromptMessages } from './prompt-test-helper';
function structuredOptions(schema: z.ZodTypeAny) {
const { responseSchemaJson, schemaHash } =
buildStructuredResponseContract(schema);
return { responseSchemaJson, schemaHash };
}
function structuredContract(
schema: z.ZodTypeAny
): RequiredStructuredOutputContract {
const contract = buildStructuredResponseContract(schema);
const requiredContract = requireStructuredOutputContract(contract);
if (!requiredContract) {
throw new Error('structured response contract is required');
}
return requiredContract;
}
class TemplateOnlyProvider extends CopilotProvider<{ apiKey: string }> {
readonly type = CopilotProviderType.OpenAI;
protected resolveModelBackendKind() {
return 'openai_responses' as const;
}
readonly structuredRequests: LlmStructuredRequest[] = [];
readonly embeddingRequests: LlmEmbeddingRequest[] = [];
readonly rerankRequests: Array<{
model: string;
query: string;
candidates: Array<{ id?: string; text: string }>;
topN?: number;
}> = [];
configured() {
return true;
}
override getDriverSpec(): ProviderDriverSpec {
return {
createBackendConfig: () => ({
base_url: 'https://api.openai.com',
auth_token: 'test-key',
}),
mapError: (error: unknown) => error,
structured: {},
embedding: {
defaultDimensions: 8,
},
rerank: {},
};
}
}
test('template-only provider should reuse base structured, embedding and rerank drivers', async t => {
const provider = new TemplateOnlyProvider();
const originalStructured = (serverNativeModule as any).llmStructuredDispatch;
const originalEmbedding = (serverNativeModule as any).llmEmbeddingDispatch;
const originalRerank = (serverNativeModule as any).llmRerankDispatch;
(serverNativeModule as any).llmStructuredDispatch = (
_protocol: string,
_backendConfigJson: string,
requestJson: string
) => {
provider.structuredRequests.push(
JSON.parse(requestJson) as LlmStructuredRequest
);
return JSON.stringify({
id: 'structured_1',
model: 'gpt-5-mini',
output_text: '{"summary":"native"}',
output_json: { summary: 'native' },
usage: {
prompt_tokens: 3,
completion_tokens: 2,
total_tokens: 5,
},
finish_reason: 'stop',
});
};
(serverNativeModule as any).llmEmbeddingDispatch = (
_protocol: string,
_backendConfigJson: string,
requestJson: string
) => {
const request = JSON.parse(requestJson) as LlmEmbeddingRequest;
provider.embeddingRequests.push(request);
return JSON.stringify({
model: request.model,
embeddings: request.inputs.map((_, index) => [index + 0.1, index + 0.2]),
});
};
(serverNativeModule as any).llmRerankDispatch = (
_protocol: string,
_backendConfigJson: string,
requestJson: string
) => {
const request = JSON.parse(requestJson) as LlmRerankRequest;
provider.rerankRequests.push(request);
return JSON.stringify({
model: request.model,
scores: request.candidates.map((_candidate, index) =>
index === 0 ? 0.9 : 0.1
),
});
};
t.teardown(() => {
(serverNativeModule as any).llmStructuredDispatch = originalStructured;
(serverNativeModule as any).llmEmbeddingDispatch = originalEmbedding;
(serverNativeModule as any).llmRerankDispatch = originalRerank;
});
const structured = await getProviderRuntimeHost(provider).run.structured(
{ modelId: 'gpt-5-mini' },
singleUserPromptMessages('summarize this'),
structuredOptions(z.object({ summary: z.string() })),
structuredContract(z.object({ summary: z.string() }))
);
const embeddings = await getProviderRuntimeHost(provider).run.embedding(
{ modelId: 'text-embedding-3-small' },
['alpha', 'beta'],
{
dimensions: 8,
}
);
const scores = await getProviderRuntimeHost(provider).run.rerank(
{ modelId: 'gpt-4o-mini' },
{
query: 'alpha',
candidates: [
{ id: 'alpha', text: 'alpha result' },
{ id: 'beta', text: 'beta result' },
],
topK: 1,
}
);
t.is(structured, JSON.stringify({ summary: 'native' }));
t.deepEqual(embeddings, [
[0.1, 0.2],
[1.1, 1.2],
]);
t.deepEqual(scores, [0.9, 0.1]);
t.is(provider.structuredRequests.length, 1);
t.like(provider.structuredRequests[0], {
model: 'gpt-5-mini',
messages: [
{ role: 'user', content: nativeUserText('summarize this').content },
],
schema: {
type: 'object',
properties: {
summary: { type: 'string' },
},
required: ['summary'],
additionalProperties: false,
},
strict: true,
responseMimeType: 'application/json',
});
t.is(provider.structuredRequests[0]?.middleware, undefined);
t.deepEqual(provider.embeddingRequests, [
{
model: 'text-embedding-3-small',
inputs: ['alpha', 'beta'],
dimensions: 8,
taskType: 'RETRIEVAL_DOCUMENT',
},
]);
t.deepEqual(provider.rerankRequests, [
{
model: 'gpt-4o-mini',
query: 'alpha',
candidates: [
{ id: 'alpha', text: 'alpha result' },
{ id: 'beta', text: 'beta result' },
],
topN: 1,
},
]);
});
packages/backend/server/src/__tests__/copilot/runtime-boundaries.spec.ts
+431
View File
@@ -0,0 +1,431 @@
import { EventEmitter } from 'node:events';
import ava from 'ava';
import { firstValueFrom } from 'rxjs';
import type { Config, JobQueue } from '../../base';
import { ServerFeature, type ServerService } from '../../core';
import type { Models } from '../../models';
import { HistoryPromptPreloadProjector } from '../../plugins/copilot/compat/history-prompt-preload-projector';
import { CopilotController } from '../../plugins/copilot/controller';
import { ConversationPolicy } from '../../plugins/copilot/conversation/policy';
import {
chatMessageFromTurn,
type Turn,
turnFromChatMessage,
} from '../../plugins/copilot/core';
import { CopilotCronJobs } from '../../plugins/copilot/cron';
import {
CopilotFeatureGuard,
CopilotFeatureService,
} from '../../plugins/copilot/feature';
import type { PromptService } from '../../plugins/copilot/prompt';
import type { ResolvedPrompt } from '../../plugins/copilot/prompt/spec';
import { TextStreamParser } from '../../plugins/copilot/providers/utils';
import {
projectActionEventToChatEvent,
projectActionResultToAssistantTurn,
} from '../../plugins/copilot/runtime/action-output-projector';
import type { ActionStreamHost } from '../../plugins/copilot/runtime/hosts/action-stream-host';
import type { TurnOrchestrator } from '../../plugins/copilot/runtime/turn-orchestrator';
import { ChatSession } from '../../plugins/copilot/session';
import type { CopilotStorage } from '../../plugins/copilot/storage';
const test = ava;
test('copilot config controls the server feature and request admission', t => {
const config = { copilot: { enabled: false } } as Config;
const features = new Set<ServerFeature>();
const server = {
enableFeature: (feature: ServerFeature) => features.add(feature),
disableFeature: (feature: ServerFeature) => features.delete(feature),
} as unknown as ServerService;
const feature = new CopilotFeatureService(config, server);
const guard = new CopilotFeatureGuard(feature);
feature.onConfigInit();
t.false(features.has(ServerFeature.Copilot));
t.throws(() => guard.canActivate(), { message: 'Copilot is disabled.' });
config.copilot.enabled = true;
feature.onConfigChanged({ updates: { copilot: { enabled: true } } });
t.true(features.has(ServerFeature.Copilot));
t.true(guard.canActivate());
config.copilot.enabled = false;
feature.onConfigChanged({ updates: { copilot: { enabled: false } } });
t.false(features.has(ServerFeature.Copilot));
});
const prompt: ResolvedPrompt = {
name: 'Chat With AFFiNE AI',
config: {},
paramKeys: [],
params: {},
};
function turn(
conversationId: string,
role: Turn['role'],
content: string,
extra: Partial<Turn> = {}
): Turn {
return {
conversationId,
role,
content,
attachments: [],
renderTrace: [],
toolEvents: [],
metadata: {},
createdAt: new Date('2026-01-01T00:00:00.000Z'),
...extra,
};
}
test('chat session preserves prompt params, attachments, stash and revert semantics', async t => {
const saved: Turn[][] = [];
const session = new ChatSession(
{
sessionId: 'session-1',
userId: 'user-1',
workspaceId: 'workspace-1',
docId: 'doc-1',
prompt,
turns: [turn('session-1', 'user', 'persisted')],
},
(_prompt, turns, params) => [
{ role: 'system', content: `hello ${params.word}` },
...turns,
],
async state => {
saved.push(state.turns);
}
);
session.pushTurn(
turn('session-1', 'assistant', 'answer', {
attachments: [
{
kind: 'file_handle',
fileHandle: 'file-1',
mimeType: 'application/pdf',
},
],
metadata: { word: 'world' },
})
);
t.is(session.stashTurns.length, 1);
t.deepEqual(session.finish({ word: 'direct' }), [
{ role: 'system', content: 'hello direct' },
{
role: 'user',
content: 'persisted',
attachments: undefined,
params: undefined,
},
{
role: 'assistant',
content: 'answer',
attachments: [
{
kind: 'file_handle',
fileHandle: 'file-1',
mimeType: 'application/pdf',
},
],
params: { word: 'world' },
},
]);
await session.save();
t.is(session.stashTurns.length, 0);
t.deepEqual(
saved[0].map(item => item.content),
['answer']
);
session.pushTurn(turn('session-1', 'user', 'retry'));
session.pushTurn(turn('session-1', 'assistant', 'retry answer'));
session.revertLatestMessage(false);
t.deepEqual(
session.finish({ word: 'direct' }).map(item => item.content),
['hello direct', 'persisted', 'answer', 'retry']
);
session.revertLatestMessage(true);
t.deepEqual(
session.finish({ word: 'direct' }).map(item => item.content),
['hello direct', 'persisted', 'answer']
);
});
test('chat message adapters preserve and canonicalize assistant render trace', t => {
const message = {
id: 'message-1',
role: 'assistant' as const,
content: 'Final answer',
params: { schemaVersion: 'v1' },
streamObjects: [
{ type: 'reasoning' as const, textDelta: 'Plan ' },
{ type: 'reasoning' as const, textDelta: 'first' },
{
type: 'tool-call' as const,
toolCallId: 'call-1',
toolName: 'doc_read',
args: { docId: 'doc-1' },
},
{
type: 'tool-result' as const,
toolCallId: 'call-1',
toolName: 'doc_read',
args: { docId: 'doc-1' },
result: { markdown: '# AFFiNE' },
},
{ type: 'text-delta' as const, textDelta: 'Final answer' },
],
createdAt: new Date('2026-01-01T00:00:00.000Z'),
};
const converted = turnFromChatMessage(message, 'session-1');
t.deepEqual(converted.renderTrace, [
{ type: 'reasoning', textDelta: 'Plan first' },
{
type: 'tool-result',
toolCallId: 'call-1',
toolName: 'doc_read',
args: { docId: 'doc-1' },
result: { markdown: '# AFFiNE' },
},
{ type: 'text-delta', textDelta: 'Final answer' },
]);
t.deepEqual(
converted.toolEvents.map(event => event.type),
['tool_result']
);
t.deepEqual(chatMessageFromTurn(converted), {
...message,
attachments: undefined,
streamObjects: converted.renderTrace,
});
});
test('action output projection preserves public SSE and assistant-turn contracts', t => {
const session = new ChatSession(
{
sessionId: 'session-1',
userId: 'user-1',
workspaceId: 'workspace-1',
docId: 'doc-1',
prompt,
turns: [],
},
() => []
);
t.deepEqual(
projectActionEventToChatEvent('message-1', {
type: 'action_done',
actionId: 'slides.outline',
actionVersion: 'v1',
status: 'succeeded',
runId: 'run-1',
result: { content: '- Launch deck' },
}),
{ type: 'message', id: 'message-1', data: '- Launch deck' }
);
t.like(
projectActionResultToAssistantTurn({
session,
actionId: 'image.filter.remove-background',
result: {},
artifacts: [{ url: 'https://example.com/result.png' }],
wasAborted: false,
}),
{
conversationId: 'session-1',
role: 'assistant',
attachments: ['https://example.com/result.png'],
}
);
t.is(
projectActionResultToAssistantTurn({
session,
actionId: 'transcript.audio',
result: {},
wasAborted: false,
}),
null
);
});
test('text stream parser keeps reasoning and tool output distinct from answer text', t => {
const parser = new TextStreamParser();
const output = [
parser.parse({ type: 'reasoning-delta', text: 'Think' }),
parser.parse({
type: 'tool-call',
toolCallId: 'call-1',
toolName: 'web_search_exa',
input: { query: 'AFFiNE' },
}),
parser.parse({
type: 'tool-result',
toolCallId: 'call-1',
toolName: 'web_search_exa',
input: { query: 'AFFiNE' },
output: [{ title: 'AFFiNE', url: 'https://affine.pro' }],
}),
parser.parse({ type: 'text-delta', text: 'Answer' }),
].join('');
t.true(output.includes('Think'));
t.true(output.includes('Searching the web "AFFiNE"'));
t.true(output.includes('[AFFiNE](https://affine.pro)'));
t.true(output.endsWith('\nAnswer'));
t.throws(
() => parser.parse({ type: 'error', error: { message: 'failed' } }),
{ message: 'failed' }
);
});
test('history prompt preload excludes system messages and precedes durable history', t => {
const projector = new HistoryPromptPreloadProjector({
finish: () => [
{ role: 'system', content: 'hidden system' },
{ role: 'user', content: 'preloaded question' },
],
} as unknown as PromptService);
const createdAt = new Date('2026-01-01T00:00:00.000Z');
const history = {
conversation: {
id: 'session-1',
userId: 'user-1',
workspaceId: 'workspace-1',
docId: 'doc-1',
pinned: false,
parentId: null,
title: null,
createdAt,
updatedAt: createdAt,
},
prompt,
turns: [
turn('session-1', 'user', 'hello', { metadata: { tone: 'brief' } }),
],
};
t.deepEqual(
projector.project(history, true, true).map(item => item.content),
['preloaded question']
);
t.deepEqual(projector.project(history, true, false), []);
t.true(projector.project(history, true, true)[0].createdAt! < createdAt);
});
test('title policy and cron scheduling retain background-job invariants', async t => {
const policy = new ConversationPolicy({} as Models, {} as never);
t.true(
policy.shouldGenerateTitle({
title: null,
turns: [
turn('session-1', 'user', 'Question'),
turn('session-1', 'assistant', 'Answer'),
],
})
);
t.false(
policy.shouldGenerateTitle({
title: 'Existing',
turns: [turn('session-1', 'user', 'Question')],
})
);
const calls: unknown[][] = [];
const jobs = {
add: async (...args: unknown[]) => calls.push(args),
} as unknown as JobQueue;
const models = {
copilotSession: {
toBeGenerateTitle: async () => [{ id: 'session-1' }, { id: 'session-2' }],
},
} as unknown as Models;
const cron = new CopilotCronJobs(models, jobs);
await cron.dailyCleanupJob();
await cron.generateMissingTitles();
t.deepEqual(calls, [
[
'copilot.session.cleanupEmptySessions',
{},
{ jobId: 'daily-copilot-cleanup-empty-sessions' },
],
[
'copilot.session.generateMissingTitles',
{},
{ jobId: 'daily-copilot-generate-missing-titles' },
],
[
'copilot.workspace.cleanupTrashedDocEmbeddings',
{},
{ jobId: 'daily-copilot-cleanup-trashed-doc-embeddings' },
],
[
'copilot.session.generateTitle',
{ sessionId: 'session-1' },
{ priority: 100 },
],
[
'copilot.session.generateTitle',
{ sessionId: 'session-2' },
{ priority: 100 },
],
]);
});
test('controller projects successful streams and preparation failures to SSE events', async t => {
const request = { socket: new EventEmitter() } as never;
const orchestrator = {
streamText: async () => ({
messageId: 'message-1',
model: 'route-selected',
finalMessage: [],
stream: (async function* () {
yield 'hello';
})(),
}),
} as unknown as TurnOrchestrator;
const actions = {
stream: async () => {
throw new Error('action preparation failed');
},
} as unknown as ActionStreamHost;
const controller = new CopilotController(
{ copilot: { unsplash: {} } } as Config,
orchestrator,
actions,
{} as CopilotStorage
);
t.deepEqual(
await firstValueFrom(
await controller.chatStream(
{ id: 'user-1' } as never,
request,
'session-1',
{}
)
),
{ type: 'message', id: 'message-1', data: 'hello' }
);
t.like(
await firstValueFrom(
await controller.actionStream(
{ id: 'user-1' } as never,
request,
'session-1',
{}
)
),
{ type: 'error' }
);
});
packages/backend/server/src/__tests__/copilot/tool-call-loop.spec.ts
-615
View File
@@ -1,615 +0,0 @@
import serverNativeModule from '@affine/server-native';
import test from 'ava';
import { z } from 'zod';
import type { DocReader } from '../../core/doc';
import type { PermissionAccess } from '../../core/permission';
import type { Models } from '../../models';
import {
LlmRequest,
type LlmToolCallbackRequest,
type LlmToolCallbackResponse,
type LlmToolLoopStreamEvent,
llmValidateContract,
} from '../../native';
import {
buildToolContracts,
parseToolContract,
parseToolLoopStreamEvent,
} from '../../plugins/copilot/runtime/contracts';
import {
createToolExecutionCallback,
createToolLoopBridge,
} from '../../plugins/copilot/runtime/tool/bridge';
import {
buildBlobContentGetter,
createBlobReadTool,
} from '../../plugins/copilot/tools/blob-read';
import {
buildDocKeywordSearchGetter,
createDocKeywordSearchTool,
} from '../../plugins/copilot/tools/doc-keyword-search';
import {
buildDocContentGetter,
createDocReadTool,
} from '../../plugins/copilot/tools/doc-read';
import {
buildDocSearchGetter,
createDocSemanticSearchTool,
} from '../../plugins/copilot/tools/doc-semantic-search';
import {
DOCUMENT_SYNC_PENDING_MESSAGE,
LOCAL_WORKSPACE_SYNC_REQUIRED_MESSAGE,
} from '../../plugins/copilot/tools/doc-sync';
import { defineTool } from '../../plugins/copilot/tools/tool';
import {
nativeMessages,
nativeUserText,
singleUserPromptMessages,
} from './prompt-test-helper';
test('defineTool should freeze json schema at definition time', t => {
const tool = defineTool({
description: 'Read doc',
inputSchema: z.object({
doc_id: z.string(),
limit: z.number().optional(),
}),
execute: async () => ({}),
});
t.deepEqual(tool.jsonSchema, {
type: 'object',
properties: {
doc_id: { type: 'string' },
limit: { type: 'number' },
},
additionalProperties: false,
required: ['doc_id'],
});
});
test('buildToolContracts should project precomputed json schema', t => {
const toolSet = {
doc_read: defineTool({
description: 'Read doc',
inputSchema: z.object({
doc_id: z.string(),
limit: z.number().optional(),
}),
execute: async () => ({}),
}),
};
const extracted = buildToolContracts(toolSet);
t.deepEqual(extracted, [
{
name: 'doc_read',
description: 'Read doc',
parameters: {
type: 'object',
properties: {
doc_id: { type: 'string' },
limit: { type: 'number' },
},
additionalProperties: false,
required: ['doc_id'],
},
},
]);
});
test('buildToolContracts should reject tool definitions without json schema', t => {
const error = t.throws(() =>
buildToolContracts({
doc_read: {
description: 'Read doc',
inputSchema: z.object({ doc_id: z.string() }),
execute: async () => ({}),
} as never,
})
);
t.regex(error.message, /missing precomputed jsonSchema/);
});
test('defineTool should prefer explicit json schema when provided', t => {
const extracted = buildToolContracts({
doc_read: defineTool({
description: 'Read doc',
jsonSchema: {
type: 'object',
properties: {
doc_id: { type: 'string' },
},
required: ['doc_id'],
},
inputSchema: z.object({
doc_id: z.string(),
ignored: z.number(),
}),
execute: async () => ({}),
}),
});
t.deepEqual(extracted, [
{
name: 'doc_read',
description: 'Read doc',
parameters: {
type: 'object',
properties: {
doc_id: { type: 'string' },
},
required: ['doc_id'],
},
},
]);
});
test('ToolContract should freeze stable tool schema and callback payloads', t => {
const tool = parseToolContract({
name: 'doc_read',
description: 'Read doc',
parameters: {
type: 'object',
properties: {
doc_id: { type: 'string' },
},
required: ['doc_id'],
},
});
const result = llmValidateContract<LlmToolCallbackResponse>(
'toolCallbackResponse',
{
callId: 'call_1',
name: 'doc_read',
args: { doc_id: 'a1' },
output: { markdown: '# a1' },
}
);
const request = llmValidateContract<LlmToolCallbackRequest>(
'toolCallbackRequest',
{
callId: 'call_1',
name: 'doc_read',
args: { doc_id: 'a1' },
}
);
t.is(tool.name, 'doc_read');
t.deepEqual(request.args, { doc_id: 'a1' });
t.deepEqual(result.args, { doc_id: 'a1' });
});
test('ToolLoopStreamEvent should reject malformed tool_result metadata at decode boundary', t => {
const event = parseToolLoopStreamEvent({
type: 'tool_result',
call_id: 'call_1',
name: 'doc_read',
arguments: { doc_id: 'a1' },
output: { markdown: '# a1' },
});
t.is(event.type, 'tool_result');
const error = t.throws(() =>
parseToolLoopStreamEvent({
type: 'tool_result',
call_id: 'call_1',
output: { markdown: '# a1' },
})
);
t.truthy(error);
});
test('createNativeToolExecutionCallback should preserve tool execution ABI', async t => {
const callback = createToolExecutionCallback(
{
doc_read: {
inputSchema: z.object({ doc_id: z.string() }),
execute: async args => ({ markdown: `# ${String(args.doc_id)}` }),
},
},
{ messages: singleUserPromptMessages('read doc') }
);
const result = await callback({
callId: 'call_1',
name: 'doc_read',
args: { doc_id: 'a1' },
rawArgumentsText: '{"doc_id":"a1"}',
});
t.deepEqual(result, {
callId: 'call_1',
name: 'doc_read',
args: { doc_id: 'a1' },
rawArgumentsText: '{"doc_id":"a1"}',
argumentParseError: undefined,
output: { markdown: '# a1' },
});
});
test('createNativeToolLoopBridge should preserve native callback and stream ABI', async t => {
const capturedRequests: LlmRequest[] = [];
const originalMessages = singleUserPromptMessages('read doc');
const signal = new AbortController().signal;
let executedArgs: Record<string, unknown> | null = null;
let executedMessages: unknown;
let executedSignal: AbortSignal | undefined;
const original = (serverNativeModule as any).llmDispatchToolLoopStream;
(serverNativeModule as any).llmDispatchToolLoopStream = (
_protocol: string,
_backendConfigJson: string,
requestJson: string,
maxSteps: number,
callback: (error: Error | null, eventJson: string) => void,
toolCallback: (error: Error | null, requestJson: string) => Promise<string>
) => {
capturedRequests.push(JSON.parse(requestJson) as LlmRequest);
t.is(maxSteps, 4);
void (async () => {
callback(
null,
JSON.stringify({
type: 'tool_call',
call_id: 'call_1',
name: 'doc_read',
arguments: { doc_id: 'a1' },
})
);
const result = JSON.parse(
await toolCallback(
null,
JSON.stringify({
callId: 'call_1',
name: 'doc_read',
args: { doc_id: 'a1' },
rawArgumentsText: '{"doc_id":"a1"}',
})
)
) as {
callId: string;
name: string;
args: Record<string, unknown>;
rawArgumentsText?: string;
argumentParseError?: string;
output: unknown;
isError?: boolean;
};
callback(
null,
JSON.stringify({
type: 'tool_result',
call_id: result.callId,
name: result.name,
arguments: result.args,
arguments_text: result.rawArgumentsText,
arguments_error: result.argumentParseError,
output: result.output,
is_error: result.isError,
})
);
callback(null, JSON.stringify({ type: 'text_delta', text: 'done' }));
callback(null, JSON.stringify({ type: 'done', finish_reason: 'stop' }));
callback(null, '__AFFINE_LLM_STREAM_END__');
})();
return {
abort() {},
};
};
t.teardown(() => {
(serverNativeModule as any).llmDispatchToolLoopStream = original;
});
const bridge = createToolLoopBridge(
{
protocol: 'openai_chat',
backendConfig: {
base_url: 'https://api.openai.com',
auth_token: 'test-key',
},
},
{
doc_read: {
inputSchema: z.object({ doc_id: z.string() }),
execute: async (args, options) => {
executedArgs = args;
executedMessages = options.messages;
executedSignal = options.signal;
return { markdown: '# doc' };
},
},
},
4
);
const events: LlmToolLoopStreamEvent[] = [];
for await (const event of bridge(
{
model: 'gpt-5-mini',
stream: false,
messages: nativeMessages(nativeUserText('read doc')),
},
signal,
[...originalMessages]
)) {
events.push(event);
}
t.deepEqual(executedArgs, { doc_id: 'a1' });
t.deepEqual(executedMessages, originalMessages);
t.is(executedSignal, signal);
t.true(capturedRequests[0]?.stream);
t.deepEqual(
events.map(event => event.type),
['tool_call', 'tool_result', 'text_delta', 'done']
);
});
test('doc_read should return specific sync errors for unavailable docs', async t => {
const cases = [
{
name: 'local workspace without cloud sync',
workspace: null,
authors: null,
markdown: null,
expected: {
type: 'error',
name: 'Workspace Sync Required',
message: LOCAL_WORKSPACE_SYNC_REQUIRED_MESSAGE,
},
docReaderCalled: false,
},
{
name: 'cloud workspace document not synced to server yet',
workspace: { id: 'ws-1' },
authors: null,
markdown: null,
expected: {
type: 'error',
name: 'Document Sync Pending',
message: DOCUMENT_SYNC_PENDING_MESSAGE('doc-1'),
},
docReaderCalled: false,
},
{
name: 'cloud workspace document markdown not ready yet',
workspace: { id: 'ws-1' },
authors: {
createdAt: new Date('2026-01-01T00:00:00.000Z'),
updatedAt: new Date('2026-01-01T00:00:00.000Z'),
createdByUser: null,
updatedByUser: null,
},
markdown: null,
expected: {
type: 'error',
name: 'Document Sync Pending',
message: DOCUMENT_SYNC_PENDING_MESSAGE('doc-1'),
},
docReaderCalled: true,
},
] as const;
const ac = {
user: () => ({
workspace: () => ({ doc: () => ({ can: async () => true }) }),
}),
} as unknown as PermissionAccess;
for (const testCase of cases) {
let docReaderCalled = false;
const docReader = {
getDocMarkdown: async () => {
docReaderCalled = true;
return testCase.markdown;
},
} as unknown as DocReader;
const models = {
workspace: {
get: async () => testCase.workspace,
},
doc: {
getAuthors: async () => testCase.authors,
},
} as unknown as Models;
const getDoc = buildDocContentGetter(ac, docReader, models);
const tool = createDocReadTool(
getDoc.bind(null, {
user: 'user-1',
workspace: 'workspace-1',
})
);
const result = await tool.execute?.({ doc_id: 'doc-1' }, {});
t.is(docReaderCalled, testCase.docReaderCalled, testCase.name);
t.deepEqual(result, testCase.expected, testCase.name);
}
});
test('document search tools should return sync error for local workspace', async t => {
const ac = {
user: () => ({
workspace: () => ({
can: async () => true,
docs: async () => [],
}),
}),
} as unknown as PermissionAccess;
const models = {
workspace: {
get: async () => null,
},
} as unknown as Models;
let keywordSearchCalled = false;
const indexerService = {
searchDocsByKeyword: async () => {
keywordSearchCalled = true;
return [];
},
} as unknown as Parameters<typeof buildDocKeywordSearchGetter>[1];
let semanticSearchCalled = false;
const contextService = {
matchWorkspaceAll: async () => {
semanticSearchCalled = true;
return [];
},
} as unknown as Parameters<typeof buildDocSearchGetter>[1];
const keywordTool = createDocKeywordSearchTool(
buildDocKeywordSearchGetter(ac, indexerService, models).bind(null, {
user: 'user-1',
workspace: 'workspace-1',
})
);
const semanticTool = createDocSemanticSearchTool(
buildDocSearchGetter(ac, contextService, undefined, models).bind(null, {
user: 'user-1',
workspace: 'workspace-1',
})
);
const keywordResult = await keywordTool.execute?.({ query: 'hello' }, {});
const semanticResult = await semanticTool.execute?.({ query: 'hello' }, {});
t.false(keywordSearchCalled);
t.false(semanticSearchCalled);
t.deepEqual(keywordResult, {
type: 'error',
name: 'Workspace Sync Required',
message: LOCAL_WORKSPACE_SYNC_REQUIRED_MESSAGE,
});
t.deepEqual(semanticResult, {
type: 'error',
name: 'Workspace Sync Required',
message: LOCAL_WORKSPACE_SYNC_REQUIRED_MESSAGE,
});
});
test('doc_semantic_search should return empty array when nothing matches', async t => {
const ac = {
user: () => ({
workspace: () => ({
can: async () => true,
docs: async () => [],
}),
}),
} as unknown as PermissionAccess;
const models = {
workspace: {
get: async () => ({ id: 'workspace-1' }),
},
} as unknown as Models;
const contextService = {
matchWorkspaceAll: async () => [],
} as unknown as Parameters<typeof buildDocSearchGetter>[1];
const semanticTool = createDocSemanticSearchTool(
buildDocSearchGetter(ac, contextService, undefined, models).bind(null, {
user: 'user-1',
workspace: 'workspace-1',
})
);
const result = await semanticTool.execute?.({ query: 'hello' }, {});
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 PermissionAccess;
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: () => ({
workspace: () => ({
allowLocal: () => ({
can: async () => true,
}),
}),
}),
} as unknown as PermissionAccess;
const blobTool = createBlobReadTool(
buildBlobContentGetter(ac, null).bind(null, {
user: 'user-1',
workspace: 'workspace-1',
})
);
const result = await blobTool.execute?.({ blob_id: 'blob-1' }, {});
t.deepEqual(result, {
type: 'error',
name: 'Blob Read Failed',
message:
'Missing workspace, user, blob id, or copilot context for blob_read.',
});
});
packages/backend/server/src/__tests__/copilot/transcript-contract.spec.ts
+38 -166
View File
@@ -74,7 +74,7 @@ function createTranscriptPromptService() {
async function buildNativeTranscriptResult(input: any, runId: string) {
await input.onRunCreated?.({ runId, attempt: 1 });
const nativeInput = input.nativeInput;
const nativeInput = { input: input.inputSnapshot };
return {
nativeInput,
result: {
@@ -103,9 +103,7 @@ async function buildNativeTranscriptResult(input: any, runId: string) {
openQuestions: [],
blockers: [],
},
providerMeta: { provider: 'gemini', model: 'gemini-3.5-flash-lite' },
version: 'transcript-result-v1',
strategy: 'gemini',
},
};
}
@@ -127,7 +125,7 @@ function createSuccessfulTranscriptBridge(
});
yield {
type: 'action_done' as const,
actionId: 'transcript.audio.gemini',
actionId: 'transcript.audio',
actionVersion: 'v1',
status: 'succeeded' as const,
runId,
@@ -142,12 +140,9 @@ function createCopilotTranscriptionService(...deps: unknown[]) {
deps[0] as never,
deps[1] as never,
deps[2] as never,
deps[3] as never,
deps[4] as never,
deps[5] as never,
(deps[6] ?? {
assertQuotaOrByok: Sinon.stub().resolves(undefined),
}) as never,
(deps[6] ?? { assertRoute: Sinon.stub().resolves() }) as never,
(deps[7] ?? { publish: Sinon.stub() }) as never
);
}
@@ -221,47 +216,6 @@ test('settleTask unlocks ready transcript task result idempotently', async t =>
Sinon.assert.calledOnceWithExactly(settle, 'task-1');
});
test('settleTask checks copilot quota before unlocking ready task', async t => {
const payload = TranscriptPayloadSchema.parse({
normalizedTranscript: '00:00:05 A: Kickoff',
});
const settle = Sinon.stub().resolves({
id: 'task-1',
status: 'settled',
protectedResult: payload,
});
const assertQuotaOrByok = Sinon.stub().rejects(new Error('quota exceeded'));
const service = createCopilotTranscriptionService(
{
copilotTranscriptTask: {
getWithUser: Sinon.stub().resolves({
id: 'task-1',
status: 'ready',
protectedResult: payload,
}),
settle,
},
} as never,
{} as never,
{} as never,
{} as never,
{} as never,
{} as never,
{ assertQuotaOrByok } as never
);
await t.throwsAsync(
() => service.settleTask('user-1', 'workspace-1', 'task-1'),
{ message: /quota exceeded/ }
);
Sinon.assert.calledOnceWithMatch(assertQuotaOrByok, {
userId: 'user-1',
workspaceId: 'workspace-1',
featureKind: 'transcript',
});
Sinon.assert.notCalled(settle);
});
test('retryTask rejects ready transcript tasks', async t => {
const service = createCopilotTranscriptionService(
{
@@ -312,6 +266,7 @@ test('retryTask rejects settled transcript tasks', async t => {
test('retryTask reuses failed task and queues a new action attempt', async t => {
const queuedJobs: unknown[] = [];
const assertRoute = Sinon.stub().resolves();
const markRunning = Sinon.stub().resolves({
id: 'task-1',
status: 'running',
@@ -319,7 +274,6 @@ test('retryTask reuses failed task and queues a new action attempt', async t =>
const payload = TranscriptPayloadSchema.parse({
normalizedTranscript: '00:00:05 A: Kickoff',
summaryJson: null,
providerMeta: { provider: 'gemini', model: 'gemini-3.5-flash-lite' },
});
const service = createCopilotTranscriptionService(
{
@@ -327,7 +281,6 @@ test('retryTask reuses failed task and queues a new action attempt', async t =>
getWithUser: Sinon.stub().resolves({
id: 'task-1',
status: 'failed',
strategy: 'gemini',
actionRunId: 'run-failed',
protectedResult: payload,
}),
@@ -344,7 +297,8 @@ test('retryTask reuses failed task and queues a new action attempt', async t =>
resolveTranscriptionModel: Sinon.stub().resolves('gemini-3.5-flash-lite'),
} as never,
{} as never,
{} as never
{} as never,
{ assertRoute } as never
);
const result = await service.retryTask('user-1', 'workspace-1', 'task-1');
@@ -358,54 +312,24 @@ test('retryTask reuses failed task and queues a new action attempt', async t =>
retryOf: 'run-failed',
});
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 = createCopilotTranscriptionService(
Sinon.assert.calledOnceWithExactly(
assertRoute,
'transcript.audio',
{},
{
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-3.5-flash-lite'),
} as never,
{} as never,
{} as never,
{ assertQuotaOrByok } as never
user: 'user-1',
workspace: 'workspace-1',
featureKind: 'transcript',
builtInRouteId: 'Transcript audio structured',
}
);
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[] = [];
const queuedJobs: unknown[] = [];
const assertRoute = Sinon.stub().resolves();
const service = createCopilotTranscriptionService(
{
copilotTranscriptTask: {
@@ -432,7 +356,8 @@ for (const status of ['ready', 'settled']) {
),
} as never,
{} as never,
{} as never
{} as never,
{ assertRoute } as never
);
const result = await service.submitTask(
@@ -445,72 +370,25 @@ for (const status of ['ready', 'settled']) {
t.is(result.id, 'task-next');
t.like(createdTasks[0] as Record<string, unknown>, {
blobId: 'blob-1',
recipeId: 'transcript.audio.gemini',
recipeId: 'transcript.audio',
});
t.like(queuedJobs[0] as Record<string, unknown>, {
name: 'copilot.transcript.task.submit',
});
Sinon.assert.calledOnceWithExactly(
assertRoute,
'transcript.audio',
{},
{
user: 'user-1',
workspace: 'workspace-1',
featureKind: 'transcript',
builtInRouteId: 'Transcript audio structured',
}
);
});
}
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-3.5-flash-lite'
);
const service = createCopilotTranscriptionService(
{
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 = createCopilotTranscriptionService(
{
copilotTranscriptTask: {
getWithUser: Sinon.stub().resolves(null),
},
} as never,
{} as never,
{} as never,
{
resolveTranscriptionModel: Sinon.stub().resolves('gemini-3.5-flash-lite'),
} as never,
{} as never,
{} as never
);
await t.throwsAsync(
() =>
service.submitTask('user-1', 'workspace-1', 'blob-1', [], {
strategy: 'local-asr',
}),
{ message: /not available/ }
);
});
test('transcriptTask runs native transcript recipe through action bridge when available', async t => {
const payload = TranscriptPayloadSchema.parse({
sourceAudio: { blobId: 'blob-1', mimeType: 'audio/opus' },
@@ -559,28 +437,23 @@ test('transcriptTask runs native transcript recipe through action bridge when av
await service.transcriptTask({
taskId: 'task-1',
payload,
modelId: 'gemini-3.5-flash-lite',
});
t.like(bridgeInputs[0] as Record<string, unknown>, {
actionId: 'transcript.audio.gemini',
actionId: 'transcript.audio',
actionVersion: 'v1',
});
t.like(
(bridgeInputs[0] as { prepareStructuredRoutes: Record<string, unknown> })
.prepareStructuredRoutes,
{
stepId: 'transcribe',
modelId: 'gemini-3.5-flash-lite',
}
);
t.like((bridgeInputs[0] as { step: Record<string, unknown> }).step, {
slot: 'transcript.audio',
builtInRouteId: 'Transcript audio structured',
});
const messages = (
bridgeInputs[0] as {
prepareStructuredRoutes: {
step: {
messages: { content?: string; attachments?: unknown[] }[];
};
}
).prepareStructuredRoutes.messages;
).step.messages;
t.false(messages[0].content?.includes('data:image/png'));
t.like(JSON.parse(messages[0].content ?? '{}'), {
infos: [{ mimeType: 'audio/opus', index: 0 }],
@@ -630,7 +503,7 @@ test('transcriptTask fails task when native action bridge reports an error event
await buildNativeTranscriptResult(input, 'run-bridge');
yield {
type: 'error' as const,
actionId: 'transcript.audio.gemini',
actionId: 'transcript.audio',
actionVersion: 'v1',
status: 'failed' as const,
runId: 'run-bridge',
@@ -645,7 +518,6 @@ test('transcriptTask fails task when native action bridge reports an error event
service.transcriptTask({
taskId: 'task-1',
payload,
modelId: 'gemini-3.5-flash-lite',
}),
{ message: /native_failed/ }
);
packages/backend/server/src/__tests__/mocks/copilot.mock.ts
+83 -638
View File
@@ -1,656 +1,101 @@
import { randomBytes } from 'node:crypto';
import serverNativeModule from '@affine/server-native';
import type { ProviderMiddlewareConfig } from '../../plugins/copilot/config';
import {
CopilotChatOptions,
CopilotEmbeddingOptions,
type CopilotProviderModel,
CopilotProviderType,
CopilotStructuredOptions,
ModelConditions,
ModelFullConditions,
ModelOutputType,
PromptMessage,
StreamObject,
} from '../../plugins/copilot/providers';
import {
DEFAULT_DIMENSIONS,
OpenAIProvider,
} from '../../plugins/copilot/providers/openai';
import type { ProviderModelRuntimeContext } from '../../plugins/copilot/providers/provider-model-runtime';
import {
type CopilotProviderExecution,
createNativeExecutionDriverSpec,
type ProviderDriverSpec,
} from '../../plugins/copilot/providers/provider-runtime-contract';
import type { ProviderRuntimeContexts } from '../../plugins/copilot/runtime/provider-runtime-context';
import { sleep } from '../utils/utils';
import { EMBEDDING_DIMENSIONS } from '../../models';
const LLM_STREAM_END_MARKER = '__AFFINE_LLM_STREAM_END__';
const MOCK_NATIVE_TEXT = 'generate text to text';
const MOCK_NATIVE_STREAM_TEXT = 'generate text to text stream';
const STREAM_END = '__AFFINE_COPILOT_STREAM_END__';
const TEXT = 'generate text to text';
const STREAM_TEXT = 'generate text to text stream';
function mockUsage() {
return {
prompt_tokens: 1,
completion_tokens: 1,
total_tokens: 2,
};
function structuredValue(schema: unknown, key?: string): unknown {
if (!schema || typeof schema !== 'object') return TEXT;
const value = schema as Record<string, unknown>;
if (Array.isArray(value.enum)) return value.enum[0];
if (Array.isArray(value.anyOf)) return structuredValue(value.anyOf[0], key);
if (Array.isArray(value.oneOf)) return structuredValue(value.oneOf[0], key);
if (value.type === 'object') {
return Object.fromEntries(
Object.entries((value.properties as Record<string, unknown>) ?? {}).map(
([name, property]) => [name, structuredValue(property, name)]
)
);
}
if (value.type === 'array') return [structuredValue(value.items, key)];
if (value.type === 'boolean') return true;
if (value.type === 'number' || value.type === 'integer') return 1;
if (key === 'title') return 'Weekly Sync';
if (key === 'speaker' || key === 'a') return 'A';
if (key === 'text' || key === 'transcription' || key === 't') {
return 'Hello, everyone.';
}
return TEXT;
}
function buildMockDispatchResponse(model: string, text: string) {
return {
id: 'mock-dispatch',
model,
message: {
role: 'assistant',
content: [{ type: 'text', text }],
},
usage: mockUsage(),
finish_reason: 'stop',
};
}
function buildMockStructuredValue(schema: any, key?: string): any {
if (!schema || typeof schema !== 'object') {
return key === 'title' ? 'Weekly Sync' : MOCK_NATIVE_TEXT;
}
if (Array.isArray(schema.anyOf) && schema.anyOf.length > 0) {
return buildMockStructuredValue(schema.anyOf[0], key);
}
if (Array.isArray(schema.oneOf) && schema.oneOf.length > 0) {
return buildMockStructuredValue(schema.oneOf[0], key);
}
if (Array.isArray(schema.enum) && schema.enum.length > 0) {
return schema.enum[0];
}
switch (schema.type) {
case 'object': {
const properties =
schema.properties && typeof schema.properties === 'object'
? schema.properties
: {};
return Object.fromEntries(
Object.entries(properties).map(([key, value]) => [
key,
buildMockStructuredValue(value, key),
])
);
}
case 'array':
return [buildMockStructuredValue(schema.items, key)];
case 'boolean':
return true;
case 'number':
case 'integer':
switch (key) {
case 'durationMinutes':
return 45;
case 's':
return 30;
case 'e':
return 53;
default:
return 1;
}
case 'null':
return null;
case 'string':
default:
switch (key) {
case 'title':
return 'Weekly Sync';
case 'description':
return 'Send recap';
case 'owner':
return 'A';
case 'deadline':
return 'Friday';
case 'speaker':
case 'a':
return 'A';
case 'attendees':
return 'A';
case 'start':
return '00:00:42';
case 'end':
return '00:01:05';
case 'text':
case 'transcription':
case 't':
return 'Hello, everyone.';
case 'keyPoints':
return 'Reviewed launch status';
case 'decisions':
return 'Ship on Monday';
case 'openQuestions':
return 'Need final QA sign-off';
case 'blockers':
return 'Waiting on analytics';
case 'summary':
return 'Reviewed launch status';
default:
return MOCK_NATIVE_TEXT;
}
}
}
function parseFirstRoute(routesJson: string) {
const routes = JSON.parse(routesJson) as Array<{
provider_id?: string;
model?: string;
request?: {
model?: string;
operation?: string;
prompt?: string;
schema?: unknown;
function executionResult(input: { slot: string; request: unknown }) {
const request = input.request as Record<string, unknown>;
let result: unknown;
if (input.slot === 'index.embedding') {
const inputs = request.inputs as unknown[];
const dimensions =
(request.dimensions as number | undefined) ?? EMBEDDING_DIMENSIONS;
result = {
embeddings: inputs.map(() =>
Array.from({ length: dimensions }, (_, index) => index + 1)
),
};
}>;
return routes[0];
}
function buildMockStructuredResponse(model: string, schema: unknown) {
const output_json = buildMockStructuredValue(schema);
return {
id: 'mock-structured-dispatch',
model,
output_text: JSON.stringify(output_json),
output_json,
usage: mockUsage(),
finish_reason: 'stop',
};
}
function emitMockTextStream(
model: string,
callback: (error: Error | null, eventJson: string) => void
) {
callback(null, JSON.stringify({ type: 'message_start', model }));
for (const text of MOCK_NATIVE_STREAM_TEXT) {
callback(null, JSON.stringify({ type: 'text_delta', text }));
} else if (input.slot === 'search.rerank') {
const candidates = request.candidates as unknown[];
result = {
scores: candidates.map((_, index) => candidates.length - index),
};
} else if (input.slot === 'image.generate') {
result = {
images: [
{
data_base64: Buffer.from('generated image').toString('base64'),
media_type: 'image/jpeg',
},
],
};
} else if (input.slot.includes('structured')) {
const outputJson = structuredValue(request.schema);
result = {
output_json: outputJson,
output_text: JSON.stringify(outputJson),
};
} else {
result = { output_text: TEXT };
}
callback(
null,
JSON.stringify({
type: 'done',
finish_reason: 'stop',
usage: mockUsage(),
})
);
callback(null, LLM_STREAM_END_MARKER);
return JSON.stringify({ events: [], result });
}
export function installMockCopilotRuntime() {
const native = serverNativeModule as Record<string, any>;
const original = {
llmDispatchPrepared: native.llmDispatchPrepared,
llmDispatchPreparedStream: native.llmDispatchPreparedStream,
llmRenderBuiltInPrompt: native.llmRenderBuiltInPrompt,
llmRenderBuiltInSessionPrompt: native.llmRenderBuiltInSessionPrompt,
llmValidateJsonSchema: native.llmValidateJsonSchema,
llmStructuredDispatch: native.llmStructuredDispatch,
llmStructuredDispatchPrepared: native.llmStructuredDispatchPrepared,
llmEmbeddingDispatch: native.llmEmbeddingDispatch,
llmEmbeddingDispatchPrepared: native.llmEmbeddingDispatchPrepared,
llmRerankDispatch: native.llmRerankDispatch,
llmRerankDispatchPrepared: native.llmRerankDispatchPrepared,
llmImageDispatchPrepared: native.llmImageDispatchPrepared,
runNativeActionRecipePreparedStream:
native.runNativeActionRecipePreparedStream,
};
native.llmDispatchPrepared = (routesJson: string) => {
const route = parseFirstRoute(routesJson);
return JSON.stringify({
provider_id: route?.provider_id ?? 'mock-provider',
response: buildMockDispatchResponse(
route?.request?.model ?? route?.model ?? 'test',
MOCK_NATIVE_TEXT
),
});
};
native.llmDispatchPreparedStream = (
routesJson: string,
callback: (error: Error | null, eventJson: string) => void
const prototype = serverNativeModule.BackendRuntime.prototype;
const execute = prototype.executeCopilot;
const stream = prototype.executeCopilotStream;
prototype.executeCopilot = async input => executionResult(input);
prototype.executeCopilotStream = async (
_input,
_maxSteps,
callback,
_toolCallback
) => {
const route = parseFirstRoute(routesJson);
emitMockTextStream(
route?.request?.model ?? route?.model ?? 'test',
callback
callback(null, JSON.stringify({ type: 'message_start', model: 'test' }));
for (const text of STREAM_TEXT) {
callback(null, JSON.stringify({ type: 'text_delta', text }));
}
callback(
null,
JSON.stringify({
type: 'done',
finish_reason: 'stop',
usage: { prompt_tokens: 1, completion_tokens: 1, total_tokens: 2 },
})
);
callback(null, STREAM_END);
return { abort() {} };
};
native.llmStructuredDispatch = (
_protocol: string,
_backendConfigJson: string,
requestJson: string
) => {
const request = JSON.parse(requestJson) as {
model?: string;
schema?: unknown;
};
return JSON.stringify(
buildMockStructuredResponse(request.model ?? 'test', request.schema)
);
};
native.llmStructuredDispatchPrepared = (routesJson: string) => {
const route = parseFirstRoute(routesJson);
return JSON.stringify({
provider_id: route?.provider_id ?? 'mock-provider',
response: buildMockStructuredResponse(
route?.request?.model ?? route?.model ?? 'test',
route?.request?.schema
),
});
};
native.llmValidateJsonSchema = (_schema: unknown, value: unknown) => value;
native.llmEmbeddingDispatch = (
_protocol: string,
_backendConfigJson: string,
requestJson: string
) => {
const request = JSON.parse(requestJson) as {
model?: string;
dimensions?: number;
};
const length = request.dimensions ?? DEFAULT_DIMENSIONS;
return JSON.stringify({
model: request.model ?? 'test',
embeddings: [
Array.from({ length }, (_value, index) => (index % 128) + 1),
],
usage: { prompt_tokens: 1, total_tokens: 1 },
});
};
native.llmEmbeddingDispatchPrepared = (routesJson: string) => {
const route = parseFirstRoute(routesJson);
const response = JSON.parse(
native.llmEmbeddingDispatch(
'',
'',
JSON.stringify(route?.request ?? { model: route?.model ?? 'test' })
)
) as Record<string, unknown>;
return JSON.stringify({
provider_id: route?.provider_id ?? 'mock-provider',
response,
});
};
native.llmRerankDispatch = (
_protocol: string,
_backendConfigJson: string,
requestJson: string
) => {
const request = JSON.parse(requestJson) as {
model?: string;
candidates?: unknown[];
};
const candidateCount = request.candidates?.length ?? 0;
return JSON.stringify({
model: request.model ?? 'test',
scores: Array.from(
{ length: candidateCount },
(_value, index) => candidateCount - index
),
});
};
native.llmRerankDispatchPrepared = (routesJson: string) => {
const route = parseFirstRoute(routesJson);
const response = JSON.parse(
native.llmRerankDispatch(
'',
'',
JSON.stringify(route?.request ?? { model: route?.model ?? 'test' })
)
) as Record<string, unknown>;
return JSON.stringify({
provider_id: route?.provider_id ?? 'mock-provider',
response,
});
};
native.llmImageDispatchPrepared = (routesJson: string) => {
const route = parseFirstRoute(routesJson);
const model = route?.request?.model ?? route?.model ?? 'test-image';
const images = [
{
url: `https://example.com/${model}.jpg`,
media_type: 'image/jpeg',
},
];
if (route?.request?.operation === 'edit' && route.request.prompt) {
images.push({
url: `https://example.com/generated/${encodeURIComponent(route.request.prompt)}.jpg`,
media_type: 'image/jpeg',
});
}
return JSON.stringify({
provider_id: route?.provider_id ?? 'mock-provider',
response: {
images,
},
});
};
native.runNativeActionRecipePreparedStream = (
input: {
recipeId: string;
recipeVersion?: string;
input?: Record<string, any>;
},
callback: (error: Error | null, eventJson: string) => void
) => {
const version = input.recipeVersion ?? 'v1';
const result = input.recipeId.startsWith('image.filter.')
? {
url: `https://example.com/${input.recipeId}.jpg`,
}
: MOCK_NATIVE_STREAM_TEXT;
const attachmentEvent = input.recipeId.startsWith('image.filter.')
? [
{
type: 'attachment',
actionId: input.recipeId,
actionVersion: version,
status: 'running',
attachment: result,
},
]
: [];
const events = [
{
type: 'action_start',
actionId: input.recipeId,
actionVersion: version,
status: 'running',
},
{
type: 'step_start',
actionId: input.recipeId,
actionVersion: version,
stepId: 'generate',
status: 'running',
},
...attachmentEvent,
{
type: 'step_end',
actionId: input.recipeId,
actionVersion: version,
stepId: 'generate',
status: 'running',
},
{
type: 'action_done',
actionId: input.recipeId,
actionVersion: version,
status: 'succeeded',
result,
trace: {
actionId: input.recipeId,
actionVersion: version,
status: 'succeeded',
lightweight: [
{ type: 'action_start', status: 'running' },
{ type: 'action_trace', status: 'succeeded' },
],
},
},
];
for (const event of events) {
callback(null, JSON.stringify(event));
}
callback(null, LLM_STREAM_END_MARKER);
return { abort() {} };
};
return () => {
Object.assign(native, original);
prototype.executeCopilot = execute;
prototype.executeCopilotStream = stream;
};
}
export class MockCopilotProvider extends OpenAIProvider {
private runtimeHostOverride?: ProviderRuntimeContexts;
protected override resolveModelRuntimeContext(): ProviderModelRuntimeContext {
const providerType = this.type as CopilotProviderType;
return {
type: providerType,
backendKind:
providerType === CopilotProviderType.Gemini
? 'gemini_api'
: 'openai_responses',
};
}
override getDriverSpec(): ProviderDriverSpec {
const spec = super.getDriverSpec();
return {
...spec,
image: {
prepareMessages: async messages => messages,
},
};
}
private resolveMockModelId(
cond: Pick<ModelFullConditions, 'modelId' | 'outputType'>
) {
if (cond.modelId === 'test') {
return 'gpt-5-mini';
}
if (cond.modelId === 'test-image') {
return 'gpt-image-1';
}
return cond.modelId;
}
private normalizeMockConditions(
cond: ModelFullConditions
): ModelFullConditions {
const modelId = this.resolveMockModelId(cond);
return modelId === cond.modelId ? cond : { ...cond, modelId };
}
protected override createDriverSpec(spec: ProviderDriverSpec) {
return createNativeExecutionDriverSpec(spec, {
createBackendConfig: spec.createBackendConfig,
mapError: spec.mapError,
checkParams: input => this.checkParams(input),
selectModel: (cond, execution) => this.selectModel(cond, execution),
getTools: this.getTools.bind(this),
getActiveProviderMiddleware: this.getActiveProviderMiddleware.bind(this),
});
}
override async match(
cond: ModelFullConditions = {},
execution?: CopilotProviderExecution
) {
return await super.match(this.normalizeMockConditions(cond), execution);
}
override resolveModel(
modelId: string,
execution?: CopilotProviderExecution
): CopilotProviderModel | undefined {
const resolvedModelId = this.resolveMockModelId({ modelId });
return resolvedModelId
? super.resolveModel(resolvedModelId, execution)
: undefined;
}
override selectModel(
cond: ModelFullConditions,
execution?: CopilotProviderExecution
): CopilotProviderModel {
return super.selectModel(this.normalizeMockConditions(cond), execution);
}
override checkParams(input: Parameters<OpenAIProvider['checkParams']>[0]) {
return super.checkParams({
...input,
cond: this.normalizeMockConditions(input.cond),
});
}
override getActiveProviderMiddleware(): ProviderMiddlewareConfig {
return {};
}
overrideRuntimeHost(runtimeHost: ProviderRuntimeContexts) {
if (!this.runtimeHostOverride) {
const runtimeHostOverride: ProviderRuntimeContexts = {
...runtimeHost,
run: {
...runtimeHost.run,
text: this.text.bind(this),
streamText: this.streamTextRuntime.bind(this),
streamObject: this.streamObjectRuntime.bind(this),
structured: this.structure.bind(this),
embedding: this.embedding.bind(this),
},
};
this.runtimeHostOverride = runtimeHostOverride;
}
return this.runtimeHostOverride;
}
private async *streamTextRuntime(
cond: ModelConditions,
messages: PromptMessage[],
options?: CopilotChatOptions
): AsyncIterableIterator<string> {
yield* this.streamText(cond, messages, options);
}
private async *streamObjectRuntime(
cond: ModelConditions,
messages: PromptMessage[],
options?: CopilotChatOptions
): AsyncIterableIterator<StreamObject> {
yield* this.streamObject(cond, messages, options);
}
async text(
cond: ModelConditions,
messages: PromptMessage[],
options: CopilotChatOptions = {}
): Promise<string> {
const fullCond = {
...cond,
outputType: ModelOutputType.Text,
};
await this.checkParams({
messages,
cond: fullCond,
options,
});
// make some time gap for history test case
await sleep(100);
return 'generate text to text';
}
async *streamText(
cond: ModelConditions,
messages: PromptMessage[],
options: CopilotChatOptions = {}
): AsyncIterable<string> {
const fullCond = { ...cond, outputType: ModelOutputType.Text };
await this.checkParams({
messages,
cond: fullCond,
options,
});
// make some time gap for history test case
await sleep(100);
const result = 'generate text to text stream';
for (const message of result) {
yield message;
if (options.signal?.aborted) {
break;
}
}
}
async structure(
cond: ModelConditions,
messages: PromptMessage[],
options: CopilotStructuredOptions = {}
): Promise<string> {
const fullCond = { ...cond, outputType: ModelOutputType.Structured };
await this.checkParams({
messages,
cond: fullCond,
options,
});
// make some time gap for history test case
await sleep(100);
return 'generate text to text';
}
// ====== text to embedding ======
async embedding(
cond: ModelConditions,
messages: string | string[],
options: CopilotEmbeddingOptions = { dimensions: DEFAULT_DIMENSIONS }
): Promise<number[][]> {
messages = Array.isArray(messages) ? messages : [messages];
const fullCond = { ...cond, outputType: ModelOutputType.Embedding };
await this.checkParams({
embeddings: messages,
cond: fullCond,
options,
});
// make some time gap for history test case
await sleep(100);
return [
Array.from(randomBytes(options.dimensions ?? DEFAULT_DIMENSIONS)).map(
v => v % 128
),
];
}
async *streamObject(
cond: ModelConditions,
messages: PromptMessage[],
options: CopilotChatOptions = {}
): AsyncIterable<StreamObject> {
const fullCond = { ...cond, outputType: ModelOutputType.Object };
await this.checkParams({
messages,
cond: fullCond,
options,
});
// make some time gap for history test case
await sleep(100);
const result = 'generate text to object stream';
for (const data of result) {
yield { type: 'text-delta', textDelta: data } as const;
if (options.signal?.aborted) {
break;
}
}
}
}
packages/backend/server/src/__tests__/mocks/index.ts
+1 -3
View File
@@ -1,11 +1,10 @@
export { createFactory } from './factory';
export * from './prompt-service.mock';
export * from './team-workspace.mock';
export * from './user.mock';
export * from './workspace.mock';
export * from './workspace-user.mock';
import { installMockCopilotRuntime, MockCopilotProvider } from './copilot.mock';
import { installMockCopilotRuntime } from './copilot.mock';
import { MockDocMeta } from './doc-meta.mock';
import { MockDocSnapshot } from './doc-snapshot.mock';
import { MockDocUser } from './doc-user.mock';
@@ -31,7 +30,6 @@ export const Mockers = {
export {
installMockCopilotRuntime,
MockCopilotProvider,
MockEventBus,
MockJobModule,
MockJobQueue,
packages/backend/server/src/__tests__/mocks/prompt-service.mock.ts
-110
View File
@@ -1,110 +0,0 @@
import { Injectable } from '@nestjs/common';
import { CopilotPromptInvalid } from '../../base';
import { llmGetBuiltInPromptSpec, llmRenderBuiltInPrompt } from '../../native';
import { PromptService } from '../../plugins/copilot/prompt';
import type { Prompt } from '../../plugins/copilot/prompt/spec';
import type {
PromptConfig,
PromptMessage,
} from '../../plugins/copilot/providers/types';
@Injectable()
export class TestingPromptService extends PromptService {
private readonly customPrompts = new Map<string, Prompt>();
private readonly builtInPromptOverrides = new Map<string, Prompt>();
reset() {
this.customPrompts.clear();
this.builtInPromptOverrides.clear();
}
async set(
name: string,
model: string,
messages: PromptMessage[],
config?: PromptConfig | null,
extraConfig?: { optionalModels: string[] }
) {
this.assertCustomPromptName(name);
const existing = this.customPrompts.get(name);
this.customPrompts.set(name, {
name,
model,
action: existing?.action,
optionalModels: existing?.optionalModels?.length
? [...existing.optionalModels, ...(extraConfig?.optionalModels ?? [])]
: extraConfig?.optionalModels,
config: config ? structuredClone(config) : undefined,
messages: this.cloneMessages(messages),
});
}
async overrideBuiltIn(
name: string,
data: {
messages?: PromptMessage[];
model?: string;
config?: PromptConfig | null;
}
) {
const current = this.loadBuiltInPrompt(name);
if (!current) {
throw new CopilotPromptInvalid(
`Built-in prompt ${name} not found in native catalog`
);
}
const { config, messages, model } = data;
const next = this.clonePrompt(current);
if (model !== undefined) {
next.model = model;
}
if (config === null) {
next.config = undefined;
} else if (config !== undefined) {
next.config = structuredClone(config);
}
if (messages) {
next.messages = this.cloneMessages(messages);
}
this.builtInPromptOverrides.set(name, next);
}
protected override lookupCompatPrompt(name: string) {
return (
this.builtInPromptOverrides.get(name) ??
this.customPrompts.get(name) ??
null
);
}
private assertCustomPromptName(name: string) {
if (this.loadBuiltInPrompt(name)) {
throw new CopilotPromptInvalid(
`Built-in prompt ${name} is owned by native catalog`
);
}
}
private loadBuiltInPrompt(name: string): Prompt | null {
const spec = llmGetBuiltInPromptSpec(name);
if (!spec) return null;
const prompt = llmRenderBuiltInPrompt({ name, renderParams: {} });
return {
name: spec.name,
action: spec.action,
model: spec.model,
optionalModels: spec.optionalModels,
config: spec.config,
messages: prompt.messages.map(message => ({
role: message.role,
content: message.content,
...(message.params ? { params: message.params } : {}),
})),
};
}
}
packages/backend/server/src/__tests__/models/__snapshots__/copilot-session.spec.ts.md
+3
View File
@@ -185,6 +185,7 @@ Generated by [AVA](https://avajs.dev).
result: 'success',
sessionType: 'forked',
update: {
promptAction: null,
promptName: 'test-prompt',
},
},
@@ -199,6 +200,7 @@ Generated by [AVA](https://avajs.dev).
result: 'success',
sessionType: 'regular',
update: {
promptAction: null,
promptName: 'test-prompt',
},
},
@@ -206,6 +208,7 @@ Generated by [AVA](https://avajs.dev).
result: 'rejected',
sessionType: 'regular',
update: {
promptAction: 'edit',
promptName: 'action-prompt',
},
},
packages/backend/server/src/__tests__/models/__snapshots__/copilot-session.spec.ts.snap
BIN
View File
Binary file not shown.
packages/backend/server/src/__tests__/models/copilot-context.spec.ts
-3
View File
@@ -48,9 +48,6 @@ let docId = 'doc1';
test.beforeEach(async t => {
await t.context.module.initTestingDB();
await t.context.db.aiPrompt.create({
data: { name: 'prompt-name', model: 'gpt-5-mini', action: null },
});
user = await t.context.user.create({
email: 'test@affine.pro',
});
packages/backend/server/src/__tests__/models/copilot-session.spec.ts
+36 -56
View File
@@ -57,18 +57,6 @@ const TEST_PROMPTS = {
} as const;
// Helper functions
const createTestPrompts = async (
_copilotSession: CopilotSessionModel,
db: PrismaClient
) => {
await db.aiPrompt.create({
data: { name: TEST_PROMPTS.NORMAL, model: 'gpt-5-mini', action: null },
});
await db.aiPrompt.create({
data: { name: TEST_PROMPTS.ACTION, model: 'gpt-5-mini', action: 'edit' },
});
};
const createTestSession = async (
t: ExecutionContext<Context>,
overrides: Partial<{
@@ -121,7 +109,6 @@ const addMessagesToSession = async (
await copilotSession.updateMessages({
sessionId,
userId: user.id,
prompt: { model: 'gpt-5-mini' },
messages: [
{
role: 'user',
@@ -154,9 +141,7 @@ const createSessionWithMessages = async (
type UpdateData = Omit<UpdateChatSessionOptions, 'userId' | 'sessionId'>;
test('should list and filter session type', async t => {
const { copilotSession, db } = t.context;
await createTestPrompts(copilotSession, db);
const { copilotSession } = t.context;
const docId = 'doc-id-1';
await createTestSession(t, { sessionId: randomUUID() });
@@ -206,7 +191,7 @@ test('should list and filter session type', async t => {
docSessions.toSorted((a, b) =>
a.promptName.localeCompare(b.promptName)
),
['id', 'userId', 'workspaceId', 'createdAt', 'updatedAt', 'tokenCost']
['id', 'userId', 'workspaceId', 'createdAt', 'updatedAt']
),
'doc sessions should only include sessions with matching docId'
);
@@ -232,8 +217,7 @@ test('should list and filter session type', async t => {
});
test('should validate session prompt compatibility', async t => {
const { copilotSession, db } = t.context;
await createTestPrompts(copilotSession, db);
const { copilotSession } = t.context;
const sessionTypes = [
{ name: 'workspace', session: { docId: null, pinned: false } },
@@ -288,8 +272,6 @@ test('should validate session prompt compatibility', async t => {
test('should pin and unpin sessions', async t => {
const { copilotSession, db } = t.context;
await createTestPrompts(copilotSession, db);
const firstSessionId = 'first-session-id';
const secondSessionId = 'second-session-id';
const thirdSessionId = 'third-session-id';
@@ -368,7 +350,6 @@ test('should pin and unpin sessions', async t => {
test('should handle session updates and type conversions', async t => {
const { copilotSession, db } = t.context;
await createTestPrompts(copilotSession, db);
const sessionId = randomUUID();
const actionSessionId = randomUUID();
@@ -414,7 +395,11 @@ test('should handle session updates and type conversions', async t => {
sessionId: forkedSessionId,
updates: [
{ pinned: true, expected: 'allow' },
{ promptName: TEST_PROMPTS.NORMAL, expected: 'allow' },
{
promptName: TEST_PROMPTS.NORMAL,
promptAction: null,
expected: 'allow',
},
{ docId: 'new-doc', expected: 'reject' },
],
},
@@ -422,8 +407,16 @@ test('should handle session updates and type conversions', async t => {
{
sessionId,
updates: [
{ promptName: TEST_PROMPTS.NORMAL, expected: 'allow' },
{ promptName: TEST_PROMPTS.ACTION, expected: 'reject' },
{
promptName: TEST_PROMPTS.NORMAL,
promptAction: null,
expected: 'allow',
},
{
promptName: TEST_PROMPTS.ACTION,
promptAction: 'edit',
expected: 'reject',
},
{ promptName: 'non-existent-prompt', expected: 'reject' },
],
},
@@ -517,7 +510,6 @@ test('should handle session updates and type conversions', async t => {
test('should handle session queries, ordering, and filtering', async t => {
const { copilotSession, db } = t.context;
await createTestPrompts(copilotSession, db);
const docId = randomUUID();
const sessionIds: string[] = [];
@@ -764,7 +756,6 @@ test('should handle session queries, ordering, and filtering', async t => {
test('should handle fork and session attachment operations', async t => {
const { copilotSession } = t.context;
await createTestPrompts(copilotSession, t.context.db);
const parentSessionId = randomUUID();
const docId = randomUUID();
@@ -812,7 +803,7 @@ test('should handle fork and session attachment operations', async t => {
pinned: forkConfig.pinned,
title: null,
parentSessionId,
prompt: { name: TEST_PROMPTS.NORMAL, action: null, model: 'gpt-5-mini' },
prompt: { name: TEST_PROMPTS.NORMAL, action: null },
messages: [
{
role: 'user',
@@ -925,7 +916,6 @@ test('should handle fork and session attachment operations', async t => {
test('should cleanup empty sessions correctly', async t => {
const { copilotSession, db } = t.context;
await createTestPrompts(copilotSession, db);
const oneDayAgo = new Date(Date.now() - 24 * 60 * 60 * 1000);
const twoHoursAgo = new Date(Date.now() - 2 * 60 * 60 * 1000);
@@ -971,20 +961,23 @@ test('should cleanup empty sessions correctly', async t => {
);
const result = await copilotSession.cleanupEmptySessions(oneDayAgo);
const orderedSessionIds = [
...neverUsedSessionIds,
...emptySessionIds,
recentSessionId,
sessionWithMsgId,
];
const remainingSessions = await db.aiSession.findMany({
where: {
id: {
in: [
...neverUsedSessionIds,
...emptySessionIds,
recentSessionId,
sessionWithMsgId,
],
},
id: { in: orderedSessionIds },
},
select: { id: true, deletedAt: true, pinned: true },
});
remainingSessions.sort(
(left, right) =>
orderedSessionIds.indexOf(left.id) - orderedSessionIds.indexOf(right.id)
);
t.snapshot(
{
@@ -1005,15 +998,13 @@ test('should cleanup empty sessions correctly', async t => {
);
});
test('should append durable message and account durable costs', async t => {
test('should append durable message and account message cost', async t => {
const { copilotSession, db } = t.context;
await createTestPrompts(copilotSession, db);
const { sessionId } = await createTestSession(t);
const appended = await copilotSession.appendMessage({
sessionId,
userId: user.id,
prompt: { model: 'gpt-5-mini' },
message: {
role: 'user',
content: 'hello durable world',
@@ -1024,18 +1015,16 @@ test('should append durable message and account durable costs', async t => {
const afterAppend = await db.aiSession.findUniqueOrThrow({
where: { id: sessionId },
select: { messageCost: true, tokenCost: true },
select: { messageCost: true },
});
t.truthy(appended.id);
t.is(afterAppend.messageCost, 1);
t.true(afterAppend.tokenCost > 0);
t.deepEqual(appended.params, { foo: 'bar' });
const appendedBare = await copilotSession.appendMessage({
sessionId,
userId: user.id,
prompt: { model: 'gpt-5-mini' },
message: {
role: 'assistant',
content: 'assistant reply',
@@ -1070,14 +1059,12 @@ test('should append durable message and account durable costs', async t => {
});
test('should count action runs without double-counting legacy action sessions', async t => {
const { copilotSession, db, models } = t.context;
await createTestPrompts(copilotSession, db);
const { copilotSession, models } = t.context;
const regular = await createTestSession(t);
await copilotSession.appendMessage({
sessionId: regular.sessionId,
userId: user.id,
prompt: { model: 'gpt-5-mini' },
message: {
role: 'user',
content: 'regular message',
@@ -1124,8 +1111,7 @@ test('should count action runs without double-counting legacy action sessions',
userId: user.id,
workspaceId: workspace.id,
blobId: 'audio-1',
strategy: 'gemini',
recipeId: 'transcript.audio.gemini',
recipeId: 'transcript.audio',
recipeVersion: 'v1',
});
await models.copilotTranscriptTask.complete(transcriptTask.id, {
@@ -1146,14 +1132,12 @@ test('should count action runs without double-counting legacy action sessions',
});
test('should exclude BYOK provider usage from copilot quota cost', async t => {
const { copilotSession, db, models } = t.context;
await createTestPrompts(copilotSession, db);
const { copilotSession, models } = t.context;
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',
@@ -1163,7 +1147,6 @@ test('should exclude BYOK provider usage from copilot quota cost', async t => {
const secondMessage = await copilotSession.appendMessage({
sessionId: regular.sessionId,
userId: user.id,
prompt: { model: 'gpt-5-mini' },
message: {
role: 'user',
content: 'second BYOK message',
@@ -1173,7 +1156,6 @@ test('should exclude BYOK provider usage from copilot quota cost', async t => {
await copilotSession.appendMessage({
sessionId: regular.sessionId,
userId: user.id,
prompt: { model: 'gpt-5-mini' },
message: {
role: 'user',
content: 'quota-backed message',
@@ -1194,8 +1176,7 @@ test('should exclude BYOK provider usage from copilot quota cost', async t => {
userId: user.id,
workspaceId: workspace.id,
blobId: 'pending-audio',
strategy: 'gemini',
recipeId: 'transcript.audio.gemini',
recipeId: 'transcript.audio',
recipeVersion: 'v1',
});
await models.copilotUsage.create({
@@ -1251,7 +1232,6 @@ test('should exclude BYOK provider usage from copilot quota cost', async t => {
test('should get sessions for title generation correctly', async t => {
const { copilotSession, db } = t.context;
await createTestPrompts(copilotSession, db);
// create valid sessions with messages
const sessionIds: string[] = [randomUUID(), randomUUID()];
packages/backend/server/src/__tests__/utils/copilot.ts
-5
View File
@@ -734,7 +734,6 @@ type ChatMessage = {
type History = {
sessionId: string;
pinned: boolean;
tokens: number;
action: string | null;
createdAt: string;
messages: ChatMessage[];
@@ -773,7 +772,6 @@ export async function getHistories(
histories(docId: $docId, options: $options) {
sessionId
pinned
tokens
action
createdAt
messages {
@@ -811,7 +809,6 @@ export async function getWorkspaceSessions(
histories(docId: null, options: $options) {
sessionId
pinned
tokens
action
createdAt
messages {
@@ -858,7 +855,6 @@ export async function getDocSessions(
histories(docId: $docId, options: $options) {
sessionId
pinned
tokens
action
createdAt
messages {
@@ -912,7 +908,6 @@ export async function getPinnedSessions(
}) {
sessionId
pinned
tokens
action
createdAt
messages {
packages/backend/server/src/__tests__/utils/testing-module.ts
+1
View File
@@ -127,6 +127,7 @@ export async function createTestingModule(
},
},
copilot: {
enabled: true,
storage: {
provider: 'assetpack',
bucket: 'copilot',
packages/backend/server/src/core/backend-runtime/__tests__/provider.spec.ts
+51 -4
View File
@@ -1,41 +1,88 @@
import { generateKeyPairSync } from 'node:crypto';
import test from 'ava';
import Sinon from 'sinon';
import type { Config } from '../../../base';
import { BackendRuntimeProvider } from '../provider';
const privateKey = generateKeyPairSync('ec', {
namedCurve: 'P-256',
}).privateKey.export({ format: 'pem', type: 'pkcs8' }) as string;
const config = { crypto: { privateKey } } as Config;
test('backend-runtime provider starts once, runs migrations once, and reports health', async t => {
const provider = new BackendRuntimeProvider();
const provider = new BackendRuntimeProvider(config);
const runtime = {
start: Sinon.stub().resolves(),
stop: Sinon.stub().resolves(),
runMigrations: Sinon.stub().resolves(),
reloadConfig: Sinon.stub().resolves(),
health: Sinon.stub().resolves({
started: true,
databaseConnected: true,
}),
};
(provider as any).runtime = runtime;
(provider as unknown as { runtime: typeof runtime }).runtime = runtime;
await provider.start();
await provider.start();
await provider.onConfigChanged({ updates: { mailer: {} } });
await provider.onConfigChanged({ updates: { copilot: {} } });
const health = await provider.health();
await provider.stop();
t.is(runtime.start.callCount, 2);
t.is(runtime.runMigrations.callCount, 1);
t.true(runtime.reloadConfig.calledOnceWithExactly(privateKey));
t.true(health.databaseConnected);
t.is(runtime.stop.callCount, 1);
});
test('backend-runtime provider measures explicit typed methods', async t => {
const provider = new BackendRuntimeProvider();
const provider = new BackendRuntimeProvider(config);
const runtime = {
cleanupExpiredRuntimeStates: Sinon.stub().resolves(3),
assertCopilotRoute: Sinon.stub().resolves(),
};
(provider as any).runtime = runtime;
(provider as unknown as { runtime: typeof runtime }).runtime = runtime;
const result = await provider.cleanupExpiredRuntimeStates(1000);
const routeInput = {
slot: 'transcript.audio',
access: {
routeAllowed: true,
managedTier: 'Standard' as const,
serverByok: true,
localByok: false,
},
};
await provider.assertCopilotRoute(routeInput);
t.is(result, 3);
t.true(runtime.cleanupExpiredRuntimeStates.calledOnceWithExactly(1000));
t.true(runtime.assertCopilotRoute.calledOnceWithExactly(routeInput));
});
test('backend-runtime provider aborts a stream handle that resolves after iterator cancellation', async t => {
const provider = new BackendRuntimeProvider(config);
const abort = Sinon.stub();
let resolveHandle!: (handle: { abort: () => void }) => void;
const runtime = {
executeCopilotStream: Sinon.stub().returns(
new Promise<{ abort: () => void }>(resolve => {
resolveHandle = resolve;
})
),
};
(provider as unknown as { runtime: typeof runtime }).runtime = runtime;
const stream = provider.streamCopilot({} as never, async () => '', {
maxSteps: 1,
});
await stream.return?.();
resolveHandle({ abort });
await Promise.resolve();
t.true(abort.calledOnce);
});
packages/backend/server/src/core/backend-runtime/provider.ts
+219 -2
View File
@@ -3,13 +3,76 @@ import {
Logger,
type OnApplicationBootstrap,
type OnApplicationShutdown,
Optional,
} from '@nestjs/common';
import { Config, OnEvent } from '../../base';
import { wrapCallMetric } from '../../base/metrics';
import { BackendRuntime, type BackendRuntimeHealth } from '../../native';
import {
BackendRuntime,
type BackendRuntimeHealth,
type ByokLocalLeaseOutput,
type ByokProbeResultOutput,
type ByokProfileOutput,
type CopilotExecuteInput,
type CopilotRouteCheckInput,
type CreateByokLocalLeaseInput,
type CreateByokProfileInput,
type ProbeByokDraftInput,
type ProbeByokProfileInput,
type ReorderByokProfilesInput,
type ReplaceByokProfileInput,
type RotateByokCredentialInput,
} from '../../native';
type RuntimeInstance = InstanceType<typeof BackendRuntime>;
class RuntimeEventStream<T> implements AsyncIterableIterator<T> {
private readonly values: T[] = [];
private readonly readers: Array<(result: IteratorResult<T>) => void> = [];
private ended = false;
private abort?: () => void;
attach(abort: () => void) {
if (this.ended) {
abort();
return;
}
this.abort = abort;
}
push(value?: T) {
if (this.ended) return;
if (value === undefined) {
this.ended = true;
for (const reader of this.readers.splice(0)) {
reader({ value: undefined, done: true });
}
return;
}
const reader = this.readers.shift();
if (reader) reader({ value, done: false });
else this.values.push(value);
}
[Symbol.asyncIterator]() {
return this;
}
async next(): Promise<IteratorResult<T>> {
const value = this.values.shift();
if (value !== undefined) return { value, done: false };
if (this.ended) return { value: undefined, done: true };
return await new Promise(resolve => this.readers.push(resolve));
}
async return(): Promise<IteratorResult<T>> {
this.abort?.();
this.push();
return { value: undefined, done: true };
}
}
export type RuntimeQuotaTargetDomainInput = {
domain: string;
count: number;
@@ -196,9 +259,13 @@ export class BackendRuntimeProvider
implements OnApplicationBootstrap, OnApplicationShutdown
{
private readonly logger = new Logger(BackendRuntimeProvider.name);
private readonly runtime: RuntimeInstance = new BackendRuntime();
private readonly runtime: RuntimeInstance;
private migrationsStarted = false;
constructor(@Optional() private readonly config?: Config) {
this.runtime = new BackendRuntime(this.config?.crypto.privateKey);
}
async onApplicationBootstrap() {
await this.start();
}
@@ -219,6 +286,14 @@ export class BackendRuntimeProvider
this.logger.log('backend runtime stopped');
}
@OnEvent('config.changed')
async onConfigChanged({ updates }: Events['config.changed']) {
if (!updates.copilot && !updates.crypto && !updates.db) {
return;
}
await this.runtime.reloadConfig(this.config?.crypto.privateKey);
}
async health(): Promise<BackendRuntimeHealth> {
return await this.runtime.health();
}
@@ -298,6 +373,148 @@ export class BackendRuntimeProvider
);
}
async listByokProfiles(workspaceId: string): Promise<ByokProfileOutput[]> {
return await this.measured('listByokProfiles', runtime =>
runtime.listByokProfiles(workspaceId)
);
}
async createByokProfile(
input: CreateByokProfileInput
): Promise<ByokProfileOutput> {
return await this.measured('createByokProfile', runtime =>
runtime.createByokProfile(input)
);
}
async replaceByokProfile(
input: ReplaceByokProfileInput
): Promise<ByokProfileOutput> {
return await this.measured('replaceByokProfile', runtime =>
runtime.replaceByokProfile(input)
);
}
async rotateByokCredential(
input: RotateByokCredentialInput
): Promise<ByokProfileOutput> {
return await this.measured('rotateByokCredential', runtime =>
runtime.rotateByokCredential(input)
);
}
async probeByokProfile(
input: ProbeByokProfileInput
): Promise<ByokProbeResultOutput> {
return await this.measured('probeByokProfile', runtime =>
runtime.probeByokProfile(input)
);
}
async probeByokDraft(
input: ProbeByokDraftInput
): Promise<ByokProbeResultOutput> {
return await this.measured('probeByokDraft', runtime =>
runtime.probeByokDraft(input)
);
}
async deleteByokProfile(workspaceId: string, profileId: string) {
return await this.measured('deleteByokProfile', runtime =>
runtime.deleteByokProfile(workspaceId, profileId)
);
}
async reorderByokProfiles(
input: ReorderByokProfilesInput
): Promise<ByokProfileOutput[]> {
return await this.measured('reorderByokProfiles', runtime =>
runtime.reorderByokProfiles(input)
);
}
async createByokLocalLease(
input: CreateByokLocalLeaseInput
): Promise<ByokLocalLeaseOutput> {
return await this.measured('createByokLocalLease', runtime =>
runtime.createByokLocalLease(input)
);
}
async executeCopilot(input: CopilotExecuteInput) {
const output = await this.measured('executeCopilot', runtime =>
runtime.executeCopilot(input)
);
return JSON.parse(output) as {
events: Array<{
type: 'route_selected' | 'route_failed' | 'usage';
route: {
profileId: string;
source: 'server' | 'local' | 'affine_cloud';
provider: string;
model: string;
};
errorKind?: string;
usage?: unknown;
}>;
result: unknown;
};
}
async assertCopilotRoute(input: CopilotRouteCheckInput) {
await this.measured('assertCopilotRoute', runtime =>
runtime.assertCopilotRoute(input)
);
}
streamCopilot<TEvent>(
input: CopilotExecuteInput,
toolCallback: (request: string) => Promise<string>,
options: { maxSteps: number; signal?: AbortSignal }
): AsyncIterableIterator<TEvent> {
const stream = new RuntimeEventStream<TEvent>();
const endMarker = '__AFFINE_COPILOT_STREAM_END__';
void this.runtime
.executeCopilotStream(
input,
options.maxSteps,
(error, value) => {
if (error) {
stream.push({
type: 'error',
errorKind: 'callback',
message: error.message,
} as TEvent);
} else if (value === endMarker) {
stream.push();
} else {
stream.push(JSON.parse(value) as TEvent);
}
},
async (error, request) => {
if (error) throw error;
return await toolCallback(request);
}
)
.then(handle => {
stream.attach(() => handle.abort());
if (options.signal?.aborted) handle.abort();
else
options.signal?.addEventListener('abort', () => handle.abort(), {
once: true,
});
})
.catch(error => {
stream.push({
type: 'error',
errorKind: 'setup',
message: error instanceof Error ? error.message : String(error),
} as TEvent);
stream.push();
});
return stream;
}
async isInviteAbuseUserQuarantinedOrBanned(userId: string) {
return await this.measured('isInviteAbuseUserQuarantinedOrBanned', rt =>
this.quotaRuntime(rt).isInviteAbuseUserQuarantinedOrBanned(userId)
packages/backend/server/src/core/realtime/__tests__/registry.spec.ts
+19 -10
View File
@@ -8,9 +8,9 @@ import { z } from 'zod';
import { CANARY_CLIENT_VERSION_MAX_AGE_DAYS } from '../../../base';
import { Flavor } from '../../../env';
import { PublicDocMode } from '../../../models';
import { CopilotEmbeddingRealtimeProvider } from '../../../plugins/copilot/context';
import type { CopilotTranscriptionReader } from '../../../plugins/copilot/transcript';
import { CopilotTranscriptRealtimeProvider } from '../../../plugins/copilot/transcript';
import { CopilotEmbeddingRealtimeProvider } from '../../../plugins/copilot/context/realtime';
import type { CopilotTranscriptionReader } from '../../../plugins/copilot/transcript/reader';
import { CopilotTranscriptRealtimeProvider } from '../../../plugins/copilot/transcript/realtime';
import type { CurrentUser } from '../../auth';
import { CommentRealtimeProvider } from '../../comment/realtime';
import { NotificationRealtimeProvider } from '../../notification/realtime';
@@ -440,12 +440,14 @@ test('front and sync realtime gateway required handlers are registered by lightw
{} as never,
{} as never,
registry,
{} as never,
{} as never
).onModuleInit();
new CopilotTranscriptRealtimeProvider(
{} as never,
{} as never,
registry
registry,
{} as never
).onModuleInit();
new QuotaStateRealtimeProvider(
{} as never,
@@ -1000,12 +1002,14 @@ test('copilot embedding realtime provider uses lightweight model reads', async t
const publisher = {
publish: (...args: unknown[]) => published.push(args),
} as unknown as RealtimePublisher;
const config = { copilot: { enabled: true } };
const provider = new CopilotEmbeddingRealtimeProvider(
ac,
models as never,
registry,
publisher
publisher,
config as never
);
provider.onModuleInit();
@@ -1018,6 +1022,13 @@ test('copilot embedding realtime provider uses lightweight model reads', async t
embedded: 3,
}
);
config.copilot.enabled = false;
await t.throwsAsync(
registry
.getRequest('workspace.embedding.progress.get')
.handle(user, { workspaceId: 'space' }),
{ message: 'Copilot is disabled.' }
);
t.is(
registry
.getTopic('workspace.embedding.progress.changed')
@@ -1068,11 +1079,9 @@ test('copilot transcript realtime provider registers task live query handlers',
},
} as unknown as CopilotTranscriptionReader;
new CopilotTranscriptRealtimeProvider(
ac,
transcript,
registry
).onModuleInit();
new CopilotTranscriptRealtimeProvider(ac, transcript, registry, {
copilot: { enabled: true },
} as never).onModuleInit();
t.deepEqual(
await registry.getRequest('copilot.transcript.task.get').handle(user, {
packages/backend/server/src/models/copilot-byok.ts
+1 -117
View File
@@ -1,135 +1,19 @@
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({
await this.db.aiWorkspaceByokConfig.updateMany({
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 },
packages/backend/server/src/models/copilot-session.ts
+13 -60
View File
@@ -1,6 +1,6 @@
import { Injectable } from '@nestjs/common';
import { Transactional } from '@nestjs-cls/transactional';
import { AiPromptRole, Prisma } from '@prisma/client';
import { AiSessionMessageRole, Prisma } from '@prisma/client';
import { omit } from 'lodash-es';
import {
@@ -9,7 +9,6 @@ import {
CopilotSessionInvalidInput,
CopilotSessionNotFound,
} from '../base';
import { getTokenEncoder } from '../native';
import type { PromptAttachment } from '../plugins/copilot/providers/types';
import {
type ChatMessage as CopilotChatMessage,
@@ -26,7 +25,6 @@ export enum SessionType {
type ChatPrompt = {
name: string;
action?: string | null;
model: string;
};
type ChatAttachment = PromptAttachment;
@@ -95,12 +93,11 @@ export type ForkSessionOptions = Omit<
ChatSession,
'messages' | 'promptName' | 'promptAction'
> & {
prompt: { name: string; action: string | null | undefined; model: string };
prompt: { name: string; action: string | null | undefined };
messages: ChatMessage[];
};
type UpdateChatSessionMessage = ChatSessionBaseState & {
prompt: { model: string };
messages: ChatMessage[];
};
@@ -108,7 +105,7 @@ export type UpdateChatSessionOptions = ChatSessionBaseState &
Pick<
Partial<ChatSession>,
'docId' | 'pinned' | 'promptName' | 'promptAction' | 'title'
> & { promptModel?: string };
>;
export type UpdateChatSession = ChatSessionBaseState & UpdateChatSessionOptions;
@@ -144,20 +141,6 @@ export class CopilotSessionModel extends BaseModel {
};
}
private async ensurePromptCompatRecord(prompt: ChatPrompt) {
await this.db.aiPrompt.upsert({
where: { name: prompt.name },
update: {},
create: {
name: prompt.name,
action: prompt.action,
model: prompt.model,
optionalModels: [],
config: {},
},
});
}
private sanitizeString<T extends string | null | undefined>(value: T): T {
if (typeof value !== 'string') {
return value;
@@ -354,7 +337,7 @@ export class CopilotSessionModel extends BaseModel {
private isCountedUserMessage(
message: Pick<StoredChatMessage, 'role'>
): boolean {
return message.role === AiPromptRole.user;
return message.role === AiSessionMessageRole.user;
}
getSessionType(session: Pick<ChatSession, 'docId' | 'pinned'>): SessionType {
@@ -418,7 +401,6 @@ export class CopilotSessionModel extends BaseModel {
reuseChat = false
): Promise<string> {
const { prompt, ...rest } = state;
await this.ensurePromptCompatRecord(prompt);
return await this.models.copilotSession.create(
{ ...rest, promptName: prompt.name, promptAction: prompt.action ?? null },
reuseChat
@@ -507,7 +489,6 @@ export class CopilotSessionModel extends BaseModel {
pinned: true,
title: true,
promptName: true,
tokenCost: true,
createdAt: true,
updatedAt: true,
messages: {
@@ -536,7 +517,6 @@ export class CopilotSessionModel extends BaseModel {
pinned: true,
title: true,
promptName: true,
tokenCost: true,
createdAt: true,
updatedAt: true,
});
@@ -605,7 +585,6 @@ export class CopilotSessionModel extends BaseModel {
pinned: true,
title: true,
promptName: true,
tokenCost: true,
createdAt: true,
updatedAt: true,
messages: options.withMessages
@@ -682,25 +661,11 @@ export class CopilotSessionModel extends BaseModel {
let nextPromptAction: string | null | undefined;
if (promptName) {
if (options.promptModel) {
await this.ensurePromptCompatRecord({
name: promptName,
action: options.promptAction,
model: options.promptModel,
});
}
nextPromptAction = options.promptAction;
if (nextPromptAction === undefined) {
const prompt = await this.db.aiPrompt.findFirst({
where: { name: promptName },
select: { action: true },
});
if (!prompt) {
throw new CopilotSessionInvalidInput(
`Prompt ${promptName} not found or not available for session ${sessionId}`
);
}
nextPromptAction = prompt.action ?? null;
throw new CopilotSessionInvalidInput(
`Prompt action is required when changing prompt ${promptName}`
);
}
if (nextPromptAction) {
throw new CopilotSessionInvalidInput(
@@ -809,12 +774,6 @@ export class CopilotSessionModel extends BaseModel {
return message ? this.toPublicMessage(message) : null;
}
private calculateTokenSize(messages: any[], model: string): number {
const encoder = getTokenEncoder(model);
const content = messages.map(m => m.content).join('');
return encoder?.count(content) || 0;
}
@Transactional()
async updateMessages(state: UpdateChatSessionMessage) {
const { sessionId, userId, messages } = state;
@@ -825,10 +784,6 @@ export class CopilotSessionModel extends BaseModel {
if (messages.length) {
const sanitizedMessages = messages.map(m => this.sanitizeMessage(m));
const tokenCost = this.calculateTokenSize(
sanitizedMessages,
state.prompt.model
);
await this.db.aiSessionMessage.createMany({
data: sanitizedMessages.map(m => ({
compatSubmissionId: m.compatSubmissionId || undefined,
@@ -848,7 +803,6 @@ export class CopilotSessionModel extends BaseModel {
where: { id: sessionId },
data: {
messageCost: { increment: userMessages.length },
tokenCost: { increment: tokenCost },
},
});
}
@@ -858,7 +812,6 @@ export class CopilotSessionModel extends BaseModel {
async appendMessage(state: {
sessionId: string;
userId: string;
prompt: { model: string };
message: ChatMessage;
}) {
const haveSession = await this.has(state.sessionId, state.userId);
@@ -867,8 +820,6 @@ export class CopilotSessionModel extends BaseModel {
}
const message = this.sanitizeMessage(state.message);
const tokenCost = this.calculateTokenSize([message], state.prompt.model);
const created = await this.db.aiSessionMessage.create({
data: {
sessionId: state.sessionId,
@@ -896,8 +847,9 @@ export class CopilotSessionModel extends BaseModel {
where: { id: state.sessionId },
data: {
messageCost:
message.role === AiPromptRole.user ? { increment: 1 } : undefined,
tokenCost: { increment: tokenCost },
message.role === AiSessionMessageRole.user
? { increment: 1 }
: undefined,
},
});
@@ -970,8 +922,9 @@ export class CopilotSessionModel extends BaseModel {
});
const ids = messages
.slice(
messages.findLastIndex(({ role }) => role === AiPromptRole.user) +
(removeLatestUserMessage ? 0 : 1)
messages.findLastIndex(
({ role }) => role === AiSessionMessageRole.user
) + (removeLatestUserMessage ? 0 : 1)
)
.map(({ id }) => id);
packages/backend/server/src/models/copilot-transcript-task.ts
+1 -7
View File
@@ -24,12 +24,7 @@ export class CopilotTranscriptTaskModel extends BaseModel {
async create(
input: Pick<
Prisma.AiTranscriptTaskCreateArgs['data'],
| 'userId'
| 'workspaceId'
| 'blobId'
| 'strategy'
| 'recipeId'
| 'recipeVersion'
'userId' | 'workspaceId' | 'blobId' | 'recipeId' | 'recipeVersion'
> &
Partial<Prisma.AiTranscriptTaskCreateArgs['data']>
) {
@@ -39,7 +34,6 @@ export class CopilotTranscriptTaskModel extends BaseModel {
workspaceId: input.workspaceId,
blobId: input.blobId,
status: 'pending',
strategy: input.strategy,
recipeId: input.recipeId,
recipeVersion: input.recipeVersion,
inputSnapshot: nullableJson(input.inputSnapshot),
packages/backend/server/src/native.ts
+67 -1037
View File
File diff suppressed because it is too large Load Diff
packages/backend/server/src/plugins/copilot/access/index.ts
-1
View File
@@ -1,2 +1 @@
export * from './feature-coverage';
export * from './policy';
packages/backend/server/src/plugins/copilot/access/policy.ts
-106
View File
@@ -1,106 +0,0 @@
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);
}
}
}
packages/backend/server/src/plugins/copilot/availability.ts
+8
View File
@@ -0,0 +1,8 @@
import type { Config } from '../../base/config';
import { ActionForbidden } from '../../base/error/errors.gen';
export function assertCopilotEnabled(config: Config) {
if (!config.copilot.enabled) {
throw new ActionForbidden('Copilot is disabled.');
}
}
packages/backend/server/src/plugins/copilot/byok/index.ts
-1
View File
@@ -1,4 +1,3 @@
export { ByokEntitlementPolicy } from './policy';
export { WorkspaceByokResolver } from './resolver';
export { type ByokProviderRequestContext, ByokService } from './service';
export * from './types';
packages/backend/server/src/plugins/copilot/byok/policy.ts
+10
View File
@@ -103,6 +103,16 @@ export class ByokEntitlementPolicy {
}
}
async assertEntitled(workspaceId: string, userId?: string) {
const [serverEntitled, localEntitled] = await this.hasEntitlement(
workspaceId,
userId
);
if (!serverEntitled && !localEntitled) {
throw new ActionForbidden('BYOK requires Pro, Team, or Believer.');
}
}
private async hasWorkspaceTeamPlan(workspaceId: string) {
try {
const state =
packages/backend/server/src/plugins/copilot/byok/probe.ts
-94
View File
@@ -1,94 +0,0 @@
import { BadRequestException } from '@nestjs/common';
import type { safeFetch } from '../../../base';
import { ByokProvider } from './types';
const TEST_TIMEOUT_MS = 10_000;
export const PROVIDER_PROBE_MAX_BYTES = 1024 * 1024;
type ProbeFetch = typeof safeFetch;
export async function runProviderProbe(
probeFetch: ProbeFetch,
provider: ByokProvider,
apiKey: string,
endpoint: string | null,
allowPrivateEndpoint: boolean
) {
const request = buildProbeRequest(provider, apiKey, endpoint);
const response = await probeFetch(
request.url,
{
method: request.method,
headers: request.headers,
},
{
timeoutMs: TEST_TIMEOUT_MS,
maxRedirects: 3,
maxBytes: PROVIDER_PROBE_MAX_BYTES,
allowedHeaders: Object.keys(request.headers),
allowHttp: endpoint?.startsWith('http:') ?? false,
allowPrivateTargetOrigin: allowPrivateEndpoint,
}
);
if (!response.ok) {
throw new BadRequestException(providerProbeFailureMessage(response.status));
}
}
function buildProbeRequest(
provider: ByokProvider,
apiKey: string,
endpoint: string | null
): {
method: 'GET';
url: string;
headers: Record<string, string>;
} {
switch (provider) {
case ByokProvider.openai:
return {
method: 'GET',
url: `${endpoint ?? 'https://api.openai.com/v1'}/models`,
headers: { Authorization: `Bearer ${apiKey}` },
};
case ByokProvider.anthropic:
return {
method: 'GET',
url: `${endpoint ?? 'https://api.anthropic.com/v1'}/models`,
headers: {
'x-api-key': apiKey,
'anthropic-version': '2023-06-01',
},
};
case ByokProvider.gemini:
return {
method: 'GET',
url: `${endpoint ?? 'https://generativelanguage.googleapis.com/v1beta'}/models`,
headers: { 'x-goog-api-key': apiKey },
};
case ByokProvider.fal:
return {
method: 'GET',
url: 'https://api.fal.ai/v1/models?limit=10',
headers: { Authorization: `Key ${apiKey}` },
};
}
}
function providerProbeFailureMessage(status: number) {
switch (status) {
case 401:
return 'Provider rejected the BYOK key.';
case 403:
return 'Provider rejected the BYOK key permissions.';
case 404:
return 'Provider probe endpoint was not found.';
case 429:
return 'Provider rate limit exceeded while testing the key.';
default:
return status >= 500
? 'Provider service is unavailable.'
: `Provider key test failed with HTTP ${status}.`;
}
}
packages/backend/server/src/plugins/copilot/byok/resolver.ts
+564 -189
View File
@@ -1,3 +1,4 @@
import { BadRequestException } from '@nestjs/common';
import {
Args,
Field,
@@ -11,18 +12,139 @@ import {
} from '@nestjs/graphql';
import { SafeIntResolver } from 'graphql-scalars';
import { Throttle } from '../../../base';
import { Config, Throttle } from '../../../base';
import { CurrentUser } from '../../../core/auth';
import { BackendRuntimeProvider } from '../../../core/backend-runtime';
import { PermissionAccess } from '../../../core/permission';
import { WorkspaceType } from '../../../core/workspaces';
import { Models } from '../../../models';
import { llmGetByokCatalog } from '../../../native';
import { CopilotEnabled } from '../feature';
import { ByokEntitlementPolicy } from './policy';
import { ByokKeyConfig, ByokLocalLeaseProvider, ByokService } from './service';
import { ByokKeyStorage, ByokKeyTestStatus, ByokProvider } from './types';
import {
BYOK_ALLOWED_PROVIDERS,
ByokProvider,
ByokProviderSource,
} from './types';
@ObjectType()
export class WorkspaceByokKeyConfigType implements ByokKeyConfig {
class WorkspaceByokCapabilityType {
@Field(() => [String])
input!: string[];
@Field(() => [String])
output!: string[];
@Field(() => [String])
features!: string[];
@Field(() => [String])
attachmentKinds!: string[];
@Field(() => [String])
attachmentSources!: string[];
}
@ObjectType()
class WorkspaceByokModelDeclarationType {
@Field(() => String)
modelId!: string;
@Field(() => Boolean)
enabled!: boolean;
@Field(() => [WorkspaceByokCapabilityType])
capabilities!: WorkspaceByokCapabilityType[];
}
@ObjectType()
class WorkspaceByokEndpointType {
@Field(() => String)
kind!: string;
@Field(() => String, { nullable: true })
url!: string | null;
}
@ObjectType()
class WorkspaceByokProfileDefinitionType {
@Field(() => SafeIntResolver)
version!: number;
@Field(() => WorkspaceByokEndpointType)
endpoint!: WorkspaceByokEndpointType;
@Field(() => [WorkspaceByokModelDeclarationType])
models!: WorkspaceByokModelDeclarationType[];
}
@ObjectType()
class WorkspaceByokProbeStatusType {
@Field(() => String)
kind!: string;
@Field(() => Date, { nullable: true })
testedAt!: Date | null;
@Field(() => String, { nullable: true })
errorKind!: string | null;
}
@ObjectType()
class WorkspaceByokModelProbeCheckType {
@Field(() => String)
operation!: string;
@Field(() => WorkspaceByokProbeStatusType)
status!: WorkspaceByokProbeStatusType;
}
@ObjectType()
class WorkspaceByokModelProbeType {
@Field(() => String)
modelId!: string;
@Field(() => [WorkspaceByokModelProbeCheckType])
checks!: WorkspaceByokModelProbeCheckType[];
}
@ObjectType()
class WorkspaceByokValidationType {
@Field(() => String)
definitionFingerprint!: string;
@Field(() => SafeIntResolver)
credentialGeneration!: number;
@Field(() => WorkspaceByokProbeStatusType)
connection!: WorkspaceByokProbeStatusType;
@Field(() => [WorkspaceByokModelProbeType])
models!: WorkspaceByokModelProbeType[];
}
@ObjectType()
class WorkspaceByokProbeResultType {
@Field(() => String)
definitionFingerprint!: string;
@Field(() => Boolean)
stale!: boolean;
@Field(() => WorkspaceByokProbeStatusType)
connection!: WorkspaceByokProbeStatusType;
@Field(() => [WorkspaceByokModelProbeType])
models!: WorkspaceByokModelProbeType[];
}
@ObjectType()
export class WorkspaceByokProfileType {
@Field(() => ID)
id!: string;
profileId!: string;
@Field(() => String)
workspaceId!: string;
@Field(() => ByokProvider)
provider!: ByokProvider;
@@ -33,59 +155,53 @@ export class WorkspaceByokKeyConfigType implements ByokKeyConfig {
@Field(() => String, { nullable: true })
description!: string | null;
@Field(() => ByokKeyStorage)
storage!: ByokKeyStorage;
@Field(() => Boolean)
configured!: boolean;
@Field(() => WorkspaceByokProfileDefinitionType)
definition!: WorkspaceByokProfileDefinitionType;
@Field(() => Boolean)
enabled!: boolean;
@Field(() => String, { nullable: true })
endpoint!: string | null;
@Field(() => Boolean)
endpointEditable!: boolean;
@Field(() => SafeIntResolver)
sortOrder!: number;
@Field(() => [String])
capabilities!: string[];
@Field(() => SafeIntResolver)
revision!: number;
@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;
@Field(() => WorkspaceByokValidationType, { nullable: true })
validation!: WorkspaceByokValidationType | null;
}
@ObjectType()
class WorkspaceByokCapabilityWarningType {
class WorkspaceByokCatalogModelType {
@Field(() => String)
featureKind!: string;
modelId!: string;
@Field(() => String)
reason!: string;
displayName!: string;
@Field(() => [ByokProvider])
requiredProviders!: ByokProvider[];
@Field(() => Boolean)
recommended!: boolean;
@Field(() => [WorkspaceByokCapabilityType])
capabilities!: WorkspaceByokCapabilityType[];
}
@ObjectType()
class WorkspaceByokCatalogProviderType {
@Field(() => ByokProvider)
provider!: ByokProvider;
@Field(() => [WorkspaceByokCatalogModelType])
models!: WorkspaceByokCatalogModelType[];
}
@ObjectType()
class WorkspaceByokCatalogType {
@Field(() => String)
version!: string;
@Field(() => [WorkspaceByokCatalogProviderType])
providers!: WorkspaceByokCatalogProviderType[];
}
@ObjectType()
@@ -102,29 +218,20 @@ class WorkspaceByokSettingsType {
@Field(() => Boolean)
localEntitled!: boolean;
@Field(() => [String])
entitlementRequired!: string[];
@Field(() => [WorkspaceByokKeyConfigType])
keys!: WorkspaceByokKeyConfigType[];
@Field(() => [WorkspaceByokProfileType])
profiles!: WorkspaceByokProfileType[];
@Field(() => [ByokProvider])
allowedProviders!: ByokProvider[];
@Field(() => Boolean)
localStorageSupported!: boolean;
@Field(() => Boolean)
customEndpointSupported!: boolean;
@Field(() => Boolean)
privateEndpointSupported!: boolean;
@Field(() => Boolean)
hasAiPlan!: boolean;
@Field(() => [WorkspaceByokCapabilityWarningType])
warnings!: WorkspaceByokCapabilityWarningType[];
@Field(() => WorkspaceByokCatalogType)
catalog!: WorkspaceByokCatalogType;
}
@ObjectType()
@@ -139,18 +246,6 @@ class WorkspaceByokUsagePointType {
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)
@@ -161,10 +256,58 @@ class CreateWorkspaceByokLocalLeaseResultType {
}
@InputType()
class UpsertWorkspaceByokConfigInput {
@Field(() => ID, { nullable: true })
id?: string;
class WorkspaceByokCapabilityInput {
@Field(() => [String])
input!: string[];
@Field(() => [String])
output!: string[];
@Field(() => [String])
features!: string[];
@Field(() => [String])
attachmentKinds!: string[];
@Field(() => [String])
attachmentSources!: string[];
}
@InputType()
class WorkspaceByokModelDeclarationInput {
@Field(() => String)
modelId!: string;
@Field(() => Boolean)
enabled!: boolean;
@Field(() => [WorkspaceByokCapabilityInput])
capabilities!: WorkspaceByokCapabilityInput[];
}
@InputType()
class WorkspaceByokEndpointInput {
@Field(() => String)
kind!: string;
@Field(() => String, { nullable: true })
url!: string | null;
}
@InputType()
class WorkspaceByokProfileDefinitionInput {
@Field(() => SafeIntResolver)
version!: number;
@Field(() => WorkspaceByokEndpointInput)
endpoint!: WorkspaceByokEndpointInput;
@Field(() => [WorkspaceByokModelDeclarationInput])
models!: WorkspaceByokModelDeclarationInput[];
}
@InputType()
class CreateWorkspaceByokProfileInput {
@Field(() => String)
workspaceId!: string;
@@ -175,59 +318,125 @@ class UpsertWorkspaceByokConfigInput {
name!: string;
@Field(() => String, { nullable: true })
description?: string | null;
description!: string | null;
@Field(() => ByokKeyStorage)
storage!: ByokKeyStorage;
@Field(() => String)
credential!: string;
@Field(() => String, { nullable: true })
apiKey?: string | null;
@Field(() => WorkspaceByokProfileDefinitionInput)
definition!: WorkspaceByokProfileDefinitionInput;
@Field(() => String, { nullable: true })
endpoint?: string | null;
@Field(() => SafeIntResolver, { nullable: true })
sortOrder?: number | null;
@Field(() => Boolean, { nullable: true })
enabled?: boolean | null;
@Field(() => Boolean)
enabled!: boolean;
}
@InputType()
class TestWorkspaceByokConfigInput {
class ReplaceWorkspaceByokProfileInput {
@Field(() => String)
workspaceId!: string;
@Field(() => ID)
profileId!: string;
@Field(() => SafeIntResolver)
expectedRevision!: number;
@Field(() => String)
name!: string;
@Field(() => String, { nullable: true })
description!: string | null;
@Field(() => WorkspaceByokProfileDefinitionInput)
definition!: WorkspaceByokProfileDefinitionInput;
@Field(() => String, { nullable: true })
credential!: string | null;
@Field(() => Boolean)
enabled!: boolean;
}
@InputType()
class RotateWorkspaceByokCredentialInput {
@Field(() => String)
workspaceId!: string;
@Field(() => ID)
profileId!: string;
@Field(() => SafeIntResolver)
expectedRevision!: number;
@Field(() => String)
credential!: string;
}
@InputType()
class WorkspaceByokProbeCheckInput {
@Field(() => String)
modelId!: string;
@Field(() => String)
operation!: string;
}
@InputType()
class ProbeWorkspaceByokProfileInput {
@Field(() => String)
workspaceId!: string;
@Field(() => ID)
profileId!: string;
@Field(() => [WorkspaceByokProbeCheckInput])
checks!: WorkspaceByokProbeCheckInput[];
}
@InputType()
class ProbeWorkspaceByokDraftInput {
@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;
credential!: string | null;
@Field(() => ID, { nullable: true })
configId?: string | null;
profileId!: string | null;
@Field(() => SafeIntResolver, { nullable: true })
expectedRevision!: number | null;
@Field(() => WorkspaceByokProfileDefinitionInput)
definition!: WorkspaceByokProfileDefinitionInput;
@Field(() => [WorkspaceByokProbeCheckInput])
checks!: WorkspaceByokProbeCheckInput[];
}
@InputType()
class ReorderWorkspaceByokConfigsInput {
class WorkspaceByokProfileOrderInput {
@Field(() => ID)
profileId!: string;
@Field(() => SafeIntResolver)
expectedRevision!: number;
}
@InputType()
class ReorderWorkspaceByokProfilesInput {
@Field(() => String)
workspaceId!: string;
@Field(() => ByokKeyStorage)
storage!: ByokKeyStorage;
@Field(() => [ID])
ids!: string[];
@Field(() => [WorkspaceByokProfileOrderInput])
profiles!: WorkspaceByokProfileOrderInput[];
}
@InputType()
class CreateWorkspaceByokLocalLeaseProviderInput implements ByokLocalLeaseProvider {
class CreateWorkspaceByokLocalLeaseProviderInput {
@Field(() => ByokProvider)
provider!: ByokProvider;
@@ -235,19 +444,16 @@ class CreateWorkspaceByokLocalLeaseProviderInput implements ByokLocalLeaseProvid
name!: string;
@Field(() => String, { nullable: true })
description?: string | null;
description!: string | null;
@Field(() => String)
apiKey!: string;
credential!: string;
@Field(() => String, { nullable: true })
endpoint?: string | null;
@Field(() => WorkspaceByokProfileDefinitionInput)
definition!: WorkspaceByokProfileDefinitionInput;
@Field(() => SafeIntResolver, { nullable: true })
sortOrder?: number | null;
@Field(() => Boolean, { nullable: true })
enabled?: boolean | null;
@Field(() => Boolean)
enabled!: boolean;
}
@InputType()
@@ -259,12 +465,15 @@ class CreateWorkspaceByokLocalLeaseInput {
providers!: CreateWorkspaceByokLocalLeaseProviderInput[];
}
@CopilotEnabled()
@Resolver(() => WorkspaceType)
export class WorkspaceByokResolver {
constructor(
private readonly ac: PermissionAccess,
private readonly entitlement: ByokEntitlementPolicy,
private readonly byok: ByokService
private readonly runtime: BackendRuntimeProvider,
private readonly models: Models,
private readonly config: Config
) {}
@ResolveField(() => WorkspaceByokSettingsType, {
@@ -275,13 +484,35 @@ export class WorkspaceByokResolver {
@CurrentUser() user: CurrentUser,
@Parent() workspace: WorkspaceType
) {
await this.ac
.user(user.id)
.workspace(workspace.id)
.allowLocal()
.assert('Workspace.Settings.Read');
await this.assertRead(user.id, workspace.id);
await this.entitlement.assertManagementAccess(workspace.id, user.id);
return await this.byok.getSettings(workspace.id, user.id);
const [serverEntitled, localEntitled] =
await this.entitlement.hasEntitlement(workspace.id, user.id);
const profiles = serverEntitled
? await this.runtime.listByokProfiles(workspace.id)
: [];
const customEndpointSupported =
this.config.copilot.byok.allowCustomEndpoint;
const catalog = llmGetByokCatalog();
return {
workspaceId: workspace.id,
entitled: serverEntitled || localEntitled,
serverEntitled,
localEntitled,
profiles: profiles.map(profile => projectProfile(profile)),
allowedProviders: [...BYOK_ALLOWED_PROVIDERS],
customEndpointSupported,
privateEndpointSupported:
customEndpointSupported &&
this.config.copilot.byok.allowPrivateEndpoint,
catalog: {
...catalog,
providers: catalog.providers.map(provider => ({
...provider,
provider: provider.provider as ByokProvider,
})),
},
};
}
@ResolveField(() => [WorkspaceByokUsagePointType], {
@@ -294,100 +525,131 @@ export class WorkspaceByokResolver {
@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.assertRead(user.id, workspace.id);
await this.entitlement.assertManagementAccess(workspace.id, user.id);
return await this.byok.getUsage(workspace.id, from, to);
return await this.models.copilotUsage.aggregateByDay({
workspaceId: workspace.id,
from,
to,
providerSources: [ByokProviderSource.Server, ByokProviderSource.Local],
});
}
@Mutation(() => WorkspaceByokProfileType)
@Throttle('strict')
@Mutation(() => TestWorkspaceByokConfigResultType)
async testWorkspaceByokConfig(
async createWorkspaceByokProfile(
@CurrentUser() user: CurrentUser,
@Args('input') input: TestWorkspaceByokConfigInput
@Args('input') input: CreateWorkspaceByokProfileInput
) {
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.assertUpdate(user.id, input.workspaceId);
await this.entitlement.assertServerEntitled(input.workspaceId);
requireExplicitDescription(input);
return projectProfile(
await this.runtime.createByokProfile({
...input,
description: input.description ?? undefined,
definition: nativeDefinition(input.definition),
actorUserId: user.id,
})
);
}
@Mutation(() => WorkspaceByokProfileType)
@Throttle('strict')
async replaceWorkspaceByokProfile(
@CurrentUser() user: CurrentUser,
@Args('input') input: ReplaceWorkspaceByokProfileInput
) {
await this.assertUpdate(user.id, input.workspaceId);
await this.entitlement.assertServerEntitled(input.workspaceId);
requireExplicitDescription(input);
return projectProfile(
await this.runtime.replaceByokProfile({
...input,
description: input.description ?? undefined,
credential: input.credential ?? undefined,
definition: nativeDefinition(input.definition),
actorUserId: user.id,
})
);
}
@Mutation(() => WorkspaceByokProfileType)
@Throttle('strict')
async rotateWorkspaceByokCredential(
@CurrentUser() user: CurrentUser,
@Args('input') input: RotateWorkspaceByokCredentialInput
) {
await this.assertUpdate(user.id, input.workspaceId);
await this.entitlement.assertServerEntitled(input.workspaceId);
return projectProfile(
await this.runtime.rotateByokCredential({
...input,
actorUserId: user.id,
})
);
}
@Mutation(() => WorkspaceByokProbeResultType)
@Throttle('strict')
async probeWorkspaceByokProfile(
@CurrentUser() user: CurrentUser,
@Args('input') input: ProbeWorkspaceByokProfileInput
) {
await this.assertUpdate(user.id, input.workspaceId);
await this.entitlement.assertServerEntitled(input.workspaceId);
return projectProbeResult(await this.runtime.probeByokProfile(input));
}
@Mutation(() => WorkspaceByokProbeResultType)
@Throttle('strict')
async probeWorkspaceByokDraft(
@CurrentUser() user: CurrentUser,
@Args('input') input: ProbeWorkspaceByokDraftInput
) {
await this.assertUpdate(user.id, input.workspaceId);
if (input.profileId) {
await this.entitlement.assertServerEntitled(input.workspaceId);
} else {
await this.entitlement.assertLocalEntitled(input.workspaceId, user.id);
await this.entitlement.assertEntitled(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 });
return projectProbeResult(
await this.runtime.probeByokDraft({
...input,
credential: input.credential ?? undefined,
profileId: input.profileId ?? undefined,
expectedRevision: input.expectedRevision ?? undefined,
definition: nativeDefinition(input.definition),
})
);
}
@Mutation(() => Boolean)
@Throttle('strict')
async deleteWorkspaceByokConfig(
async deleteWorkspaceByokProfile(
@CurrentUser() user: CurrentUser,
@Args('id', { type: () => ID }) id: string,
@Args('profileId', { type: () => ID }) profileId: 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.assertUpdate(user.id, workspaceId);
await this.entitlement.assertServerEntitled(workspaceId);
return await this.byok.deleteConfig(workspaceId, id, user.id);
return await this.runtime.deleteByokProfile(workspaceId, profileId);
}
@Mutation(() => Boolean)
@Mutation(() => [WorkspaceByokProfileType])
@Throttle('strict')
async clearWorkspaceByokConfigs(
async reorderWorkspaceByokProfiles(
@CurrentUser() user: CurrentUser,
@Args('workspaceId', { type: () => String }) workspaceId: string,
@Args('provider', { type: () => ByokProvider, nullable: true })
provider?: ByokProvider | null
@Args('input') input: ReorderWorkspaceByokProfilesInput
) {
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);
await this.assertUpdate(user.id, input.workspaceId);
await this.entitlement.assertServerEntitled(input.workspaceId);
return (
await this.runtime.reorderByokProfiles({
...input,
actorUserId: user.id,
})
).map(profile => projectProfile(profile));
}
@Mutation(() => CreateWorkspaceByokLocalLeaseResultType)
@@ -403,6 +665,119 @@ export class WorkspaceByokResolver {
.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 });
input.providers.forEach(requireExplicitDescription);
const result = await this.runtime.createByokLocalLease({
...input,
providers: input.providers.map(provider => ({
...provider,
description: provider.description ?? undefined,
definition: nativeDefinition(provider.definition),
})),
userId: user.id,
});
return {
leaseId: result.leaseId,
expiresAt: new Date(result.expiresAtMs),
};
}
private async assertRead(userId: string, workspaceId: string) {
await this.ac
.user(userId)
.workspace(workspaceId)
.allowLocal()
.assert('Workspace.Settings.Read');
}
private async assertUpdate(userId: string, workspaceId: string) {
await this.ac
.user(userId)
.workspace(workspaceId)
.allowLocal()
.assert('Workspace.Settings.Update');
await this.entitlement.assertManagementAccess(workspaceId, userId);
}
}
function requireExplicitDescription(input: { description: string | null }) {
if (!Object.hasOwn(input, 'description')) {
throw new BadRequestException('description must be provided explicitly.');
}
}
function nativeDefinition(input: WorkspaceByokProfileDefinitionInput) {
return {
...input,
endpoint: {
...input.endpoint,
url: input.endpoint.url ?? undefined,
},
};
}
function projectProbe(probe: {
kind: string;
testedAtMs?: number;
errorKind?: string;
}) {
return {
kind: probe.kind,
testedAt: probe.testedAtMs ? new Date(probe.testedAtMs) : null,
errorKind: probe.errorKind ?? null,
};
}
function projectProbeResult(result: {
definitionFingerprint: string;
stale: boolean;
connection: {
kind: string;
testedAtMs?: number;
errorKind?: string;
};
models: Array<{
modelId: string;
checks: Array<{
operation: string;
status: {
kind: string;
testedAtMs?: number;
errorKind?: string;
};
}>;
}>;
}) {
return {
...result,
connection: projectProbe(result.connection),
models: result.models.map(model => ({
...model,
checks: model.checks.map(check => ({
...check,
status: projectProbe(check.status),
})),
})),
};
}
function projectProfile(
profile: Awaited<ReturnType<BackendRuntimeProvider['createByokProfile']>>
) {
return {
...profile,
provider: profile.provider as ByokProvider,
validation: profile.validation
? {
...profile.validation,
connection: projectProbe(profile.validation.connection),
models: profile.validation.models.map(model => ({
...model,
checks: model.checks.map(check => ({
...check,
status: projectProbe(check.status),
})),
})),
}
: null,
};
}
packages/backend/server/src/plugins/copilot/byok/service.ts
-833
View File
@@ -1,833 +0,0 @@
import { createHash, createHmac, randomUUID } from 'node:crypto';
import { BadRequestException, Injectable } from '@nestjs/common';
import {
BadRequest,
Cache,
Config,
CryptoHelper,
metrics,
safeFetch,
} from '../../../base';
import { Models } from '../../../models';
import type { CopilotProviderProfile } from '../config';
import { ByokEntitlementPolicy } from './policy';
import { runProviderProbe } from './probe';
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;
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;
privateEndpointSupported: 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 {
private readonly probeFetch = safeFetch;
constructor(
private readonly models: Models,
private readonly crypto: CryptoHelper,
private readonly cache: Cache,
private readonly entitlement: ByokEntitlementPolicy,
private readonly config: Config
) {}
get customEndpointSupported() {
return env.selfhosted && this.config.copilot.byok.allowCustomEndpoint;
}
get privateEndpointSupported() {
return (
this.customEndpointSupported &&
this.config.copilot.byok.allowPrivateEndpoint
);
}
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,
privateEndpointSupported: this.privateEndpointSupported,
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,
privateEndpointSupported: this.privateEndpointSupported,
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,
privateEndpointSupported: this.privateEndpointSupported,
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 runProviderProbe(
this.probeFetch,
input.provider,
apiKey,
endpoint,
this.privateEndpointSupported
);
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 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 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}`;
}
}
packages/backend/server/src/plugins/copilot/compat/history-projector.ts
+3 -7
View File
@@ -1,5 +1,5 @@
import { Injectable } from '@nestjs/common';
import { AiPromptRole } from '@prisma/client';
import { AiSessionMessageRole } from '@prisma/client';
import type { Conversation, Turn } from '../core';
import { chatMessageFromTurn } from '../core';
@@ -16,7 +16,6 @@ export type CanonicalConversationHistory = {
conversation: Conversation;
turns: Turn[];
prompt: ResolvedPrompt;
tokenCost: number;
};
export type CanonicalConversationMeta = Omit<
@@ -35,7 +34,7 @@ export class CompatHistoryProjector {
private projectSessionBase(
history: CanonicalConversationMeta
): Omit<ChatHistory, 'messages'> {
const { conversation, prompt, tokenCost } = history;
const { conversation, prompt } = history;
return {
userId: conversation.userId,
sessionId: conversation.id,
@@ -45,10 +44,7 @@ export class CompatHistoryProjector {
pinned: conversation.pinned,
title: conversation.title,
action: prompt.action || null,
model: prompt.model,
optionalModels: prompt.optionalModels || [],
promptName: prompt.name,
tokens: tokenCost,
createdAt: conversation.createdAt,
updatedAt: conversation.updatedAt,
};
@@ -84,7 +80,7 @@ export class CompatHistoryProjector {
.concat(messages)
.filter(
message =>
message.role !== AiPromptRole.user ||
message.role !== AiSessionMessageRole.user ||
!!message.content.trim() ||
!!message.attachments?.length
)
packages/backend/server/src/plugins/copilot/compat/history-prompt-preload-projector.ts
+4 -2
View File
@@ -1,5 +1,5 @@
import { Injectable } from '@nestjs/common';
import { AiPromptRole } from '@prisma/client';
import { AiSessionMessageRole } from '@prisma/client';
import { PromptService } from '../prompt/service';
import type { ChatMessage } from '../types';
@@ -24,7 +24,9 @@ export class HistoryPromptPreloadProjector {
history.turns[0] ? history.turns[0].metadata : {},
history.conversation.id
)
.filter(({ role }) => role !== AiPromptRole.system) as ChatMessage[];
.filter(
({ role }) => role !== AiSessionMessageRole.system
) as ChatMessage[];
preload.forEach((message, index) => {
message.createdAt = new Date(
packages/backend/server/src/plugins/copilot/config.ts
+10 -172
View File
@@ -5,32 +5,9 @@ import {
StorageJSONSchema,
StorageProviderConfig,
} from '../../base';
import {
AnthropicOfficialConfig,
AnthropicVertexConfig,
} from './providers/anthropic';
import { CloudflareWorkersAIConfig } from './providers/cloudflare';
import type { FalConfig } from './providers/fal';
import { GeminiGenerativeConfig, GeminiVertexConfig } from './providers/gemini';
import { OpenAIConfig } from './providers/openai';
import {
CopilotProviderType,
ModelOutputType,
VertexSchema,
} from './providers/types';
import { CopilotProviderType } from './providers/types';
export type CopilotProviderConfigMap = {
[CopilotProviderType.OpenAI]: OpenAIConfig;
[CopilotProviderType.CloudflareWorkersAi]: CloudflareWorkersAIConfig;
[CopilotProviderType.FAL]: FalConfig;
[CopilotProviderType.Gemini]: GeminiGenerativeConfig;
[CopilotProviderType.GeminiVertex]: GeminiVertexConfig;
[CopilotProviderType.Anthropic]: AnthropicOfficialConfig;
[CopilotProviderType.AnthropicVertex]: AnthropicVertexConfig;
};
export type ProviderSpecificConfig =
CopilotProviderConfigMap[keyof CopilotProviderConfigMap];
export type ProviderSpecificConfig = Record<string, unknown>;
export const RustRequestMiddlewareValues = [
'normalize_messages',
@@ -65,24 +42,13 @@ type CopilotProviderProfileCommon = {
displayName?: string;
priority?: number;
enabled?: boolean;
models?: string[];
models: string[];
middleware?: ProviderMiddlewareConfig;
};
type CopilotProviderProfileVariant<T extends CopilotProviderType> = {
type: T;
config: CopilotProviderConfigMap[T];
};
export type CopilotProviderProfile = CopilotProviderProfileCommon &
{
[Type in CopilotProviderType]: CopilotProviderProfileVariant<Type>;
}[CopilotProviderType];
export type CopilotProviderDefaults = Partial<
Record<Exclude<ModelOutputType, typeof ModelOutputType.Rerank>, string>
> & {
fallback?: string;
export type CopilotProviderProfile = CopilotProviderProfileCommon & {
type: CopilotProviderType;
config: ProviderSpecificConfig;
};
const CopilotProviderProfileBaseShape = z.object({
@@ -90,7 +56,7 @@ const CopilotProviderProfileBaseShape = z.object({
displayName: z.string().optional(),
priority: z.number().optional(),
enabled: z.boolean().optional(),
models: z.array(z.string()).optional(),
models: z.array(z.string().min(1)).min(1),
middleware: z
.object({
rust: z
@@ -106,79 +72,9 @@ const CopilotProviderProfileBaseShape = z.object({
.optional(),
});
const OpenAIConfigShape = z.object({
apiKey: z.string(),
baseURL: z.string().optional(),
oldApiStyle: z.boolean().optional(),
});
const FalConfigShape = z.object({
apiKey: z.string(),
});
const CloudflareWorkersAIConfigShape = z.object({
apiToken: z.string(),
accountId: z.string().optional(),
baseURL: z.string().optional(),
});
const GeminiGenerativeConfigShape = z.object({
apiKey: z.string(),
baseURL: z.string().optional(),
});
const VertexProviderConfigShape = z.object({
location: z.string().optional(),
project: z.string().optional(),
baseURL: z.string().optional(),
googleAuthOptions: z.any().optional(),
fetch: z.any().optional(),
});
const AnthropicOfficialConfigShape = z.object({
apiKey: z.string(),
baseURL: z.string().optional(),
});
const CopilotProviderProfileShape = z.discriminatedUnion('type', [
CopilotProviderProfileBaseShape.extend({
type: z.literal(CopilotProviderType.OpenAI),
config: OpenAIConfigShape,
}),
CopilotProviderProfileBaseShape.extend({
type: z.literal(CopilotProviderType.FAL),
config: FalConfigShape,
}),
CopilotProviderProfileBaseShape.extend({
type: z.literal(CopilotProviderType.CloudflareWorkersAi),
config: CloudflareWorkersAIConfigShape,
}),
CopilotProviderProfileBaseShape.extend({
type: z.literal(CopilotProviderType.Gemini),
config: GeminiGenerativeConfigShape,
}),
CopilotProviderProfileBaseShape.extend({
type: z.literal(CopilotProviderType.GeminiVertex),
config: VertexProviderConfigShape,
}),
CopilotProviderProfileBaseShape.extend({
type: z.literal(CopilotProviderType.Anthropic),
config: AnthropicOfficialConfigShape,
}),
CopilotProviderProfileBaseShape.extend({
type: z.literal(CopilotProviderType.AnthropicVertex),
config: VertexProviderConfigShape,
}),
]);
const CopilotProviderDefaultsShape = z.object({
[ModelOutputType.Text]: z.string().optional(),
[ModelOutputType.Object]: z.string().optional(),
[ModelOutputType.Embedding]: z.string().optional(),
[ModelOutputType.Image]: z.string().optional(),
[ModelOutputType.Rerank]: z.string().optional(),
[ModelOutputType.Structured]: z.string().optional(),
fallback: z.string().optional(),
const CopilotProviderProfileShape = CopilotProviderProfileBaseShape.extend({
type: z.nativeEnum(CopilotProviderType),
config: z.record(z.string(), z.unknown()),
});
declare global {
@@ -202,14 +98,6 @@ declare global {
storage: ConfigItem<StorageProviderConfig>;
providers: {
profiles: ConfigItem<CopilotProviderProfile[]>;
defaults: ConfigItem<CopilotProviderDefaults>;
openai: ConfigItem<OpenAIConfig>;
cloudflareWorkersAi: ConfigItem<CloudflareWorkersAIConfig>;
fal: ConfigItem<FalConfig>;
gemini: ConfigItem<GeminiGenerativeConfig>;
geminiVertex: ConfigItem<GeminiVertexConfig>;
anthropic: ConfigItem<AnthropicOfficialConfig>;
anthropicVertex: ConfigItem<AnthropicVertexConfig>;
};
};
}
@@ -245,56 +133,6 @@ defineModuleConfig('copilot', {
default: [],
shape: z.array(CopilotProviderProfileShape),
},
'providers.defaults': {
desc: 'The default provider ids for model output types and global fallback.',
default: {},
shape: CopilotProviderDefaultsShape,
},
'providers.openai': {
desc: 'The config for the openai provider.',
default: {
apiKey: '',
baseURL: 'https://api.openai.com/v1',
},
link: 'https://github.com/openai/openai-node',
},
'providers.cloudflareWorkersAi': {
desc: 'The config for the Cloudflare Workers AI provider.',
default: {
apiToken: '',
accountId: '',
},
},
'providers.fal': {
desc: 'The config for the fal provider.',
default: {
apiKey: '',
},
},
'providers.gemini': {
desc: 'The config for the gemini provider.',
default: {
apiKey: '',
baseURL: 'https://generativelanguage.googleapis.com/v1beta',
},
},
'providers.geminiVertex': {
desc: 'The config for the gemini provider in Google Vertex AI.',
default: {},
schema: VertexSchema,
},
'providers.anthropic': {
desc: 'The config for the anthropic provider.',
default: {
apiKey: '',
baseURL: 'https://api.anthropic.com/v1',
},
},
'providers.anthropicVertex': {
desc: 'The config for the anthropic provider in Google Vertex AI.',
default: {},
schema: VertexSchema,
},
unsplash: {
desc: 'The config for the unsplash key.',
default: {
packages/backend/server/src/plugins/copilot/context/realtime.ts
+6 -2
View File
@@ -1,7 +1,8 @@
import { Injectable, OnModuleInit } from '@nestjs/common';
import { z } from 'zod';
import { OnEvent } from '../../../base';
import { Config } from '../../../base/config';
import { OnEvent } from '../../../base/event';
import { PermissionAccess } from '../../../core/permission';
import {
RealtimePublisher,
@@ -10,6 +11,7 @@ import {
registerRealtimeLiveQuery,
} from '../../../core/realtime';
import { Models } from '../../../models';
import { assertCopilotEnabled } from '../availability';
export function workspaceEmbeddingRoom(workspaceId: string) {
return realtimeWorkspaceEmbeddingProgressRoom(workspaceId);
@@ -21,7 +23,8 @@ export class CopilotEmbeddingRealtimeProvider implements OnModuleInit {
private readonly ac: PermissionAccess,
private readonly models: Models,
private readonly registry: RealtimeRegistry,
private readonly publisher: RealtimePublisher
private readonly publisher: RealtimePublisher,
private readonly config: Config
) {}
onModuleInit() {
@@ -118,6 +121,7 @@ export class CopilotEmbeddingRealtimeProvider implements OnModuleInit {
}
private async assertCopilot(userId: string, workspaceId: string) {
assertCopilotEnabled(this.config);
await this.ac
.user(userId)
.workspace(workspaceId)
packages/backend/server/src/plugins/copilot/context/resolver.ts
+3
View File
@@ -50,6 +50,7 @@ import {
Models,
} from '../../../models';
import { CopilotEmbeddingJob } from '../embedding/job';
import { CopilotEnabled } from '../feature';
import { COPILOT_LOCKER, CopilotType } from '../resolver';
import { ChatSessionService } from '../session';
import { CopilotStorage } from '../storage';
@@ -286,6 +287,7 @@ class ContextMatchedDocChunk implements DocChunkSimilarity {
}
@Throttle()
@CopilotEnabled()
@Resolver(() => CopilotType)
export class CopilotContextRootResolver {
constructor(
@@ -435,6 +437,7 @@ export class CopilotContextRootResolver {
}
@Throttle()
@CopilotEnabled()
@Resolver(() => CopilotContextType)
export class CopilotContextResolver {
constructor(
packages/backend/server/src/plugins/copilot/context/service.ts
+1
View File
@@ -1,3 +1,4 @@
/* oxlint-disable import/no-cycle -- Context embedding reuses the shared capability runtime. */
import { Injectable, OnApplicationBootstrap } from '@nestjs/common';
import {
packages/backend/server/src/plugins/copilot/controller.ts
+2
View File
@@ -36,6 +36,7 @@ import {
UnsplashIsNotConfigured,
} from '../../base';
import { CurrentUser, Public } from '../../core/auth';
import { CopilotEnabled } from './feature';
import {
ActionStreamHost,
projectActionEventToChatEvent,
@@ -52,6 +53,7 @@ export interface ChatEvent {
const PING_INTERVAL = 5000;
@CopilotEnabled()
@Controller('/api/copilot')
export class CopilotController implements BeforeApplicationShutdown {
private readonly logger = new Logger(CopilotController.name);
packages/backend/server/src/plugins/copilot/conversation/store.ts
-9
View File
@@ -83,7 +83,6 @@ export class ConversationStore {
conversation: Conversation;
turns: Turn[];
promptName: string;
tokenCost: number;
}
| undefined
> {
@@ -96,7 +95,6 @@ export class ConversationStore {
conversation: this.toConversation(session),
turns: this.toTurns(session),
promptName: session.promptName,
tokenCost: session.tokenCost,
};
}
@@ -104,7 +102,6 @@ export class ConversationStore {
| {
conversation: Conversation;
promptName: string;
tokenCost: number;
}
| undefined
> {
@@ -124,7 +121,6 @@ export class ConversationStore {
updatedAt: session.updatedAt,
},
promptName: session.promptName,
tokenCost: session.tokenCost,
};
}
@@ -146,7 +142,6 @@ export class ConversationStore {
turnFromChatMessage(message, session.id)
),
promptName: session.promptName,
tokenCost: session.tokenCost,
}));
}
@@ -168,14 +163,12 @@ export class ConversationStore {
updatedAt: session.updatedAt,
} satisfies Conversation,
promptName: session.promptName,
tokenCost: session.tokenCost,
}));
}
async appendTurns(input: {
sessionId: string;
userId: string;
prompt: { model: string };
turns: Turn[];
}) {
return await this.models.copilotSession.updateMessages({
@@ -190,14 +183,12 @@ export class ConversationStore {
async appendTurn(input: {
sessionId: string;
userId: string;
prompt: { model: string };
turn: Turn;
compatSubmissionId?: string;
}) {
const message = await this.models.copilotSession.appendMessage({
sessionId: input.sessionId,
userId: input.userId,
prompt: input.prompt,
message: (() => {
const { id: _id, ...message } = chatMessageFromTurn(input.turn);
return { ...message, compatSubmissionId: input.compatSubmissionId };
packages/backend/server/src/plugins/copilot/cron.ts
+6 -3
View File
@@ -5,6 +5,7 @@ import { JOB_SIGNAL, JobQueue, OneDay, OnJob } from '../../base';
import { Models } from '../../models';
const CLEANUP_EMBEDDING_JOB_BATCH_SIZE = 100;
const BACKGROUND_COPILOT_JOB_PRIORITY = 100;
declare global {
interface Jobs {
@@ -71,9 +72,11 @@ export class CopilotCronJobs {
const sessions = await this.models.copilotSession.toBeGenerateTitle();
for (const session of sessions) {
await this.jobs.add('copilot.session.generateTitle', {
sessionId: session.id,
});
await this.jobs.add(
'copilot.session.generateTitle',
{ sessionId: session.id },
{ priority: BACKGROUND_COPILOT_JOB_PRIORITY }
);
}
this.logger.log(
`Scheduled title generation for ${sessions.length} sessions`
packages/backend/server/src/plugins/copilot/embedding/client.ts
+20 -24
View File
@@ -1,6 +1,7 @@
/* oxlint-disable import/no-cycle -- Embedding delegates to the shared capability runtime. */
import { createHash } from 'node:crypto';
import { Injectable, Logger } from '@nestjs/common';
import { forwardRef, Inject, Injectable, Logger } from '@nestjs/common';
import { CopilotFailedToGenerateEmbedding } from '../../../base/error/errors.gen';
import {
@@ -10,7 +11,6 @@ import {
} from '../../../models';
import { type CopilotRerankRequest } from '../providers/types';
import { CapabilityRuntime } from '../runtime/capability-runtime';
import { TaskPolicy } from '../runtime/task-policy';
import {
type EmbeddingCallOptionsInput,
EmbeddingClient,
@@ -18,20 +18,20 @@ import {
type ReRankResult,
} from './types';
type EmbeddingRuntime = Pick<
CapabilityRuntime,
'embeddingConfigured' | 'embed' | 'rerank'
>;
class ProductionEmbeddingClient extends EmbeddingClient {
private readonly logger = new Logger(ProductionEmbeddingClient.name);
constructor(
private readonly taskPolicy: TaskPolicy,
private readonly runtime: CapabilityRuntime
) {
constructor(private readonly runtime: EmbeddingRuntime) {
super();
}
override async configured(): Promise<boolean> {
const result = await this.runtime.embeddingConfigured(
this.taskPolicy.resolveEmbeddingModelId()
);
const result = await this.runtime.embeddingConfigured('route-selected');
if (!result) {
this.logger.warn(
'Copilot embedding client is not configured properly, please check your configuration.'
@@ -45,7 +45,7 @@ class ProductionEmbeddingClient extends EmbeddingClient {
options?: EmbeddingCallOptionsInput
): Promise<Embedding[]> {
const normalizedOptions = normalizeEmbeddingCallOptions(options);
const modelId = this.taskPolicy.resolveEmbeddingModelId();
const modelId = 'route-selected';
const embeddings = await this.runtime.embed(modelId, input, {
dimensions: EMBEDDING_DIMENSIONS,
signal: normalizedOptions.signal,
@@ -94,17 +94,13 @@ class ProductionEmbeddingClient extends EmbeddingClient {
})),
};
const ranks = await this.runtime.rerank(
this.taskPolicy.resolveRerankModelId(),
rerankRequest,
{
signal: normalizedOptions.signal,
user: normalizedOptions.userId,
workspace: normalizedOptions.workspaceId,
byokLeaseId: normalizedOptions.byokLeaseId,
featureKind: 'rerank',
}
);
const ranks = await this.runtime.rerank('route-selected', rerankRequest, {
signal: normalizedOptions.signal,
user: normalizedOptions.userId,
workspace: normalizedOptions.workspaceId,
byokLeaseId: normalizedOptions.byokLeaseId,
featureKind: 'rerank',
});
try {
return ranks.map((score, i) => {
@@ -206,12 +202,12 @@ export class CopilotEmbeddingClientService {
private client: EmbeddingClient | undefined;
constructor(
private readonly taskPolicy: TaskPolicy,
private readonly runtime: CapabilityRuntime
@Inject(forwardRef(() => CapabilityRuntime))
private readonly runtime: EmbeddingRuntime
) {}
async refresh() {
const client = new ProductionEmbeddingClient(this.taskPolicy, this.runtime);
const client = new ProductionEmbeddingClient(this.runtime);
await client.configured();
this.client = client;
return this.client;
packages/backend/server/src/plugins/copilot/feature.ts
+54
View File
@@ -0,0 +1,54 @@
import { CanActivate, Injectable, UseGuards } from '@nestjs/common';
import { Config } from '../../base/config';
import { OnEvent } from '../../base/event';
import { ServerFeature, ServerService } from '../../core/config';
import { assertCopilotEnabled } from './availability';
@Injectable()
export class CopilotFeatureService {
constructor(
private readonly config: Config,
private readonly server: ServerService
) {}
get enabled() {
return this.config.copilot.enabled;
}
@OnEvent('config.init')
onConfigInit() {
this.syncServerFeature();
}
@OnEvent('config.changed')
onConfigChanged(event: Events['config.changed']) {
if ('copilot' in event.updates) {
this.syncServerFeature();
}
}
assertEnabled() {
assertCopilotEnabled(this.config);
}
private syncServerFeature() {
if (this.enabled) {
this.server.enableFeature(ServerFeature.Copilot);
} else {
this.server.disableFeature(ServerFeature.Copilot);
}
}
}
@Injectable()
export class CopilotFeatureGuard implements CanActivate {
constructor(private readonly feature: CopilotFeatureService) {}
canActivate() {
this.feature.assertEnabled();
return true;
}
}
export const CopilotEnabled = () => UseGuards(CopilotFeatureGuard);
packages/backend/server/src/plugins/copilot/index.ts
+12 -4
View File
@@ -11,6 +11,7 @@ import { StorageModule } from '../../core/storage';
import { WorkspaceModule } from '../../core/workspaces';
import { IndexerModule } from '../indexer';
import { CopilotController } from './controller';
import { CopilotFeatureGuard, CopilotFeatureService } from './feature';
import { WorkspaceMcpController } from './mcp/controller';
import { McpCredentialService } from './mcp/credential';
import { McpCredentialResolver } from './mcp/resolver';
@@ -34,20 +35,27 @@ const COPILOT_SHARED_IMPORTS = [
];
@Module({
imports: [...COPILOT_SHARED_IMPORTS],
imports: [ServerConfigModule],
providers: [CopilotFeatureService, CopilotFeatureGuard],
exports: [CopilotFeatureService, CopilotFeatureGuard],
})
export class CopilotAvailabilityModule {}
@Module({
imports: [...COPILOT_SHARED_IMPORTS, CopilotAvailabilityModule],
providers: [...COPILOT_KERNEL_PROVIDERS],
exports: [...COPILOT_KERNEL_PROVIDERS],
exports: [CopilotAvailabilityModule, ...COPILOT_KERNEL_PROVIDERS],
})
export class CopilotKernelModule {}
@Module({
imports: [PermissionModule],
imports: [PermissionModule, CopilotAvailabilityModule],
providers: [...COPILOT_TRANSCRIPT_REALTIME_PROVIDERS],
})
export class CopilotRealtimeModule {}
@Module({
imports: [PermissionModule],
imports: [PermissionModule, CopilotAvailabilityModule],
providers: [...COPILOT_CONTEXT_REALTIME_PROVIDERS],
})
export class CopilotEmbeddingRealtimeModule {}
packages/backend/server/src/plugins/copilot/mcp/controller.ts
+2
View File
@@ -16,6 +16,7 @@ import type { Request, Response } from 'express';
import { ActionForbidden, Throttle } from '../../../base';
import { Public } from '../../../core/auth';
import { extractTokenFromHeader } from '../../../core/auth/input';
import { CopilotEnabled } from '../feature';
import { McpCredentialService } from './credential';
import { WorkspaceMcpProvider, type WorkspaceMcpServer } from './provider';
@@ -46,6 +47,7 @@ const SUPPORTED_PROTOCOL_VERSIONS = new Set([
'2024-10-07',
]);
@CopilotEnabled()
@Controller('/api/workspaces/:workspaceId/mcp')
export class WorkspaceMcpController {
private readonly logger = new Logger(WorkspaceMcpController.name);
packages/backend/server/src/plugins/copilot/mcp/resolver.ts
+2
View File
@@ -15,6 +15,7 @@ import { McpAccessMode } from '@prisma/client';
import { CurrentUser } from '../../../core/auth';
import { PermissionAccess } from '../../../core/permission';
import { CopilotEnabled } from '../feature';
import { McpCredentialService } from './credential';
registerEnumType(McpAccessMode, { name: 'McpAccessMode' });
@@ -89,6 +90,7 @@ class CreateMcpCredentialInput {
expirationDays!: number;
}
@CopilotEnabled()
@Resolver()
export class McpCredentialResolver {
constructor(
packages/backend/server/src/plugins/copilot/module-providers.ts
+4 -34
View File
@@ -1,9 +1,4 @@
import { CopilotAccessPolicy } from './access';
import {
ByokEntitlementPolicy,
ByokService,
WorkspaceByokResolver,
} from './byok';
import { ByokEntitlementPolicy, 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';
@@ -25,30 +20,18 @@ import {
} from './embedding';
import { WorkspaceMcpProvider } from './mcp/provider';
import { PromptService } from './prompt';
import {
CopilotProviderFactory,
CopilotProviderLifecycleService,
CopilotProviderRegistryService,
CopilotProviders,
} from './providers';
import { CopilotResolver, UserCopilotResolver } from './resolver';
import { ActionRuntimeBridge } from './runtime/action-runtime-bridge';
import { CapabilityRuntime } from './runtime/capability-runtime';
import { CopilotExecutionMetrics } from './runtime/execution-metrics';
import { ExecutionPlanBuilder } from './runtime/execution-plan';
import { CopilotRuntimeEventConsumer } from './runtime/copilot-runtime-event-consumer';
import { ActionStreamHost } from './runtime/hosts/action-stream-host';
import { AttachmentAdmissionHost } from './runtime/hosts/attachment-admission';
import { AttachmentMaterializer } from './runtime/hosts/attachment-materializer';
import { CapabilityPolicyHost } from './runtime/hosts/capability-policy-host';
import { ConversationHost } from './runtime/hosts/conversation-host';
import { ImageResultHost } from './runtime/hosts/image-result-host';
import { ResponsePostprocessor } from './runtime/hosts/response-postprocessor';
import { ToolExecutorHost } from './runtime/hosts/tool-executor-host';
import { TurnPersistence } from './runtime/hosts/turn-persistence';
import { ModelSelectionPolicy } from './runtime/model-selection-policy';
import { NativeExecutionEngine } from './runtime/native-execution-engine';
import { PromptRuntime } from './runtime/prompt-runtime';
import { TaskPolicy } from './runtime/task-policy';
import { ToolRuntime } from './runtime/tool-runtime';
import { TurnOrchestrator } from './runtime/turn-orchestrator';
import { ChatSessionService } from './session';
@@ -65,21 +48,14 @@ import {
CopilotWorkspaceService,
} from './workspace';
export const COPILOT_PROVIDER_PROVIDERS = [
...CopilotProviders,
CopilotProviderRegistryService,
CopilotProviderFactory,
CopilotProviderLifecycleService,
];
export const COPILOT_PROVIDER_PROVIDERS: [] = [];
export const COPILOT_RUNTIME_PROVIDERS = [
ByokEntitlementPolicy,
ByokService,
ChatSessionService,
ConversationStore,
ConversationInboxService,
ConversationPolicy,
CopilotAccessPolicy,
HistoryAttachmentUrlProjector,
CompatHistoryProjector,
HistoryPromptPreloadProjector,
@@ -88,18 +64,12 @@ export const COPILOT_RUNTIME_PROVIDERS = [
CopilotContextService,
CopilotEmbeddingClientService,
PromptService,
ModelSelectionPolicy,
ActionRuntimeBridge,
CopilotExecutionMetrics,
ExecutionPlanBuilder,
CopilotRuntimeEventConsumer,
PromptRuntime,
CapabilityPolicyHost,
ConversationHost,
CapabilityRuntime,
NativeExecutionEngine,
TaskPolicy,
ToolRuntime,
ToolExecutorHost,
AttachmentMaterializer,
AttachmentAdmissionHost,
ActionStreamHost,
packages/backend/server/src/plugins/copilot/prompt/native-contract.ts
+1 -117
View File
@@ -1,36 +1,17 @@
import {
llmCollectPromptMetadata,
llmCountPromptTokens,
llmGetBuiltInPromptSpec,
llmListBuiltInPromptSpecs,
llmRenderBuiltInPrompt,
llmRenderBuiltInSessionPrompt,
llmRenderPrompt,
llmRenderSessionPrompt,
type NativeBuiltInPromptRenderRequest as NativeBuiltInPromptRenderContract,
type NativeBuiltInPromptSessionRenderRequest as NativeBuiltInPromptSessionContract,
type NativePromptCountTokensRequest as NativePromptTokenCountContract,
type NativePromptCountTokensResponse as NativePromptTokenCountResult,
type NativePromptMetadataRequest as NativePromptMetadataContract,
type NativePromptMetadataResponse as NativePromptMetadataResult,
type NativePromptRenderRequest as NativePromptRenderContract,
type NativePromptRenderResponse as NativePromptRenderResult,
type NativePromptSessionRenderRequest as NativePromptSessionContract,
type NativePromptSessionRenderResponse as NativePromptSessionResult,
} from '../../../native';
import type { PromptMessage, PromptParams } from '../providers/types';
import { projectPromptMessageForNative } from '../runtime/contracts';
import type { PromptSpec } from './spec';
export type NativePromptRenderRequest = Omit<
NativePromptRenderContract,
'messages' | 'templateParams' | 'renderParams'
> & {
messages: PromptMessage[];
templateParams: PromptParams;
renderParams: PromptParams;
};
export type NativePromptRenderResponse = Omit<
NativePromptRenderResult,
'messages'
@@ -45,46 +26,6 @@ export type NativeBuiltInPromptRenderRequest = Omit<
renderParams: PromptParams;
};
export type NativePromptCountTokensRequest = Omit<
NativePromptTokenCountContract,
'messages' | 'model'
> & {
model?: string | null;
messages: Pick<PromptMessage, 'content'>[];
};
export type NativePromptCountTokensResponse = NativePromptTokenCountResult;
export type NativePromptMetadataRequest = Omit<
NativePromptMetadataContract,
'messages'
> & {
messages: PromptMessage[];
};
export type NativePromptMetadataResponse = Omit<
NativePromptMetadataResult,
'templateParams'
> & {
templateParams: PromptParams;
};
export type NativePromptSessionRenderRequest = Omit<
NativePromptSessionContract,
'prompt' | 'turns' | 'renderParams'
> & {
prompt: Omit<
NativePromptSessionContract['prompt'],
'templateParams' | 'messages' | 'model'
> & {
model?: string | null;
templateParams: PromptParams;
messages: PromptMessage[];
};
turns: PromptMessage[];
renderParams: PromptParams;
};
export type NativePromptSessionRenderResponse = Omit<
NativePromptSessionResult,
'messages'
@@ -100,8 +41,7 @@ export type NativeBuiltInPromptSessionRenderRequest = Omit<
renderParams: PromptParams;
};
type NativePromptContractMessage =
NativePromptRenderContract['messages'][number];
type NativePromptContractMessage = NativePromptRenderResult['messages'][number];
function toNativePromptMessage(
message: PromptMessage
@@ -123,22 +63,6 @@ function fromNativePromptMessage(
};
}
export function renderPromptNative(
request: NativePromptRenderRequest
): NativePromptRenderResponse {
const normalizedMessages = request.messages.map(toNativePromptMessage);
const rendered = llmRenderPrompt({
messages: normalizedMessages,
templateParams: request.templateParams,
renderParams: request.renderParams,
});
return {
...rendered,
messages: rendered.messages.map(fromNativePromptMessage),
};
}
export function renderBuiltInPromptNative(
request: NativeBuiltInPromptRenderRequest
): NativePromptRenderResponse {
@@ -153,25 +77,6 @@ export function renderBuiltInPromptNative(
};
}
export function renderPromptSessionNative(
request: NativePromptSessionRenderRequest
): NativePromptSessionRenderResponse {
const rendered = llmRenderSessionPrompt({
...request,
prompt: {
...request.prompt,
messages: request.prompt.messages.map(toNativePromptMessage),
model: request.prompt.model ?? undefined,
},
turns: request.turns.map(toNativePromptMessage),
renderParams: request.renderParams,
});
return {
...rendered,
messages: rendered.messages.map(fromNativePromptMessage),
};
}
export function renderBuiltInPromptSessionNative(
request: NativeBuiltInPromptSessionRenderRequest
): NativePromptSessionRenderResponse {
@@ -187,29 +92,10 @@ export function renderBuiltInPromptSessionNative(
};
}
export function countPromptTokensNative(
request: NativePromptCountTokensRequest
): NativePromptCountTokensResponse {
return llmCountPromptTokens({
...request,
model: request.model ?? undefined,
});
}
export function collectPromptMetadataNative(
request: NativePromptMetadataRequest
): NativePromptMetadataResponse {
return llmCollectPromptMetadata({
messages: request.messages.map(toNativePromptMessage),
});
}
export function listBuiltInPromptSpecsNative(): PromptSpec[] {
return llmListBuiltInPromptSpecs().map(spec => ({
name: spec.name,
action: spec.action,
model: spec.model,
optionalModels: spec.optionalModels,
config: spec.config,
params: spec.params
? Object.fromEntries(
@@ -238,8 +124,6 @@ export function getBuiltInPromptSpecNative(name: string): PromptSpec | null {
return {
name: spec.name,
action: spec.action,
model: spec.model,
optionalModels: spec.optionalModels,
config: spec.config,
params: spec.params
? Object.fromEntries(
packages/backend/server/src/plugins/copilot/prompt/service.ts
+10 -95
View File
@@ -2,15 +2,11 @@ import { Injectable, Logger } from '@nestjs/common';
import type { PromptMessage, PromptParams } from '../providers/types';
import {
collectPromptMetadataNative,
countPromptTokensNative,
getBuiltInPromptSpecNative,
renderBuiltInPromptNative,
renderBuiltInPromptSessionNative,
renderPromptNative,
renderPromptSessionNative,
} from './native-contract';
import type { Prompt, PromptSpec, ResolvedPrompt } from './spec';
import type { PromptSpec, ResolvedPrompt } from './spec';
@Injectable()
export class PromptService {
@@ -20,11 +16,6 @@ export class PromptService {
}
async get(name: string): Promise<ResolvedPrompt | null> {
const compatPrompt = this.lookupCompatPrompt(name);
if (compatPrompt) {
return this.describeCompatPrompt(this.clonePrompt(compatPrompt));
}
const builtInPromptSpec = this.lookupBuiltInPromptSpec(name);
if (!builtInPromptSpec) return null;
@@ -36,17 +27,10 @@ export class PromptService {
params: PromptParams,
sessionId?: string
): PromptMessage[] {
const rendered =
prompt.source === 'built_in'
? renderBuiltInPromptNative({
name: prompt.name,
renderParams: params,
})
: renderPromptNative({
messages: this.requireCompatMessages(prompt),
templateParams: prompt.params,
renderParams: params,
});
const rendered = renderBuiltInPromptNative({
name: prompt.name,
renderParams: params,
});
this.logWarnings(rendered.warnings, sessionId);
return rendered.messages;
@@ -56,38 +40,18 @@ export class PromptService {
prompt: ResolvedPrompt,
turns: PromptMessage[],
params: PromptParams,
maxTokenSize = prompt.config?.maxTokens || 128 * 1024,
sessionId?: string
): PromptMessage[] {
const rendered =
prompt.source === 'built_in'
? renderBuiltInPromptSessionNative({
name: prompt.name,
turns,
renderParams: params,
maxTokenSize,
})
: renderPromptSessionNative({
prompt: {
action: prompt.action,
model: prompt.model,
promptTokens: this.countCompatPromptTokens(prompt),
templateParams: prompt.params,
messages: this.requireCompatMessages(prompt),
},
turns,
renderParams: params,
maxTokenSize,
});
const rendered = renderBuiltInPromptSessionNative({
name: prompt.name,
turns,
renderParams: params,
});
this.logWarnings(rendered.warnings, sessionId);
return rendered.messages;
}
protected lookupCompatPrompt(_name: string): Prompt | null {
return null;
}
protected lookupBuiltInPromptSpec(name: string): PromptSpec | null {
const spec = getBuiltInPromptSpecNative(name);
return spec ? this.clonePromptSpec(spec) : null;
@@ -104,23 +68,9 @@ export class PromptService {
}));
}
protected clonePrompt(prompt: Prompt): Prompt {
return {
...prompt,
optionalModels: prompt.optionalModels
? [...prompt.optionalModels]
: undefined,
config: prompt.config ? structuredClone(prompt.config) : undefined,
messages: this.cloneMessages(prompt.messages),
};
}
protected clonePromptSpec(spec: PromptSpec): PromptSpec {
return {
...spec,
optionalModels: spec.optionalModels
? [...spec.optionalModels]
: undefined,
config: spec.config ? structuredClone(spec.config) : undefined,
params: spec.params ? structuredClone(spec.params) : undefined,
messages: spec.messages.map(message => ({ ...message })),
@@ -132,27 +82,9 @@ export class PromptService {
return {
name: spec.name,
action: spec.action,
model: spec.model,
optionalModels: spec.optionalModels ?? [],
config: spec.config ? structuredClone(spec.config) : undefined,
paramKeys: Object.keys(params),
params,
source: 'built_in',
};
}
private describeCompatPrompt(prompt: Prompt): ResolvedPrompt {
const metadata = collectPromptMetadataNative({ messages: prompt.messages });
return {
name: prompt.name,
action: prompt.action,
model: prompt.model,
optionalModels: prompt.optionalModels ?? [],
config: prompt.config ? structuredClone(prompt.config) : undefined,
paramKeys: metadata.paramKeys,
params: metadata.templateParams,
source: 'compat',
messages: prompt.messages,
};
}
@@ -178,23 +110,6 @@ export class PromptService {
);
}
private countCompatPromptTokens(prompt: ResolvedPrompt): number {
return countPromptTokensNative({
model: prompt.model,
messages: this.requireCompatMessages(prompt).map(message => ({
content: message.content,
})),
}).tokens;
}
private requireCompatMessages(prompt: ResolvedPrompt): PromptMessage[] {
if (prompt.source === 'compat' && prompt.messages) {
return this.cloneMessages(prompt.messages);
}
throw new Error(`Prompt ${prompt.name} does not expose compat messages`);
}
private logWarnings(warnings: string[], sessionId?: string) {
if (!sessionId) {
return;
packages/backend/server/src/plugins/copilot/prompt/spec.ts
+1 -20
View File
@@ -1,28 +1,11 @@
import type {
PromptConfig,
PromptMessage,
PromptParams,
} from '../providers/types';
export type Prompt = {
name: string;
model: string;
optionalModels?: string[];
action?: string;
messages: PromptMessage[];
config?: PromptConfig;
};
import type { PromptConfig, PromptParams } from '../providers/types';
export type ResolvedPrompt = {
name: string;
model: string;
optionalModels: string[];
action?: string;
config?: PromptConfig;
paramKeys: string[];
params: PromptParams;
source: 'built_in' | 'compat';
messages?: PromptMessage[];
};
type PromptParamSpec = {
@@ -38,8 +21,6 @@ type PromptSpecMessage = {
export type PromptSpec = {
name: string;
action?: string;
model: string;
optionalModels?: string[];
config?: PromptConfig;
params?: Record<string, PromptParamSpec>;
messages: PromptSpecMessage[];
packages/backend/server/src/plugins/copilot/providers/anthropic/anthropic.ts
-100
View File
@@ -1,100 +0,0 @@
import { CopilotProviderSideError, UserFriendlyError } from '../../../../base';
import {
type LlmBackendConfig,
llmResolveRequestIntentOptions,
} from '../../../../native';
import { CopilotProvider } from '../provider';
import { hasProviderModelBehaviorFlag } from '../provider-model-runtime';
import {
type CopilotProviderExecution,
type ProviderDriverSpec,
} from '../provider-runtime-contract';
import { CopilotProviderType } from '../types';
import {
getGoogleAuth,
getVertexAnthropicBaseUrl,
type VertexAnthropicProviderConfig,
} from '../utils';
export abstract class AnthropicProvider<T> extends CopilotProvider<T> {
protected resolveModelBackendKind() {
return this.type === CopilotProviderType.AnthropicVertex
? ('anthropic_vertex' as const)
: ('anthropic' as const);
}
override getDriverSpec(): ProviderDriverSpec {
return {
createBackendConfig: execution => this.createNativeConfig(execution),
mapError: error => this.handleError(error),
chat: {
resolveRequestOptions: async context => {
const requestIntent = await llmResolveRequestIntentOptions({
protocol: context.protocol,
backendConfig: context.backendConfig,
reasoning: {
enabled: context.options.reasoning,
supported: hasProviderModelBehaviorFlag(
context.model,
'reasoning_budget_12000'
),
budgetTokens: hasProviderModelBehaviorFlag(
context.model,
'reasoning_budget_12000'
)
? 12000
: undefined,
},
});
return {
attachmentCapability: this.getAttachCapability(
context.model,
context.outputType
),
reasoning: requestIntent.reasoning,
};
},
},
structured: false,
embedding: false,
rerank: false,
};
}
private handleError(e: any) {
if (e instanceof UserFriendlyError) {
return e;
}
return new CopilotProviderSideError({
provider: this.type,
kind: 'unexpected_response',
message: e?.message || 'Unexpected anthropic response',
});
}
private async createNativeConfig(
execution?: CopilotProviderExecution
): Promise<LlmBackendConfig> {
const config = this.getConfig(execution);
if (this.type === CopilotProviderType.AnthropicVertex) {
const vertexConfig = config as VertexAnthropicProviderConfig;
const auth = await getGoogleAuth(vertexConfig, 'anthropic');
const { Authorization: authHeader } = auth.headers();
const token = authHeader.replace(/^Bearer\s+/i, '');
const baseUrl = getVertexAnthropicBaseUrl(vertexConfig) || auth.baseUrl;
return {
base_url: baseUrl || '',
auth_token: token,
headers: { Authorization: authHeader },
};
}
const officialConfig = config as { apiKey: string; baseURL?: string };
const baseUrl = officialConfig.baseURL || 'https://api.anthropic.com/v1';
return {
base_url: baseUrl.replace(/\/v1\/?$/, ''),
auth_token: officialConfig.apiKey,
};
}
}
packages/backend/server/src/plugins/copilot/providers/anthropic/index.ts
-2
View File
@@ -1,2 +0,0 @@
export * from './official';
export * from './vertex';
packages/backend/server/src/plugins/copilot/providers/anthropic/official.ts
-16
View File
@@ -1,16 +0,0 @@
import type { CopilotProviderExecution } from '../provider-runtime-contract';
import { CopilotProviderType } from '../types';
import { AnthropicProvider } from './anthropic';
export type AnthropicOfficialConfig = {
apiKey: string;
baseURL?: string;
};
export class AnthropicOfficialProvider extends AnthropicProvider<AnthropicOfficialConfig> {
override readonly type = CopilotProviderType.Anthropic;
override configured(execution?: CopilotProviderExecution): boolean {
return !!this.getConfig(execution).apiKey;
}
}
packages/backend/server/src/plugins/copilot/providers/anthropic/vertex.ts
-16
View File
@@ -1,16 +0,0 @@
import type { CopilotProviderExecution } from '../provider-runtime-contract';
import { CopilotProviderType } from '../types';
import { getVertexAnthropicBaseUrl, type VertexProviderConfig } from '../utils';
import { AnthropicProvider } from './anthropic';
export type AnthropicVertexConfig = VertexProviderConfig;
export class AnthropicVertexProvider extends AnthropicProvider<AnthropicVertexConfig> {
override readonly type = CopilotProviderType.AnthropicVertex;
override configured(execution?: CopilotProviderExecution): boolean {
const config = this.getConfig(execution);
if (!config.location || !config.googleAuthOptions) return false;
return !!config.project || !!getVertexAnthropicBaseUrl(config);
}
}
packages/backend/server/src/plugins/copilot/providers/attachments.ts
+1 -39
View File
@@ -1,42 +1,4 @@
import type {
ModelAttachmentCapability,
PromptAttachment,
PromptMessage,
} from './types';
export const IMAGE_ATTACHMENT_CAPABILITY: ModelAttachmentCapability = {
kinds: ['image'],
sourceKinds: ['url', 'data'],
allowRemoteUrls: true,
};
export const GEMINI_ATTACHMENT_CAPABILITY: ModelAttachmentCapability = {
kinds: ['image', 'audio', 'file'],
sourceKinds: ['url', 'data', 'bytes', 'file_handle'],
allowRemoteUrls: true,
};
export function promptAttachmentHasSource(
attachment: PromptAttachment
): boolean {
if (typeof attachment === 'string') {
return !!attachment.trim();
}
if ('attachment' in attachment) {
return !!attachment.attachment;
}
switch (attachment.kind) {
case 'url':
return !!attachment.url;
case 'data':
case 'bytes':
return !!attachment.data;
case 'file_handle':
return !!attachment.fileHandle;
}
}
import type { PromptAttachment, PromptMessage } from './types';
export function applyPromptAttachmentMimeTypeHintForNative(
attachment: PromptAttachment,
packages/backend/server/src/plugins/copilot/providers/cloudflare.ts
-65
View File
@@ -1,65 +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 } from './types';
export type CloudflareWorkersAIConfig = {
apiToken: string;
accountId?: string;
baseURL?: string;
};
export class CloudflareWorkersAIProvider extends CopilotProvider<CloudflareWorkersAIConfig> {
override readonly type = CopilotProviderType.CloudflareWorkersAi;
protected resolveModelBackendKind() {
return 'cloudflare_workers_ai' as const;
}
override configured(execution?: CopilotProviderExecution): boolean {
const config = this.getConfig(execution);
return !!config.apiToken && (!!config.accountId || !!config.baseURL);
}
private handleError(e: any) {
if (e instanceof UserFriendlyError) {
return e;
}
return new CopilotProviderSideError({
provider: this.type,
kind: 'unexpected_response',
message: e?.message || 'Unexpected cloudflare workers ai response',
});
}
private createNativeConfig(
execution?: CopilotProviderExecution
): LlmBackendConfig {
const config = this.getConfig(execution);
return {
base_url: this.resolveBaseUrl(execution),
auth_token: config.apiToken,
};
}
private resolveBaseUrl(execution?: CopilotProviderExecution) {
const config = this.getConfig(execution);
if (config.baseURL) {
return config.baseURL.replace(/\/v1\/?$/, '').replace(/\/$/, '');
}
const accountId = config.accountId ?? '';
return `https://api.cloudflare.com/client/v4/accounts/${accountId}/ai`;
}
override getDriverSpec(): ProviderDriverSpec {
return {
createBackendConfig: execution => this.createNativeConfig(execution),
mapError: error => this.handleError(error),
structured: false,
embedding: false,
};
}
}
packages/backend/server/src/plugins/copilot/providers/factory.ts
-527
View File
@@ -1,527 +0,0 @@
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,
} from './provider-registry';
import type {
CopilotProviderExecution,
PreparedNativeEmbeddingExecution,
PreparedNativeExecution,
PreparedNativeImageExecution,
PreparedNativeRerankExecution,
PreparedNativeStructuredExecution,
} from './provider-runtime-contract';
import { CopilotProviderRegistryService } from './registry-service';
import {
type CopilotChatOptions,
type CopilotEmbeddingOptions,
type CopilotImageOptions,
CopilotProviderType,
type CopilotRerankRequest,
type CopilotStructuredOptions,
ModelFullConditions,
ModelOutputType,
type PromptMessage,
} from './types';
export type ResolvedCopilotProvider = {
providerId: string;
provider: CopilotProvider;
execution: CopilotProviderExecution;
profile: NormalizedCopilotProviderProfile;
rawModelId?: string;
modelId?: string;
explicitProviderId?: string;
prepared?: PreparedNativeExecution;
preparedStructured?: PreparedNativeStructuredExecution;
preparedEmbedding?: PreparedNativeEmbeddingExecution;
preparedRerank?: PreparedNativeRerankExecution;
preparedImage?: PreparedNativeImageExecution;
};
type RoutePreparationResult = Partial<
Pick<
ResolvedCopilotProvider,
| 'prepared'
| 'preparedStructured'
| 'preparedEmbedding'
| 'preparedRerank'
| 'preparedImage'
| 'modelId'
>
>;
type EffectiveProviderRegistry = {
byokRegistry: CopilotProviderRegistry;
quotaBackedRegistry: CopilotProviderRegistry;
quotaBackedRoutesAvailable: boolean;
};
@Injectable()
export class CopilotProviderFactory {
constructor(
private readonly server: ServerService,
private readonly registries: CopilotProviderRegistryService,
private readonly access: CopilotAccessPolicy
) {}
private readonly logger = new Logger(CopilotProviderFactory.name);
readonly #providers = new Map<string, CopilotProvider>();
readonly #providerIdsByType = new Map<CopilotProviderType, Set<string>>();
private getRegistry() {
return this.registries.getRegistry();
}
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 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 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
);
}
private async prepareResolvedRoutes(
routes: ResolvedCopilotProvider[],
prepare: (
route: ResolvedCopilotProvider
) => Promise<RoutePreparationResult | null | undefined>
) {
const preparedRoutes = await Promise.all(
routes.map(async route => {
const prepared = await prepare(route);
return prepared ? { ...route, ...prepared } : null;
})
);
return this.filterPreparedRoutes(preparedRoutes);
}
async resolveProvider(
cond: ModelFullConditions,
filter: {
prefer?: CopilotProviderType;
} = {},
context: CopilotAccessContext = {}
): Promise<ResolvedCopilotProvider | 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 { 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.getAvailableProviderIds(registry),
preferredProviderIds: this.getPreferredProviderIds(
registry,
filter.prefer
),
});
const resolved: ResolvedCopilotProvider[] = [];
for (const providerId of route.candidateProviderIds) {
const profile = registry.profiles.get(providerId);
const provider = profile
? this.getProviderByProfile(providerId, profile)
: undefined;
if (!provider || !profile) continue;
const normalizedCond = this.normalizeCond(registry, providerId, cond);
if (
normalizedCond.modelId &&
profile.models?.length &&
!profile.models.includes(normalizedCond.modelId)
) {
continue;
}
const execution = { providerId, profile };
const matched = await provider.match(normalizedCond, execution);
if (!matched) continue;
resolved.push({
providerId,
provider,
execution,
profile,
rawModelId: route.rawModelId,
modelId: normalizedCond.modelId,
explicitProviderId: route.explicitProviderId,
});
}
return resolved;
}
async prepareRoutes(
kind: 'text' | 'streamText' | 'streamObject',
cond: ModelFullConditions,
messages: PromptMessage[],
options: CopilotChatOptions = {},
filter: {
prefer?: CopilotProviderType;
} = {}
): Promise<ResolvedCopilotProvider[]> {
const routes = await this.resolveRoutes(
cond,
filter,
this.getRequestContext(options)
);
return await this.prepareResolvedRoutes(routes, async route => {
const prepared = await getProviderRuntimeHost(
route.provider
).prepare.chat(
kind,
{ ...cond, modelId: route.modelId },
messages,
options,
route.execution
);
const normalizedPrepared = prepared?.route ? prepared : undefined;
if (!normalizedPrepared) {
return null;
}
return {
modelId: normalizedPrepared.route.model,
prepared: normalizedPrepared,
};
});
}
async prepareStructuredRoutes(
cond: ModelFullConditions,
messages: PromptMessage[],
options: CopilotStructuredOptions = {},
filter: {
prefer?: CopilotProviderType;
} = {},
responseContract?: RequiredStructuredOutputContract
): Promise<ResolvedCopilotProvider[]> {
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(
{ ...cond, modelId: route.modelId },
messages,
options,
responseContract,
route.execution
)) ?? undefined;
if (!preparedStructured) {
return null;
}
return {
modelId: preparedStructured.route.model,
preparedStructured,
};
});
}
async prepareEmbeddingRoutes(
modelId: string,
input: string | string[],
options: CopilotEmbeddingOptions = {}
): Promise<ResolvedCopilotProvider[]> {
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(
{ modelId: route.modelId },
input,
options,
route.execution
)) ?? undefined;
if (!preparedEmbedding) {
return null;
}
return {
modelId: preparedEmbedding.route.model,
preparedEmbedding,
};
});
}
async prepareRerankRoutes(
modelId: string,
request: CopilotRerankRequest,
options: CopilotChatOptions = {}
): Promise<ResolvedCopilotProvider[]> {
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(
{ modelId: route.modelId },
request,
options,
route.execution
)) ?? undefined;
if (!preparedRerank) {
return null;
}
return {
modelId: preparedRerank.route.model,
preparedRerank,
};
});
}
async prepareImageRoutes(
cond: ModelFullConditions,
messages: PromptMessage[],
options: CopilotImageOptions = {},
filter: {
prefer?: CopilotProviderType;
} = {}
): Promise<ResolvedCopilotProvider[]> {
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(
{ ...cond, modelId: route.modelId },
messages,
options,
route.execution
)) ?? undefined;
if (!preparedImage) {
return null;
}
return {
modelId: preparedImage.route.model,
preparedImage,
};
});
}
async getProvider(
cond: ModelFullConditions,
filter: {
prefer?: CopilotProviderType;
} = {}
): Promise<CopilotProvider | null> {
return (await this.resolveProvider(cond, filter))?.provider ?? null;
}
async getProviderByModel(
modelId: string,
filter: {
prefer?: CopilotProviderType;
} = {}
): Promise<CopilotProvider | null> {
this.logger.debug(`Resolving copilot provider for model: ${modelId}`);
return this.getProvider({ modelId }, filter);
}
register(providerId: string, provider: CopilotProvider) {
const existed = this.#providers.get(providerId);
if (existed?.type && existed.type !== provider.type) {
const ids = this.#providerIdsByType.get(existed.type);
ids?.delete(providerId);
if (!ids?.size) {
this.#providerIdsByType.delete(existed.type);
}
}
this.#providers.set(providerId, provider);
const ids = this.#providerIdsByType.get(provider.type) ?? new Set<string>();
ids.add(providerId);
this.#providerIdsByType.set(provider.type, ids);
this.logger.log(
`Copilot provider [${provider.type}] registered as [${providerId}].`
);
this.server.enableFeature(ServerFeature.Copilot);
}
unregister(providerId: string, provider: CopilotProvider) {
const existed = this.#providers.get(providerId);
if (!existed || existed !== provider) {
return;
}
this.#providers.delete(providerId);
const ids = this.#providerIdsByType.get(provider.type);
ids?.delete(providerId);
if (!ids?.size) {
this.#providerIdsByType.delete(provider.type);
}
this.logger.log(
`Copilot provider [${provider.type}] unregistered from [${providerId}].`
);
if (this.#providers.size === 0) {
this.server.disableFeature(ServerFeature.Copilot);
}
}
}
packages/backend/server/src/plugins/copilot/providers/fal.ts
-59
View File
@@ -1,59 +0,0 @@
import { Injectable } from '@nestjs/common';
import { CopilotProviderSideError, UserFriendlyError } from '../../../base';
import { CopilotProvider } from './provider';
import type {
CopilotProviderExecution,
ProviderDriverSpec,
} from './provider-runtime-contract';
import { CopilotProviderType } from './types';
export type FalConfig = {
apiKey: string;
};
@Injectable()
export class FalProvider extends CopilotProvider<FalConfig> {
override type = CopilotProviderType.FAL;
protected resolveModelBackendKind() {
return 'fal' as const;
}
override configured(execution?: CopilotProviderExecution): boolean {
return !!this.getConfig(execution).apiKey;
}
private createNativeConfig(execution?: CopilotProviderExecution) {
return {
base_url: 'https://fal.run',
auth_token: this.getConfig(execution).apiKey,
};
}
override getDriverSpec(): ProviderDriverSpec {
return {
createBackendConfig: execution => this.createNativeConfig(execution),
mapError: error => this.handleError(error),
chat: false,
structured: false,
embedding: false,
rerank: false,
image: {},
};
}
private handleError(e: any) {
if (e instanceof UserFriendlyError) {
// pass through user friendly errors
return e;
} else {
const error = new CopilotProviderSideError({
provider: this.type,
kind: 'unexpected_response',
message: e?.message || 'Unexpected fal response',
});
return error;
}
}
}
packages/backend/server/src/plugins/copilot/providers/gemini/gemini.ts
-250
View File
@@ -1,250 +0,0 @@
import { setTimeout as delay } from 'node:timers/promises';
import { Inject } from '@nestjs/common';
import { ZodError } from 'zod';
import {
CopilotProviderSideError,
OneMB,
UserFriendlyError,
} from '../../../../base';
import {
isInvalidStructuredOutputError,
type LlmBackendConfig,
llmResolveRequestIntentOptions,
} from '../../../../native';
import {
admittedAttachmentToPromptAttachment,
AttachmentAdmissionHost,
} from '../../runtime/hosts/attachment-admission';
import {
planAdmittedAttachmentMaterialization,
planHostUrlAttachmentMaterialization,
} from '../../runtime/hosts/attachment-materialization-planner';
import { AttachmentMaterializer } from '../../runtime/hosts/attachment-materializer';
import { CopilotProvider } from '../provider';
import { hasProviderModelBehaviorFlag } from '../provider-model-runtime';
import {
type CopilotProviderExecution,
type ProviderDriverSpec,
} from '../provider-runtime-contract';
import type { PromptAttachment, PromptMessage } from '../types';
import { promptAttachmentMimeType, promptAttachmentToUrl } from '../utils';
export const DEFAULT_DIMENSIONS = 256;
const GEMINI_REMOTE_ATTACHMENT_MAX_BYTES = 64 * OneMB;
const TRUSTED_ATTACHMENT_HOST_SUFFIXES = ['cdn.affine.pro'];
const GEMINI_RETRY_INITIAL_DELAY_MS = 2_000;
function normalizeMimeType(mediaType?: string) {
return mediaType?.split(';', 1)[0]?.trim() || 'application/octet-stream';
}
export abstract class GeminiProvider<T> extends CopilotProvider<T> {
@Inject() protected readonly attachmentMaterializer!: AttachmentMaterializer;
@Inject()
protected readonly attachmentAdmissionHost?: AttachmentAdmissionHost;
protected resolveModelBackendKind() {
return this.type === 'geminiVertex'
? ('gemini_vertex' as const)
: ('gemini_api' as const);
}
protected abstract createNativeConfig(
execution?: CopilotProviderExecution
): Promise<LlmBackendConfig>;
private handleError(e: any) {
if (e instanceof UserFriendlyError) {
return e;
} else {
return new CopilotProviderSideError({
provider: this.type,
kind: 'unexpected_response',
message: e?.message || 'Unexpected google response',
});
}
}
private getAttachmentAdmissionHost() {
return (
this.attachmentAdmissionHost ??
new AttachmentAdmissionHost(this.attachmentMaterializer)
);
}
protected async prepareMessages(
messages: PromptMessage[],
backendConfig: LlmBackendConfig,
options?: {
signal?: AbortSignal;
user?: string;
workspace?: string;
session?: string;
}
): Promise<PromptMessage[]> {
const prepared: PromptMessage[] = [];
for (const message of messages) {
options?.signal?.throwIfAborted();
if (!Array.isArray(message.attachments) || !message.attachments.length) {
prepared.push(message);
continue;
}
const attachments: PromptAttachment[] = [];
let changed = false;
for (const attachment of message.attachments) {
options?.signal?.throwIfAborted();
const rawUrl = promptAttachmentToUrl(attachment);
if (!rawUrl || rawUrl.startsWith('data:')) {
attachments.push(attachment);
continue;
}
try {
new URL(rawUrl);
} catch {
attachments.push(attachment);
continue;
}
const declaredMimeType = promptAttachmentMimeType(
attachment,
typeof message.params?.mimetype === 'string'
? message.params.mimetype
: undefined
);
const referencePlan = await planHostUrlAttachmentMaterialization(
'gemini',
backendConfig,
{
attachmentId: rawUrl,
url: rawUrl,
expectedMime: declaredMimeType
? normalizeMimeType(declaredMimeType)
: undefined,
maxSize: GEMINI_REMOTE_ATTACHMENT_MAX_BYTES,
}
);
if (referencePlan.mode === 'remote_reference') {
attachments.push(attachment);
continue;
}
const admitted =
await this.getAttachmentAdmissionHost().admitPromptAttachment(
attachment,
{
userId: options?.user ?? 'provider-runtime',
workspaceId: options?.workspace ?? 'provider-runtime',
sessionId: options?.session,
signal: options?.signal,
maxBytes: referencePlan.request.maxSize,
trustedHostSuffixes: TRUSTED_ATTACHMENT_HOST_SUFFIXES,
}
);
const materialization = planAdmittedAttachmentMaterialization(admitted);
attachments.push(
materialization.mode === 'inline'
? materialization.attachment
: admittedAttachmentToPromptAttachment(admitted)
);
changed = true;
}
prepared.push(changed ? { ...message, attachments } : message);
}
return prepared;
}
protected async waitForStructuredRetry(
delayMs: number,
signal?: AbortSignal
) {
await delay(delayMs, undefined, signal ? { signal } : undefined);
}
override getDriverSpec(): ProviderDriverSpec {
return {
createBackendConfig: execution => this.createNativeConfig(execution),
mapError: error => this.handleError(error),
chat: {
prepareMessages: async context =>
await this.prepareMessages(
context.input.messages,
context.backendConfig,
context.options
),
resolveRequestOptions: async context => {
const requestIntent = await llmResolveRequestIntentOptions({
protocol: context.protocol,
backendConfig: context.backendConfig,
reasoning: {
enabled: context.options.reasoning,
supported:
hasProviderModelBehaviorFlag(
context.model,
'reasoning_medium'
) ||
hasProviderModelBehaviorFlag(context.model, 'reasoning_high'),
effort: hasProviderModelBehaviorFlag(
context.model,
'reasoning_high'
)
? 'high'
: 'medium',
includeReasoning:
hasProviderModelBehaviorFlag(
context.model,
'reasoning_medium'
) ||
hasProviderModelBehaviorFlag(context.model, 'reasoning_high'),
},
});
return {
attachmentCapability: this.getAttachCapability(
context.model,
context.outputType
),
include: requestIntent.include,
reasoning: requestIntent.reasoning,
};
},
},
structured: {
prepareMessages: (inputMessages, backendConfig, structuredOptions) =>
this.prepareMessages(inputMessages, backendConfig, structuredOptions),
shouldRetry: async ({ error, attempt, options: structuredOptions }) => {
const isParsingError =
isInvalidStructuredOutputError(error) || error instanceof ZodError;
const retryableError =
isParsingError || !(error instanceof UserFriendlyError);
const maxRetries = Math.max(structuredOptions.maxRetries ?? 3, 0);
if (!retryableError || attempt >= maxRetries) {
return false;
}
if (!isParsingError) {
await this.waitForStructuredRetry(
GEMINI_RETRY_INITIAL_DELAY_MS * 2 ** attempt,
structuredOptions.signal
);
}
return true;
},
},
embedding: {
defaultDimensions: DEFAULT_DIMENSIONS,
taskType: 'RETRIEVAL_DOCUMENT',
},
rerank: false,
image: {
prepareMessages: (inputMessages, backendConfig, imageOptions) =>
this.prepareMessages(inputMessages, backendConfig, imageOptions),
},
};
}
}
packages/backend/server/src/plugins/copilot/providers/gemini/generative.ts
-28
View File
@@ -1,28 +0,0 @@
import type { LlmBackendConfig } from '../../../../native';
import type { CopilotProviderExecution } from '../provider-runtime-contract';
import { CopilotProviderType } from '../types';
import { GeminiProvider } from './gemini';
export type GeminiGenerativeConfig = {
apiKey: string;
baseURL?: string;
};
export class GeminiGenerativeProvider extends GeminiProvider<GeminiGenerativeConfig> {
override readonly type = CopilotProviderType.Gemini;
override configured(execution?: CopilotProviderExecution): boolean {
return !!this.getConfig(execution).apiKey;
}
protected override async createNativeConfig(
execution?: CopilotProviderExecution
): Promise<LlmBackendConfig> {
const config = this.getConfig(execution);
return {
base_url: (
config.baseURL || 'https://generativelanguage.googleapis.com/v1beta'
).replace(/\/$/, ''),
auth_token: config.apiKey,
};
}
}
packages/backend/server/src/plugins/copilot/providers/gemini/index.ts
-2
View File
@@ -1,2 +0,0 @@
export * from './generative';
export * from './vertex';
packages/backend/server/src/plugins/copilot/providers/gemini/vertex.ts
-34
View File
@@ -1,34 +0,0 @@
import type { LlmBackendConfig } from '../../../../native';
import type { CopilotProviderExecution } from '../provider-runtime-contract';
import { CopilotProviderType } from '../types';
import {
getGoogleAuth,
getVertexGoogleBaseUrl,
type VertexProviderConfig,
} from '../utils';
import { GeminiProvider } from './gemini';
export type GeminiVertexConfig = VertexProviderConfig;
export class GeminiVertexProvider extends GeminiProvider<GeminiVertexConfig> {
override readonly type = CopilotProviderType.GeminiVertex;
override configured(execution?: CopilotProviderExecution): boolean {
const config = this.getConfig(execution);
return !!getVertexGoogleBaseUrl(config) && !!config.googleAuthOptions;
}
protected async resolveVertexAuth(execution?: CopilotProviderExecution) {
return await getGoogleAuth(this.getConfig(execution), 'google');
}
protected override async createNativeConfig(
execution?: CopilotProviderExecution
): Promise<LlmBackendConfig> {
const auth = await this.resolveVertexAuth(execution);
const { Authorization: authHeader } = auth.headers();
return {
base_url: auth.baseUrl || '',
auth_token: authHeader.replace(/^Bearer\s+/i, ''),
};
}
}
packages/backend/server/src/plugins/copilot/providers/index.ts
-14
View File
@@ -1,14 +0,0 @@
export {
AnthropicOfficialProvider,
AnthropicVertexProvider,
} from './anthropic';
export { CloudflareWorkersAIProvider } from './cloudflare';
export { CopilotProviderFactory } from './factory';
export { FalProvider } from './fal';
export { GeminiGenerativeProvider, GeminiVertexProvider } from './gemini';
export { CopilotProviderLifecycleService } from './lifecycle-service';
export { OpenAIProvider } from './openai';
export type { CopilotProvider } from './provider';
export { CopilotProviders } from './provider-tokens';
export { CopilotProviderRegistryService } from './registry-service';
export * from './types';
packages/backend/server/src/plugins/copilot/providers/lifecycle-service.ts
-90
View File
@@ -1,90 +0,0 @@
import { Injectable, Type } from '@nestjs/common';
import { ModuleRef } from '@nestjs/core';
import { OnEvent } from '../../../base';
import { CopilotProviderFactory } from './factory';
import type { CopilotProvider } from './provider';
import type { CopilotProviderExecution } from './provider-runtime-contract';
import { CopilotProviders } from './provider-tokens';
import { CopilotProviderRegistryService } from './registry-service';
@Injectable()
export class CopilotProviderLifecycleService {
private readonly registeredByProvider = new WeakMap<
CopilotProvider,
Set<string>
>();
constructor(
private readonly moduleRef: ModuleRef,
private readonly factory: CopilotProviderFactory,
private readonly registries: CopilotProviderRegistryService
) {}
private getProviders(): CopilotProvider[] {
return CopilotProviders.flatMap(token => {
const provider = this.moduleRef.get(token as Type<CopilotProvider>, {
strict: false,
});
return provider ? [provider] : [];
});
}
private getRegisteredProviderIds(provider: CopilotProvider) {
const current = this.registeredByProvider.get(provider);
if (current) {
return current;
}
const next = new Set<string>();
this.registeredByProvider.set(provider, next);
return next;
}
private async syncProvider(provider: CopilotProvider) {
const registry = this.registries.getRegistry();
const configuredIds = new Set<string>();
for (const providerId of registry.byType.get(provider.type) ?? []) {
const profile = registry.profiles.get(providerId);
if (!profile) {
continue;
}
const execution: CopilotProviderExecution = { providerId, profile };
if (!provider.configured(execution)) {
this.factory.unregister(providerId, provider);
continue;
}
configuredIds.add(providerId);
this.factory.register(providerId, provider);
}
const previous = this.getRegisteredProviderIds(provider);
for (const providerId of previous) {
if (!configuredIds.has(providerId)) {
this.factory.unregister(providerId, provider);
}
}
this.registeredByProvider.set(provider, configuredIds);
}
async syncProviders() {
for (const provider of this.getProviders()) {
await this.syncProvider(provider);
}
}
@OnEvent('config.init')
async onConfigInit() {
await this.syncProviders();
}
@OnEvent('config.changed')
async onConfigChanged(event: Events['config.changed']) {
if ('copilot' in event.updates) {
await this.syncProviders();
}
}
}
packages/backend/server/src/plugins/copilot/providers/openai.ts
-172
View File
@@ -1,172 +0,0 @@
import { Inject } from '@nestjs/common';
import {
CopilotProviderSideError,
OneMB,
UserFriendlyError,
} from '../../../base';
import {
type LlmBackendConfig,
llmResolveRequestIntentOptions,
} from '../../../native';
import {
admittedAttachmentToPromptAttachment,
AttachmentAdmissionHost,
} from '../runtime/hosts/attachment-admission';
import { AttachmentMaterializer } from '../runtime/hosts/attachment-materializer';
import { CopilotProvider } from './provider';
import { hasProviderModelBehaviorFlag } from './provider-model-runtime';
import type {
CopilotProviderExecution,
ProviderDriverSpec,
} from './provider-runtime-contract';
import {
CopilotProviderType,
type PromptAttachment,
type PromptMessage,
} from './types';
import { promptAttachmentToUrl } from './utils';
export const DEFAULT_DIMENSIONS = 256;
export type OpenAIConfig = {
apiKey: string;
baseURL?: string;
oldApiStyle?: boolean;
};
export class OpenAIProvider extends CopilotProvider<OpenAIConfig> {
readonly type = CopilotProviderType.OpenAI;
@Inject() protected readonly attachmentMaterializer!: AttachmentMaterializer;
@Inject()
protected readonly attachmentAdmissionHost?: AttachmentAdmissionHost;
protected resolveModelBackendKind(execution?: CopilotProviderExecution) {
return this.getConfig(execution).oldApiStyle
? ('openai_chat' as const)
: ('openai_responses' 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 openai response',
});
}
protected createNativeConfig(
execution?: CopilotProviderExecution
): LlmBackendConfig {
const config = this.getConfig(execution);
const baseUrl = config.baseURL || 'https://api.openai.com/v1';
return {
base_url: baseUrl.replace(/\/v1\/?$/, ''),
auth_token: config.apiKey,
};
}
private getAttachmentAdmissionHost() {
return (
this.attachmentAdmissionHost ??
new AttachmentAdmissionHost(this.attachmentMaterializer)
);
}
private async prepareImageMessages(
messages: PromptMessage[],
options: {
signal?: AbortSignal;
user?: string;
workspace?: string;
session?: string;
}
) {
const prepared: PromptMessage[] = [];
for (const message of messages) {
options.signal?.throwIfAborted();
if (!Array.isArray(message.attachments) || !message.attachments.length) {
prepared.push(message);
continue;
}
let changed = false;
const attachments: PromptAttachment[] = [];
for (const attachment of message.attachments) {
options.signal?.throwIfAborted();
const url = promptAttachmentToUrl(attachment);
if (!url || url.startsWith('data:')) {
attachments.push(attachment);
continue;
}
const admitted =
await this.getAttachmentAdmissionHost().admitPromptAttachment(
attachment,
{
userId: options.user ?? 'provider-runtime',
workspaceId: options.workspace ?? 'provider-runtime',
sessionId: options.session,
signal: options.signal,
maxBytes: 50 * OneMB,
}
);
attachments.push(admittedAttachmentToPromptAttachment(admitted));
changed = true;
}
prepared.push(changed ? { ...message, attachments } : message);
}
return prepared;
}
override getDriverSpec(): ProviderDriverSpec {
return {
createBackendConfig: execution => this.createNativeConfig(execution),
mapError: error => this.handleError(error),
chat: {
resolveRequestOptions: async context => {
const requestIntent = await llmResolveRequestIntentOptions({
protocol: context.protocol,
backendConfig: context.backendConfig,
include: context.options.webSearch ? ['citations'] : undefined,
reasoning: {
enabled: context.options.reasoning,
supported: hasProviderModelBehaviorFlag(
context.model,
'reasoning_supported'
),
},
});
return {
attachmentCapability: this.getAttachCapability(
context.model,
context.outputType
),
include: requestIntent.include,
reasoning: requestIntent.reasoning,
};
},
},
structured: {},
embedding: {
defaultDimensions: DEFAULT_DIMENSIONS,
taskType: 'RETRIEVAL_DOCUMENT',
},
image: {
prepareMessages: async (messages, _backendConfig, options) =>
await this.prepareImageMessages(messages, options),
},
};
}
}
packages/backend/server/src/plugins/copilot/providers/provider-middleware.ts
-73
View File
@@ -1,73 +0,0 @@
import type { ProviderMiddlewareConfig } from '../config';
import { CopilotProviderType } from './types';
const DEFAULT_NODE_TEXT_MIDDLEWARE: NonNullable<
NonNullable<ProviderMiddlewareConfig['node']>['text']
> = ['citation_footnote', 'callout'];
const DEFAULT_MIDDLEWARE_BY_TYPE: Record<
CopilotProviderType,
ProviderMiddlewareConfig
> = {
[CopilotProviderType.OpenAI]: {
node: { text: DEFAULT_NODE_TEXT_MIDDLEWARE },
},
[CopilotProviderType.CloudflareWorkersAi]: {
node: { text: DEFAULT_NODE_TEXT_MIDDLEWARE },
},
[CopilotProviderType.Anthropic]: {
node: { text: DEFAULT_NODE_TEXT_MIDDLEWARE },
},
[CopilotProviderType.AnthropicVertex]: {
node: { text: DEFAULT_NODE_TEXT_MIDDLEWARE },
},
[CopilotProviderType.Gemini]: {
node: { text: DEFAULT_NODE_TEXT_MIDDLEWARE },
},
[CopilotProviderType.GeminiVertex]: {
node: { text: DEFAULT_NODE_TEXT_MIDDLEWARE },
},
[CopilotProviderType.FAL]: {},
};
function unique<T>(items: T[]) {
return [...new Set(items)];
}
function mergeArray<T>(base: T[] | undefined, override: T[] | undefined) {
if (!base?.length && !override?.length) {
return undefined;
}
return unique([...(base ?? []), ...(override ?? [])]);
}
function compactMiddlewareSection<T extends Record<string, unknown>>(
section: T
): T | undefined {
return Object.values(section).some(value => value !== undefined)
? section
: undefined;
}
export function mergeProviderMiddleware(
defaults: ProviderMiddlewareConfig,
override?: ProviderMiddlewareConfig
): ProviderMiddlewareConfig {
return {
rust: compactMiddlewareSection({
request: mergeArray(defaults.rust?.request, override?.rust?.request),
stream: mergeArray(defaults.rust?.stream, override?.rust?.stream),
}),
node: compactMiddlewareSection({
text: mergeArray(defaults.node?.text, override?.node?.text),
}),
};
}
export function resolveProviderMiddleware(
type: CopilotProviderType,
override?: ProviderMiddlewareConfig
): ProviderMiddlewareConfig {
const defaults = DEFAULT_MIDDLEWARE_BY_TYPE[type] ?? {};
return mergeProviderMiddleware(defaults, override);
}
packages/backend/server/src/plugins/copilot/providers/provider-model-runtime.ts
-385
View File
@@ -1,385 +0,0 @@
import { z } from 'zod';
import { CopilotPromptInvalid } from '../../../base';
import {
type LlmBackendConfig,
llmInferPromptModelConditions,
llmMatchModelCapabilities,
llmMatchModelRegistry,
type LlmProtocol,
llmResolveModelRegistryVariant,
} from '../../../native';
import { applyPromptAttachmentMimeTypeHintForNative } from './attachments';
import {
type CopilotChatOptions,
type CopilotImageOptions,
type CopilotModelBackendKind,
type CopilotProviderModel,
type CopilotProviderType,
type CopilotStructuredOptions,
EmbeddingMessage,
type ModelAttachmentCapability,
type ModelCapability,
type ModelFullConditions,
ModelInputType,
ModelOutputType,
type PromptAttachmentKind,
type PromptAttachmentSourceKind,
type PromptMessage,
PromptMessageSchema,
} from './types';
// Owner: backend host model-selection glue.
// Capability matching and catalog lookup are delegated to native/adapter; this
// file keeps provider prefix/default/prefer behavior and Node prompt checks.
export type ProviderModelRuntimeContext = {
type: CopilotProviderType;
backendKind: CopilotModelBackendKind;
};
export type ResolvedProviderModel = CopilotProviderModel & {
backendKind: CopilotModelBackendKind;
canonicalKey: string;
protocol?: LlmProtocol;
requestLayer?: LlmBackendConfig['request_layer'];
routeOverrides?: Partial<
Record<
ModelOutputType,
{
protocol?: LlmProtocol;
requestLayer?: LlmBackendConfig['request_layer'];
}
>
>;
behaviorFlags?: string[];
};
function unique<T>(values: Iterable<T>) {
return Array.from(new Set(values));
}
function resolveAttachmentCapability(
cap: ModelCapability,
outputType?: ModelOutputType
): ModelAttachmentCapability | undefined {
if (outputType === ModelOutputType.Structured) {
return cap.structuredAttachments ?? cap.attachments;
}
return cap.attachments;
}
function toProviderModel(
variant: NonNullable<
ReturnType<typeof llmResolveModelRegistryVariant>['variant']
>
): ResolvedProviderModel {
return {
id: variant.rawModelId,
name: variant.displayName,
backendKind: variant.backendKind,
canonicalKey: variant.canonicalKey,
protocol: variant.protocol,
requestLayer: variant.requestLayer,
routeOverrides: variant.routeOverrides,
behaviorFlags: variant.behaviorFlags,
capabilities: variant.capabilities.map(capability => ({
input: capability.input as ModelInputType[],
output: capability.output as ModelOutputType[],
attachments: capability.attachments
? {
kinds: capability.attachments.kinds as PromptAttachmentKind[],
sourceKinds: capability.attachments.sourceKinds as
| ModelAttachmentCapability['sourceKinds']
| undefined,
allowRemoteUrls: capability.attachments.allowRemoteUrls,
}
: undefined,
structuredAttachments: capability.structuredAttachments
? {
kinds: capability.structuredAttachments
.kinds as PromptAttachmentKind[],
sourceKinds: capability.structuredAttachments.sourceKinds as
| ModelAttachmentCapability['sourceKinds']
| undefined,
allowRemoteUrls: capability.structuredAttachments.allowRemoteUrls,
}
: undefined,
defaultForOutputType: capability.defaultForOutputType,
})),
};
}
export type ProviderModelSelection = {
kind: 'configured';
model: ResolvedProviderModel;
};
export function resolveProviderModelSelection(
context: ProviderModelRuntimeContext,
cond: ModelFullConditions
): ProviderModelSelection | undefined {
if (cond.modelId) {
const resolved = llmResolveModelRegistryVariant({
backendKind: context.backendKind,
modelId: cond.modelId,
}).variant;
if (!resolved) {
return;
}
const model = toProviderModel(resolved);
const matchedModelId = llmMatchModelCapabilities([model], {
...cond,
modelId: model.id,
});
if (!matchedModelId) {
return;
}
return {
kind: 'configured',
model,
};
}
const resolved = llmMatchModelRegistry({
backendKind: context.backendKind,
cond,
}).variant;
if (!resolved) {
return;
}
return {
kind: 'configured',
model: toProviderModel(resolved),
};
}
function isMultimodal(model: CopilotProviderModel) {
return model.capabilities.some(c =>
[ModelInputType.Image, ModelInputType.Audio, ModelInputType.File].some(t =>
c.input.includes(t)
)
);
}
function handleZodError(ret: z.SafeParseReturnType<any, any>) {
if (ret.success) return;
const issues = ret.error.issues.map(i => {
const path =
'root' +
(i.path.length
? `.${i.path.map(seg => (typeof seg === 'number' ? `[${seg}]` : `.${seg}`)).join('')}`
: '');
return `${i.message}${path}`;
});
throw new CopilotPromptInvalid(issues.join('; '));
}
export async function inferModelConditionsFromMessages(
messages?: PromptMessage[],
withAttachment = true
): Promise<Partial<ModelFullConditions>> {
if (!messages?.length || !withAttachment) return {};
const projectedMessages = messages.map(message => ({
role: message.role,
content: message.content,
...(Array.isArray(message.attachments) && message.attachments.length
? {
attachments: message.attachments.map(attachment =>
applyPromptAttachmentMimeTypeHintForNative(attachment, message)
),
}
: {}),
}));
const inferredCond = llmInferPromptModelConditions(projectedMessages);
return {
...(inferredCond.attachmentKinds?.length
? { attachmentKinds: unique(inferredCond.attachmentKinds) }
: {}),
...(inferredCond.attachmentSourceKinds?.length
? {
attachmentSourceKinds: unique(
inferredCond.attachmentSourceKinds
) as PromptAttachmentSourceKind[],
}
: {}),
...(inferredCond.inputTypes?.length
? { inputTypes: unique(inferredCond.inputTypes) as ModelInputType[] }
: {}),
...(inferredCond.hasRemoteAttachments
? { hasRemoteAttachments: true }
: {}),
};
}
export function mergeModelConditions(
cond: ModelFullConditions,
inferredCond: Partial<ModelFullConditions>
): ModelFullConditions {
return {
...inferredCond,
...cond,
inputTypes: unique([
...(inferredCond.inputTypes ?? []),
...(cond.inputTypes ?? []),
]),
attachmentKinds: unique([
...(inferredCond.attachmentKinds ?? []),
...(cond.attachmentKinds ?? []),
]),
attachmentSourceKinds: unique([
...(inferredCond.attachmentSourceKinds ?? []),
...(cond.attachmentSourceKinds ?? []),
]),
hasRemoteAttachments:
cond.hasRemoteAttachments ?? inferredCond.hasRemoteAttachments,
};
}
export function getAttachCapability(
model: CopilotProviderModel,
outputType: ModelOutputType
): ModelAttachmentCapability | undefined {
const capability =
model.capabilities.find(cap => cap.output.includes(outputType)) ??
model.capabilities[0];
if (!capability) {
return;
}
return resolveAttachmentCapability(capability, outputType);
}
export function matchProviderModel(
context: ProviderModelRuntimeContext,
cond: ModelFullConditions
): boolean {
return !!resolveProviderModelSelection(context, cond);
}
export function resolveProviderModel(
context: ProviderModelRuntimeContext,
modelId: string
): ResolvedProviderModel | undefined {
return resolveProviderModelSelection(context, {
modelId,
})?.model;
}
export function hasProviderModelBehaviorFlag(
model: CopilotProviderModel,
flag: string
) {
const behaviorFlags = (model as ResolvedProviderModel).behaviorFlags;
return Array.isArray(behaviorFlags) && behaviorFlags.includes(flag);
}
export function resolveProviderModelRoute(
model: CopilotProviderModel,
outputType: ModelOutputType
) {
const resolved = model as ResolvedProviderModel;
const override = resolved.routeOverrides?.[outputType];
return {
protocol: override?.protocol ?? resolved.protocol,
requestLayer: override?.requestLayer ?? resolved.requestLayer,
};
}
export function requireProviderModelSelection(
context: ProviderModelRuntimeContext,
cond: ModelFullConditions
): ResolvedProviderModel {
const selection = resolveProviderModelSelection(context, cond);
if (selection) return selection.model;
const { modelId, outputType, inputTypes } = cond;
throw new CopilotPromptInvalid(
modelId
? `Model ${modelId} does not support ${outputType ?? '<any>'} output with ${inputTypes ?? '<any>'} input`
: outputType
? `No model supports ${outputType} output with ${inputTypes ?? '<any>'} input for provider ${context.type}`
: 'Output type is required when modelId is not provided'
);
}
export async function checkProviderParams(
context: ProviderModelRuntimeContext,
{
cond,
messages,
embeddings,
options = {},
withAttachment = true,
}: {
cond: ModelFullConditions;
messages?: PromptMessage[];
embeddings?: string[];
options?:
| CopilotChatOptions
| CopilotStructuredOptions
| CopilotImageOptions;
withAttachment?: boolean;
execution?: unknown;
}
): Promise<ModelFullConditions> {
if (messages) {
const { requireContent = true, requireAttachment = false } = options;
const MessageSchema = z
.array(
PromptMessageSchema.extend({
content: requireContent
? z.string().trim().min(1)
: z.string().optional().nullable(),
})
.passthrough()
.catchall(z.union([z.string(), z.number(), z.date(), z.null()]))
)
.optional();
handleZodError(MessageSchema.safeParse(messages));
const inferredCond = await inferModelConditionsFromMessages(
messages,
withAttachment
);
const mergedCond = mergeModelConditions(cond, inferredCond);
const model = requireProviderModelSelection(context, mergedCond);
const multimodal = isMultimodal(model);
if (
multimodal &&
requireAttachment &&
!messages.some(
message =>
message.role === 'user' &&
Array.isArray(message.attachments) &&
message.attachments.length > 0
)
) {
throw new CopilotPromptInvalid('attachments required in multimodal mode');
}
if (embeddings) {
handleZodError(EmbeddingMessage.safeParse(embeddings));
}
return mergedCond;
}
const inferredCond = await inferModelConditionsFromMessages(
messages,
withAttachment
);
const mergedCond = mergeModelConditions(cond, inferredCond);
if (embeddings) {
handleZodError(EmbeddingMessage.safeParse(embeddings));
}
return mergedCond;
}
packages/backend/server/src/plugins/copilot/providers/provider-native-runtime.ts
-337
View File
@@ -1,337 +0,0 @@
import type {
LlmBackendConfig,
LlmEmbeddingRequest,
LlmProtocol,
LlmRerankRequest,
LlmStructuredRequest,
} from '../../../native';
import {
buildLlmImageRequestFromMessages,
llmEmbeddingDispatch,
llmRerankDispatch,
llmStructuredDispatch,
} from '../../../native';
import type { NodeTextMiddleware, ProviderMiddlewareConfig } from '../config';
import {
buildToolContracts,
projectPromptMessageForNative,
} from '../runtime/contracts';
import { buildNativeRequest } from '../runtime/native-request-runtime';
import type { ToolLoopBackend } from '../runtime/tool/bridge';
import type { NativeProviderAdapter } from '../runtime/tool/native-adapter';
import type { CopilotToolSet } from '../tools';
import type {
CopilotProviderExecution,
PreparedNativeEmbeddingExecution,
PreparedNativeExecution,
PreparedNativeImageExecution,
PreparedNativeRequestOptions,
PreparedNativeRerankExecution,
PreparedNativeStructuredExecution,
} from './provider-runtime-contract';
import type {
CopilotChatOptions,
CopilotImageOptions,
PromptMessage,
} from './types';
export type CreateToolAdapterOptions = {
maxSteps?: number;
nodeTextMiddleware?: NodeTextMiddleware[];
};
export type CreateNativeAdapter = (
backend: ToolLoopBackend,
tools: CopilotToolSet,
nodeTextMiddleware?: NodeTextMiddleware[],
options?: CreateToolAdapterOptions
) => NativeProviderAdapter;
export type CreatePreparedExecutionRuntimeInput = {
resolveProviderId: (execution?: CopilotProviderExecution) => string;
getTools: (
options: CopilotChatOptions,
model: string
) => Promise<CopilotToolSet>;
getActiveProviderMiddleware: (
execution?: CopilotProviderExecution
) => ProviderMiddlewareConfig;
createNativeAdapter: CreateNativeAdapter;
maxSteps: number;
};
export type PreparedExecutionRuntime = ReturnType<
typeof createPreparedExecutionRuntime
>;
export function createPreparedExecutionRuntime(
input: CreatePreparedExecutionRuntimeInput
) {
return {
buildPreparedNativeExecution: async (
prepared: PreparedNativeRequestOptions
) =>
await buildPreparedNativeExecution(
input.resolveProviderId(prepared.execution),
input.getTools,
input.getActiveProviderMiddleware,
input.maxSteps,
prepared
),
createPreparedExecutionAdapter: (prepared: PreparedNativeExecution) =>
createPreparedExecutionAdapter(
input.createNativeAdapter,
input.maxSteps,
prepared
),
buildPreparedNativeStructuredExecution: (
protocol: LlmProtocol,
backendConfig: LlmBackendConfig,
model: string,
request: LlmStructuredRequest,
execution?: CopilotProviderExecution
) =>
buildPreparedNativeStructuredExecution(
input.resolveProviderId(execution),
protocol,
backendConfig,
model,
request
),
buildPreparedNativeEmbeddingExecution: (
protocol: LlmProtocol,
backendConfig: LlmBackendConfig,
model: string,
request: LlmEmbeddingRequest,
execution?: CopilotProviderExecution
) =>
buildPreparedNativeEmbeddingExecution(
input.resolveProviderId(execution),
protocol,
backendConfig,
model,
request
),
buildPreparedNativeRerankExecution: (
protocol: LlmProtocol,
backendConfig: LlmBackendConfig,
model: string,
request: LlmRerankRequest,
execution?: CopilotProviderExecution
) =>
buildPreparedNativeRerankExecution(
input.resolveProviderId(execution),
protocol,
backendConfig,
model,
request
),
buildPreparedNativeImageExecution: (
protocol: LlmProtocol,
backendConfig: LlmBackendConfig,
model: string,
messages: PromptMessage[],
options: CopilotImageOptions = {},
execution?: CopilotProviderExecution
) =>
buildPreparedNativeImageExecution(
input.resolveProviderId(execution),
protocol,
backendConfig,
model,
messages,
options
),
};
}
export function createPreparedExecutionAdapter(
createNativeAdapter: CreateNativeAdapter,
maxSteps: number,
prepared: PreparedNativeExecution
) {
return createNativeAdapter(
{
protocol: prepared.route.protocol,
backendConfig: prepared.route.backendConfig,
},
prepared.tools,
prepared.postprocess?.nodeTextMiddleware,
{
maxSteps,
nodeTextMiddleware: prepared.postprocess?.nodeTextMiddleware,
}
);
}
export function createNativeStructuredDispatch(
backendConfig: LlmBackendConfig,
protocol: LlmProtocol
) {
return (request: LlmStructuredRequest) =>
llmStructuredDispatch(protocol, backendConfig, request);
}
export function createNativeEmbeddingDispatch(
backendConfig: LlmBackendConfig,
protocol: LlmProtocol
) {
return (request: LlmEmbeddingRequest) =>
llmEmbeddingDispatch(protocol, backendConfig, request);
}
export function createNativeRerankDispatch(
backendConfig: LlmBackendConfig,
protocol: LlmProtocol
) {
return (request: LlmRerankRequest) =>
llmRerankDispatch(protocol, backendConfig, request);
}
function buildPreparedRoute(
providerId: string,
protocol: LlmProtocol,
backendConfig: LlmBackendConfig,
model: string
): PreparedNativeExecution['route'] {
return {
providerId,
protocol,
requestLayer: backendConfig.request_layer,
model,
backendConfig,
};
}
export async function buildPreparedNativeExecution(
providerId: string,
getTools: (
options: CopilotChatOptions,
model: string
) => Promise<CopilotToolSet>,
getActiveProviderMiddleware: (
execution?: CopilotProviderExecution
) => ProviderMiddlewareConfig,
maxSteps: number,
{
protocol,
backendConfig,
model,
messages,
options = {},
execution,
withAttachment = true,
attachmentCapability,
include,
reasoning,
tools,
middleware,
}: PreparedNativeRequestOptions
): Promise<PreparedNativeExecution> {
const resolvedTools = tools ?? (await getTools(options, model));
const resolvedMiddleware =
middleware ?? getActiveProviderMiddleware(execution);
const { request } = await buildNativeRequest({
model,
messages,
options,
toolContracts: buildToolContracts(resolvedTools),
withAttachment,
attachmentCapability,
include,
reasoning,
middleware: resolvedMiddleware,
});
return {
route: buildPreparedRoute(providerId, protocol, backendConfig, model),
request,
tools: resolvedTools,
maxSteps,
postprocess: {
nodeTextMiddleware: resolvedMiddleware.node?.text,
},
};
}
type BuildPreparedNativeDispatchExecution = <
TRequest extends
| LlmStructuredRequest
| LlmEmbeddingRequest
| LlmRerankRequest,
>(
providerId: string,
protocol: LlmProtocol,
backendConfig: LlmBackendConfig,
model: string,
request: TRequest
) => {
route: PreparedNativeExecution['route'];
request: TRequest;
};
const buildPreparedNativeDispatchExecution: BuildPreparedNativeDispatchExecution =
(providerId, protocol, backendConfig, model, request) => {
return {
route: buildPreparedRoute(providerId, protocol, backendConfig, model),
request,
};
};
export const buildPreparedNativeStructuredExecution =
buildPreparedNativeDispatchExecution as (
providerId: string,
protocol: LlmProtocol,
backendConfig: LlmBackendConfig,
model: string,
request: LlmStructuredRequest
) => PreparedNativeStructuredExecution;
export const buildPreparedNativeEmbeddingExecution =
buildPreparedNativeDispatchExecution as (
providerId: string,
protocol: LlmProtocol,
backendConfig: LlmBackendConfig,
model: string,
request: LlmEmbeddingRequest
) => PreparedNativeEmbeddingExecution;
export const buildPreparedNativeRerankExecution =
buildPreparedNativeDispatchExecution as (
providerId: string,
protocol: LlmProtocol,
backendConfig: LlmBackendConfig,
model: string,
request: LlmRerankRequest
) => PreparedNativeRerankExecution;
export function buildPreparedNativeImageExecution(
providerId: string,
protocol: LlmProtocol,
backendConfig: LlmBackendConfig,
model: string,
messages: PromptMessage[],
options: CopilotImageOptions = {}
): PreparedNativeImageExecution {
const nativeMessages = messages.map(
message => projectPromptMessageForNative(message).message
);
return {
route: buildPreparedRoute(providerId, protocol, backendConfig, model),
request: buildLlmImageRequestFromMessages({
model,
protocol,
messages: nativeMessages,
options: projectImageRequestOptions(options),
}),
};
}
function projectImageRequestOptions(options: CopilotImageOptions = {}) {
return {
quality: options.quality,
seed: options.seed,
modelName: options.modelName,
loras: options.loras,
};
}
packages/backend/server/src/plugins/copilot/providers/provider-registry.ts
-287
View File
@@ -1,287 +0,0 @@
import type {
CopilotProviderConfigMap,
CopilotProviderDefaults,
CopilotProviderProfile,
ProviderMiddlewareConfig,
} from '../config';
import { resolveProviderMiddleware } from './provider-middleware';
import { CopilotProviderType, ModelOutputType } from './types';
const PROVIDER_ID_PATTERN = /^[a-zA-Z0-9-_]+$/;
const LEGACY_PROVIDER_ORDER: CopilotProviderType[] = [
CopilotProviderType.OpenAI,
CopilotProviderType.CloudflareWorkersAi,
CopilotProviderType.FAL,
CopilotProviderType.Gemini,
CopilotProviderType.GeminiVertex,
CopilotProviderType.Anthropic,
CopilotProviderType.AnthropicVertex,
];
const LEGACY_PROVIDER_PRIORITY = LEGACY_PROVIDER_ORDER.reduce(
(acc, type, index) => {
acc[type] = LEGACY_PROVIDER_ORDER.length - index;
return acc;
},
{} as Record<CopilotProviderType, number>
);
type LegacyProvidersConfig = Partial<
Record<CopilotProviderType, CopilotProviderConfigMap[CopilotProviderType]>
>;
export type CopilotProvidersConfigInput = LegacyProvidersConfig & {
profiles?: CopilotProviderProfile[] | null;
defaults?: CopilotProviderDefaults | null;
};
export type NormalizedCopilotProviderProfile = Omit<
CopilotProviderProfile,
'enabled' | 'priority' | 'middleware'
> & {
enabled: boolean;
priority: number;
middleware: ProviderMiddlewareConfig;
};
export type CopilotProviderRegistry = {
profiles: Map<string, NormalizedCopilotProviderProfile>;
defaults: CopilotProviderDefaults;
order: string[];
byType: Map<CopilotProviderType, string[]>;
};
export type ResolveModelResult = {
rawModelId?: string;
modelId?: string;
explicitProviderId?: string;
candidateProviderIds: string[];
};
type ResolveModelOptions = {
registry: CopilotProviderRegistry;
modelId?: string;
outputType?: ModelOutputType;
availableProviderIds?: Iterable<string>;
preferredProviderIds?: Iterable<string>;
};
function unique<T>(list: T[]): T[] {
return [...new Set(list)];
}
function asArray<T>(iter?: Iterable<T>): T[] {
return iter ? Array.from(iter) : [];
}
function parseModelPrefix(
registry: CopilotProviderRegistry,
modelId: string
): { providerId: string; modelId?: string } | null {
const index = modelId.indexOf('/');
if (index <= 0) {
return null;
}
const providerId = modelId.slice(0, index);
if (!registry.profiles.has(providerId)) {
return null;
}
const model = modelId.slice(index + 1);
return { providerId, modelId: model || undefined };
}
function normalizeProfile(
profile: CopilotProviderProfile
): NormalizedCopilotProviderProfile {
return {
...profile,
enabled: profile.enabled !== false,
priority: profile.priority ?? 0,
middleware: resolveProviderMiddleware(profile.type, profile.middleware),
};
}
function toLegacyProfiles(
config: CopilotProvidersConfigInput
): CopilotProviderProfile[] {
const legacyProfiles: CopilotProviderProfile[] = [];
for (const type of LEGACY_PROVIDER_ORDER) {
const legacyConfig = config[type];
if (!legacyConfig) {
continue;
}
legacyProfiles.push({
id: `${type}-default`,
type,
priority: LEGACY_PROVIDER_PRIORITY[type],
config: legacyConfig,
} as CopilotProviderProfile);
}
return legacyProfiles;
}
function mergeProfiles(
explicitProfiles: CopilotProviderProfile[],
legacyProfiles: CopilotProviderProfile[]
): CopilotProviderProfile[] {
const profiles = new Map<string, CopilotProviderProfile>();
for (const profile of explicitProfiles) {
if (!PROVIDER_ID_PATTERN.test(profile.id)) {
throw new Error(`Invalid copilot provider profile id: ${profile.id}`);
}
if (profiles.has(profile.id)) {
throw new Error(`Duplicated copilot provider profile id: ${profile.id}`);
}
profiles.set(profile.id, profile);
}
for (const profile of legacyProfiles) {
if (!profiles.has(profile.id)) {
profiles.set(profile.id, profile);
}
}
return Array.from(profiles.values());
}
function sortProfiles(profiles: NormalizedCopilotProviderProfile[]) {
return profiles.toSorted((a, b) => {
if (a.priority !== b.priority) {
return b.priority - a.priority;
}
return a.id.localeCompare(b.id);
});
}
function assertDefaults(
defaults: CopilotProviderDefaults,
profiles: Map<string, NormalizedCopilotProviderProfile>
) {
for (const providerId of Object.values(defaults)) {
if (!providerId) {
continue;
}
if (!profiles.has(providerId)) {
throw new Error(
`Copilot provider defaults references unknown providerId: ${providerId}`
);
}
}
}
export function buildProviderRegistry(
config: CopilotProvidersConfigInput
): CopilotProviderRegistry {
const explicitProfiles = config.profiles ?? [];
const legacyProfiles = toLegacyProfiles(config);
const mergedProfiles = mergeProfiles(explicitProfiles, legacyProfiles)
.map(normalizeProfile)
.filter(profile => profile.enabled);
const sortedProfiles = sortProfiles(mergedProfiles);
const profiles = new Map(
sortedProfiles.map(profile => [profile.id, profile] as const)
);
const defaults = config.defaults ?? {};
assertDefaults(defaults, profiles);
const order = sortedProfiles.map(profile => profile.id);
const byType = new Map<CopilotProviderType, string[]>();
for (const profile of sortedProfiles) {
const ids = byType.get(profile.type) ?? [];
ids.push(profile.id);
byType.set(profile.type, ids);
}
return { profiles, defaults, order, byType };
}
export function resolveModel({
registry,
modelId,
outputType,
availableProviderIds,
preferredProviderIds,
}: ResolveModelOptions): ResolveModelResult {
const available = new Set(asArray(availableProviderIds));
const preferred = new Set(asArray(preferredProviderIds));
const hasAvailableFilter = available.size > 0;
const hasPreferredFilter = preferred.size > 0;
const isAllowed = (providerId: string) => {
const profile = registry.profiles.get(providerId);
if (!profile?.enabled) {
return false;
}
if (hasAvailableFilter && !available.has(providerId)) {
return false;
}
if (hasPreferredFilter && !preferred.has(providerId)) {
return false;
}
return true;
};
const prefixed = modelId ? parseModelPrefix(registry, modelId) : null;
if (prefixed) {
return {
rawModelId: modelId,
modelId: prefixed.modelId,
explicitProviderId: prefixed.providerId,
candidateProviderIds: isAllowed(prefixed.providerId)
? [prefixed.providerId]
: [],
};
}
if (modelId) {
return {
rawModelId: modelId,
modelId,
candidateProviderIds: registry.order.filter(providerId =>
isAllowed(providerId)
),
};
}
const defaultProviderId =
outputType && outputType !== ModelOutputType.Rerank
? registry.defaults[outputType]
: undefined;
const fallbackOrder = [
...(defaultProviderId ? [defaultProviderId] : []),
registry.defaults.fallback,
...registry.order,
].filter((id): id is string => !!id);
return {
rawModelId: modelId,
modelId,
candidateProviderIds: unique(
fallbackOrder.filter(providerId => isAllowed(providerId))
),
};
}
export function stripProviderPrefix(
registry: CopilotProviderRegistry,
providerId: string,
modelId?: string
) {
if (!modelId) {
return modelId;
}
const prefixed = parseModelPrefix(registry, modelId);
if (!prefixed) {
return modelId;
}
if (prefixed.providerId !== providerId) {
return modelId;
}
return prefixed.modelId;
}
packages/backend/server/src/plugins/copilot/providers/provider-runtime-contract.ts
-456
View File
@@ -1,456 +0,0 @@
import type {
LlmBackendConfig,
LlmEmbeddingRequest,
LlmImageRequest,
LlmProtocol,
LlmRequest,
LlmRerankRequest,
LlmStructuredRequest,
} from '../../../native';
import type { NodeTextMiddleware, ProviderMiddlewareConfig } from '../config';
import type { CopilotToolSet } from '../tools';
import {
type ProviderModelRuntimeContext,
resolveProviderModelRoute,
} from './provider-model-runtime';
import type { NormalizedCopilotProviderProfile } from './provider-registry';
import {
CopilotChatOptions,
CopilotImageOptions,
CopilotProviderModel,
CopilotStructuredOptions,
ModelAttachmentCapability,
ModelConditions,
ModelFullConditions,
ModelOutputType,
PromptMessage,
} from './types';
export type NativeExecutionRoute = {
protocol: LlmProtocol;
requestLayer?: LlmBackendConfig['request_layer'];
model: string;
backendConfig: LlmBackendConfig;
};
export type CopilotProviderExecution = {
providerId: string;
profile: NormalizedCopilotProviderProfile;
};
export type PreparedNativeExecution = {
route: NativeExecutionRoute & {
providerId: string;
};
request: LlmRequest;
tools: CopilotToolSet;
maxSteps?: number;
postprocess?: {
nodeTextMiddleware?: NodeTextMiddleware[];
};
};
export type PreparedNativeStructuredExecution = {
route: NativeExecutionRoute & {
providerId: string;
};
request: LlmStructuredRequest;
};
export type PreparedNativeEmbeddingExecution = {
route: NativeExecutionRoute & {
providerId: string;
};
request: LlmEmbeddingRequest;
};
export type PreparedNativeRerankExecution = {
route: NativeExecutionRoute & {
providerId: string;
};
request: LlmRerankRequest;
};
export type PreparedNativeImageExecution = {
route: NativeExecutionRoute & {
providerId: string;
};
request: LlmImageRequest;
};
export type PreparedNativeRequestOptions = {
protocol: LlmProtocol;
backendConfig: LlmBackendConfig;
model: string;
messages: PromptMessage[];
options?: CopilotChatOptions;
execution?: CopilotProviderExecution;
withAttachment?: boolean;
attachmentCapability?: ModelAttachmentCapability;
include?: string[];
reasoning?: Record<string, unknown>;
tools?: CopilotToolSet;
middleware?: ProviderMiddlewareConfig;
};
type ProviderChatDriverPrepareResult = Omit<
PreparedNativeRequestOptions,
'execution' | 'options'
>;
type Awaitable<T> = T | Promise<T>;
type NativeBackendConfigResolver = (
execution?: CopilotProviderExecution
) => Awaitable<LlmBackendConfig>;
export type StructuredProviderDriver = {
createBackendConfig: NativeBackendConfigResolver;
prepareMessages?: (
messages: PromptMessage[],
backendConfig: LlmBackendConfig,
options: NonNullable<CopilotStructuredOptions>
) => Promise<PromptMessage[]>;
shouldRetry?: (context: {
error: unknown;
attempt: number;
options: NonNullable<CopilotStructuredOptions>;
}) => Awaitable<boolean>;
mapError: (error: unknown) => unknown;
};
export type EmbeddingProviderDriver = {
createBackendConfig: NativeBackendConfigResolver;
defaultDimensions?: number;
taskType?: string;
mapError: (error: unknown) => unknown;
};
export type RerankProviderDriver = {
createBackendConfig: NativeBackendConfigResolver;
mapError: (error: unknown) => unknown;
};
export type ImageProviderDriver = {
createBackendConfig: NativeBackendConfigResolver;
prepareMessages?: (
messages: PromptMessage[],
backendConfig: LlmBackendConfig,
options: NonNullable<CopilotImageOptions>
) => Promise<PromptMessage[]>;
mapError: (error: unknown) => unknown;
};
export type ProviderMetricLabels = Record<
string,
string | number | boolean | undefined
>;
export type ProviderExecutionDrivers = {
chat?: ProviderChatDriver;
structured?: StructuredProviderDriver;
embedding?: EmbeddingProviderDriver;
rerank?: RerankProviderDriver;
image?: ImageProviderDriver;
};
export type ProviderDriverSpec = NativeProviderDriverBase & {
chat?: NativeChatDriverOverrides | false;
structured?: NativeStructuredDriverOverrides | false;
embedding?: NativeEmbeddingDriverOverrides | false;
rerank?: NativeRerankDriverOverrides | false;
image?: NativeImageDriverOverrides | false;
};
export type ProviderRuntimeHostSeed = {
model: ProviderModelRuntimeContext;
resolveExecutionDrivers: () => ProviderExecutionDrivers | undefined;
selectModel: NativeChatDriverBase['selectModel'];
checkParams: NativeChatDriverBase['checkParams'];
getAttachCapability: (
model: CopilotProviderModel,
outputType: ModelOutputType
) => ModelAttachmentCapability | undefined;
getActiveProviderMiddleware: (
execution?: CopilotProviderExecution
) => ProviderMiddlewareConfig;
getTools: (
options: CopilotChatOptions,
model: string
) => Promise<CopilotToolSet>;
metricLabels: (
model: string,
labels?: ProviderMetricLabels,
execution?: CopilotProviderExecution
) => ProviderMetricLabels;
};
export type ProviderChatDriverPrepareInput = {
kind: 'text' | 'streamText' | 'streamObject';
cond: ModelConditions;
messages: PromptMessage[];
options: CopilotChatOptions;
execution?: CopilotProviderExecution;
};
export type ProviderChatDriver = {
prepare: (input: {
kind: ProviderChatDriverPrepareInput['kind'];
cond: ProviderChatDriverPrepareInput['cond'];
messages: ProviderChatDriverPrepareInput['messages'];
options: ProviderChatDriverPrepareInput['options'];
execution?: ProviderChatDriverPrepareInput['execution'];
}) => Promise<ProviderChatDriverPrepareResult | null>;
mapError: (error: unknown) => unknown;
};
type NativeProviderDriverBase = Pick<
StructuredProviderDriver,
'createBackendConfig' | 'mapError'
>;
type ChatToolingResult = Pick<
ProviderChatDriverPrepareResult,
'tools' | 'middleware'
>;
type NativeChatDriverBase = NativeProviderDriverBase & {
checkParams: (input: {
cond: ModelFullConditions;
messages?: PromptMessage[];
embeddings?: string[];
options?:
| CopilotChatOptions
| CopilotStructuredOptions
| CopilotImageOptions;
withAttachment?: boolean;
execution?: CopilotProviderExecution;
}) => Promise<ModelFullConditions>;
selectModel: (
cond: ModelFullConditions,
execution?: CopilotProviderExecution
) => CopilotProviderModel;
getTools?: (
options: CopilotChatOptions,
model: string
) => Promise<CopilotToolSet>;
getActiveProviderMiddleware?: (
execution?: CopilotProviderExecution
) => ProviderMiddlewareConfig;
};
type NativeStructuredDriverOverrides = Partial<StructuredProviderDriver>;
type NativeEmbeddingDriverOverrides = Partial<EmbeddingProviderDriver>;
type NativeRerankDriverOverrides = Partial<RerankProviderDriver>;
type NativeImageDriverOverrides = Partial<ImageProviderDriver>;
type NativeChatDriverContext = {
input: ProviderChatDriverPrepareInput;
outputType: ModelOutputType;
normalizedCond: ModelFullConditions;
model: CopilotProviderModel;
backendConfig: LlmBackendConfig;
protocol: LlmProtocol;
messages: PromptMessage[];
options: NonNullable<CopilotChatOptions>;
execution?: CopilotProviderExecution;
};
type NativeChatDriverOverrides = {
resolveOutputType?: (
kind: ProviderChatDriverPrepareInput['kind']
) => ModelOutputType | null;
withAttachment?: boolean;
prepareMessages?: (
context: Omit<NativeChatDriverContext, 'messages'>
) => Awaitable<PromptMessage[]>;
resolveTooling?: (
context: NativeChatDriverContext
) => Awaitable<ChatToolingResult>;
resolveRequestOptions?: (
context: NativeChatDriverContext
) => Awaitable<
Partial<
Pick<
ProviderChatDriverPrepareResult,
'withAttachment' | 'attachmentCapability' | 'include' | 'reasoning'
>
>
>;
};
export function createNativeProviderDriverFactory(
base: NativeProviderDriverBase
) {
return {
structured(
overrides: NativeStructuredDriverOverrides = {}
): StructuredProviderDriver {
return {
createBackendConfig:
overrides.createBackendConfig ?? base.createBackendConfig,
mapError: overrides.mapError ?? base.mapError,
...(overrides.prepareMessages
? { prepareMessages: overrides.prepareMessages }
: {}),
...(overrides.shouldRetry
? { shouldRetry: overrides.shouldRetry }
: {}),
};
},
embedding(
overrides: NativeEmbeddingDriverOverrides = {}
): EmbeddingProviderDriver {
return {
createBackendConfig:
overrides.createBackendConfig ?? base.createBackendConfig,
mapError: overrides.mapError ?? base.mapError,
...(overrides.defaultDimensions !== undefined
? { defaultDimensions: overrides.defaultDimensions }
: {}),
...(overrides.taskType ? { taskType: overrides.taskType } : {}),
};
},
rerank(overrides: NativeRerankDriverOverrides = {}): RerankProviderDriver {
return {
createBackendConfig:
overrides.createBackendConfig ?? base.createBackendConfig,
mapError: overrides.mapError ?? base.mapError,
};
},
image(overrides: NativeImageDriverOverrides = {}): ImageProviderDriver {
return {
createBackendConfig:
overrides.createBackendConfig ?? base.createBackendConfig,
mapError: overrides.mapError ?? base.mapError,
...(overrides.prepareMessages
? { prepareMessages: overrides.prepareMessages }
: {}),
};
},
};
}
function compileProviderChatDriver(
spec: NativeProviderDriverBase & NativeChatDriverOverrides,
base: NativeChatDriverBase
): ProviderChatDriver {
return {
prepare: async (input: ProviderChatDriverPrepareInput) => {
const options: NonNullable<CopilotChatOptions> = input.options ?? {};
const resolvedOutputType = spec.resolveOutputType?.(input.kind);
const outputType =
resolvedOutputType === undefined
? input.kind === 'streamObject'
? ModelOutputType.Object
: ModelOutputType.Text
: resolvedOutputType;
if (!outputType) {
return null;
}
const normalizedCond = await base.checkParams({
messages: input.messages,
cond: {
...input.cond,
outputType,
},
options,
execution: input.execution,
...(spec.withAttachment !== undefined
? { withAttachment: spec.withAttachment }
: {}),
});
const model = base.selectModel(normalizedCond, input.execution);
const backendConfig = await spec.createBackendConfig(input.execution);
const route = resolveProviderModelRoute(model, outputType);
if (!route.protocol) {
throw new Error(`Missing native protocol for model ${model.id}`);
}
const partialContext = {
input,
outputType,
normalizedCond,
model,
backendConfig:
route.requestLayer === backendConfig.request_layer
? backendConfig
: { ...backendConfig, request_layer: route.requestLayer },
protocol: route.protocol,
options,
execution: input.execution,
};
const messages = spec.prepareMessages
? await spec.prepareMessages(partialContext)
: input.messages;
const context = {
...partialContext,
messages,
};
const tooling = spec.resolveTooling
? await spec.resolveTooling(context)
: {
...(base.getTools
? { tools: await base.getTools(options, model.id) }
: {}),
...(base.getActiveProviderMiddleware
? {
middleware: base.getActiveProviderMiddleware(input.execution),
}
: {}),
};
const requestOptions = spec.resolveRequestOptions
? await spec.resolveRequestOptions(context)
: {};
return {
protocol: context.protocol,
backendConfig: context.backendConfig,
model: model.id,
messages,
...(spec.withAttachment === false ? { withAttachment: false } : {}),
...requestOptions,
...tooling,
};
},
mapError: spec.mapError,
};
}
export function createNativeExecutionDriverSpec(
input: ProviderDriverSpec,
runtimeBase: NativeChatDriverBase
): ProviderExecutionDrivers {
const driverBase = {
createBackendConfig: input.createBackendConfig,
mapError: input.mapError,
};
const nativeDrivers = createNativeProviderDriverFactory(driverBase);
return {
...(input.chat !== false
? {
chat: compileProviderChatDriver(
{ ...driverBase, ...input.chat },
runtimeBase
),
}
: {}),
...(input.structured !== false
? {
structured: nativeDrivers.structured(input.structured ?? undefined),
}
: {}),
...(input.embedding !== false
? {
embedding: nativeDrivers.embedding(input.embedding ?? undefined),
}
: {}),
...(input.rerank !== false
? { rerank: nativeDrivers.rerank(input.rerank ?? undefined) }
: {}),
...(input.image !== false
? { image: nativeDrivers.image(input.image ?? undefined) }
: {}),
};
}
packages/backend/server/src/plugins/copilot/providers/provider-tokens.ts
-18
View File
@@ -1,18 +0,0 @@
import {
AnthropicOfficialProvider,
AnthropicVertexProvider,
} from './anthropic';
import { CloudflareWorkersAIProvider } from './cloudflare';
import { FalProvider } from './fal';
import { GeminiGenerativeProvider, GeminiVertexProvider } from './gemini';
import { OpenAIProvider } from './openai';
export const CopilotProviders = [
OpenAIProvider,
CloudflareWorkersAIProvider,
FalProvider,
GeminiGenerativeProvider,
GeminiVertexProvider,
AnthropicOfficialProvider,
AnthropicVertexProvider,
];
packages/backend/server/src/plugins/copilot/providers/provider.ts
-249
View File
@@ -1,249 +0,0 @@
import { Inject, Injectable, Logger } from '@nestjs/common';
import { Config } from '../../../base';
import type { NodeTextMiddleware, ProviderMiddlewareConfig } from '../config';
import { ToolExecutorHost } from '../runtime/hosts/tool-executor-host';
import { mapNativeSemanticError } from '../runtime/native-errors';
import type { ToolLoopBackend } from '../runtime/tool/bridge';
import type { CopilotTool, CopilotToolSet } from '../tools';
import { resolveProviderMiddleware } from './provider-middleware';
import {
checkProviderParams,
getAttachCapability as getAttachCapabilityHelper,
matchProviderModel as matchProviderModelHelper,
type ProviderModelRuntimeContext,
requireProviderModelSelection,
resolveProviderModel,
} from './provider-model-runtime';
import {
type CopilotProviderExecution,
createNativeExecutionDriverSpec,
type ProviderDriverSpec,
type ProviderExecutionDrivers,
type ProviderRuntimeHostSeed,
} from './provider-runtime-contract';
import {
type CopilotChatOptions,
CopilotChatTools,
type CopilotImageOptions,
type CopilotModelBackendKind,
CopilotProviderModel,
CopilotProviderType,
type CopilotStructuredOptions,
type ModelAttachmentCapability,
ModelFullConditions,
ModelOutputType,
type PromptMessage,
} from './types';
export type {
CopilotProviderExecution,
ProviderDriverSpec,
ProviderExecutionDrivers,
ProviderRuntimeHostSeed,
} from './provider-runtime-contract';
@Injectable()
export abstract class CopilotProvider<C = any> {
protected readonly logger = new Logger(this.constructor.name);
protected readonly MAX_STEPS = 20;
abstract readonly type: CopilotProviderType;
protected abstract resolveModelBackendKind(
execution?: CopilotProviderExecution
): CopilotModelBackendKind;
abstract configured(execution?: CopilotProviderExecution): boolean;
@Inject() protected readonly AFFiNEConfig!: Config;
@Inject() protected readonly toolExecutorHost!: ToolExecutorHost;
get maxSteps() {
return this.MAX_STEPS;
}
protected resolveModelRuntimeContext(
execution?: CopilotProviderExecution
): ProviderModelRuntimeContext {
return {
type: this.type,
backendKind: this.resolveModelBackendKind(execution),
};
}
protected get modelRuntimeContext(): ProviderModelRuntimeContext {
return this.resolveModelRuntimeContext();
}
getDriverSpec(): ProviderDriverSpec | undefined {
return undefined;
}
getExecutionDrivers(): ProviderExecutionDrivers | undefined {
const spec = this.getDriverSpec();
return spec ? this.createDriverSpec(spec) : undefined;
}
protected createDriverSpec(
spec: ProviderDriverSpec
): ProviderExecutionDrivers {
return createNativeExecutionDriverSpec(spec, {
createBackendConfig: spec.createBackendConfig,
mapError: error => {
const mapped = mapNativeSemanticError(error);
return mapped === error ? spec.mapError(error) : mapped;
},
checkParams: input =>
checkProviderParams(
this.resolveModelRuntimeContext(input.execution),
input
),
selectModel: (cond, execution) =>
requireProviderModelSelection(
this.resolveModelRuntimeContext(execution),
cond
),
getTools: this.getTools.bind(this),
getActiveProviderMiddleware: this.getActiveProviderMiddleware.bind(this),
});
}
selectModel(
cond: ModelFullConditions,
execution?: CopilotProviderExecution
): CopilotProviderModel {
return requireProviderModelSelection(
this.resolveModelRuntimeContext(execution),
cond
);
}
checkParams(input: {
cond: ModelFullConditions;
messages?: PromptMessage[];
embeddings?: string[];
options?:
| CopilotChatOptions
| CopilotStructuredOptions
| CopilotImageOptions;
withAttachment?: boolean;
execution?: CopilotProviderExecution;
}) {
return checkProviderParams(
this.resolveModelRuntimeContext(input.execution),
input
);
}
getRuntimeHostSeed(): ProviderRuntimeHostSeed {
return {
model: this.resolveModelRuntimeContext(),
resolveExecutionDrivers: () => this.getExecutionDrivers(),
selectModel: this.selectModel.bind(this),
checkParams: this.checkParams.bind(this),
getAttachCapability: this.getAttachCapability.bind(this),
getActiveProviderMiddleware: this.getActiveProviderMiddleware.bind(this),
getTools: this.getTools.bind(this),
metricLabels: this.metricLabels.bind(this),
};
}
protected getExecutionProfile(execution?: CopilotProviderExecution) {
return execution?.profile?.type === this.type
? execution.profile
: undefined;
}
getActiveProviderMiddleware(
execution?: CopilotProviderExecution
): ProviderMiddlewareConfig {
return (
this.getExecutionProfile(execution)?.middleware ??
resolveProviderMiddleware(this.type)
);
}
metricLabels(
model: string,
labels: Record<string, string | number | boolean | undefined> = {},
execution?: CopilotProviderExecution
) {
return {
model,
providerId: execution?.providerId ?? `${this.type}-default`,
...labels,
};
}
protected get config(): C {
return this.AFFiNEConfig.copilot.providers[this.type] as C;
}
protected getConfig(execution?: CopilotProviderExecution): C {
const profile = this.getExecutionProfile(execution);
if (profile) {
return profile.config as C;
}
return this.config;
}
getAttachCapability(
model: CopilotProviderModel,
outputType: ModelOutputType
): ModelAttachmentCapability | undefined {
return getAttachCapabilityHelper(model, outputType);
}
// make it async to allow dynamic check available models in some providers
async match(
cond: ModelFullConditions = {},
execution?: CopilotProviderExecution
): Promise<boolean> {
return (
this.configured(execution) &&
matchProviderModelHelper(this.resolveModelRuntimeContext(execution), cond)
);
}
resolveModel(
modelId: string,
execution?: CopilotProviderExecution
): CopilotProviderModel | undefined {
return resolveProviderModel(
this.resolveModelRuntimeContext(execution),
modelId
);
}
protected getProviderSpecificTools(
_toolName: CopilotChatTools,
_model: string
): [string, CopilotTool?] | undefined {
return;
}
// use for tool use, shared between providers
async getTools(
options: CopilotChatOptions,
model: string
): Promise<CopilotToolSet> {
this.logger.debug(`getTools: ${JSON.stringify(options?.tools ?? [])}`);
return await this.toolExecutorHost.getTools(
options,
model,
this.getProviderSpecificTools.bind(this)
);
}
createNativeAdapter(
backend: ToolLoopBackend,
tools: CopilotToolSet,
nodeTextMiddleware?: NodeTextMiddleware[],
options: {
maxSteps?: number;
nodeTextMiddleware?: NodeTextMiddleware[];
} = {}
) {
return this.toolExecutorHost.createNativeAdapter(backend, tools, {
...options,
nodeTextMiddleware: nodeTextMiddleware ?? options.nodeTextMiddleware,
});
}
}
packages/backend/server/src/plugins/copilot/providers/registry-service.ts
-28
View File
@@ -1,28 +0,0 @@
import { Injectable } from '@nestjs/common';
import { Config } from '../../../base';
import {
buildProviderRegistry,
type CopilotProviderRegistry,
type CopilotProvidersConfigInput,
} from './provider-registry';
@Injectable()
export class CopilotProviderRegistryService {
private lastConfig?: CopilotProvidersConfigInput;
private lastRegistry?: CopilotProviderRegistry;
constructor(private readonly config: Config) {}
getRegistry(): CopilotProviderRegistry {
const providerConfig = this.config.copilot.providers;
if (this.lastConfig === providerConfig && this.lastRegistry) {
return this.lastRegistry;
}
const registry = buildProviderRegistry(providerConfig);
this.lastConfig = providerConfig;
this.lastRegistry = registry;
return registry;
}
}
packages/backend/server/src/plugins/copilot/providers/types.ts
+8 -8
View File
@@ -1,4 +1,4 @@
import { AiPromptRole } from '@prisma/client';
import { AiSessionMessageRole } from '@prisma/client';
import { z } from 'zod';
import { JSONSchema } from '../../../base';
@@ -7,7 +7,6 @@ import type {
CapabilityModelCapability,
ModelConditionsContract,
} from '../../../native';
import type { CopilotModelBackendKind } from '../runtime/contracts';
import {
type StreamObject,
StreamObjectSchema,
@@ -97,7 +96,6 @@ export const PromptToolsSchema = z
export const PromptConfigStrictSchema = z.object({
tools: PromptToolsSchema.nullable().optional(),
proModels: z.array(z.string()).nullable().optional(),
// params requirements
requireContent: z.boolean().nullable().optional(),
requireAttachment: z.boolean().nullable().optional(),
@@ -108,7 +106,7 @@ export const PromptConfigStrictSchema = z.object({
presencePenalty: z.number().nullable().optional(),
temperature: z.number().nullable().optional(),
topP: z.number().nullable().optional(),
maxTokens: z.number().nullable().optional(),
maxOutputTokens: z.number().nullable().optional(),
// fal
modelName: z.string().nullable().optional(),
loras: z
@@ -132,7 +130,7 @@ export type PromptTools = z.infer<typeof PromptToolsSchema>;
export const EmbeddingMessage = z.array(z.string().trim().min(1)).min(1);
export const ChatMessageRole = Object.values(AiPromptRole) as [
export const ChatMessageRole = Object.values(AiSessionMessageRole) as [
'system',
'assistant',
'user',
@@ -268,6 +266,8 @@ const CopilotProviderOptionsSchema = z.object({
billingUnitId: z.string().optional(),
taskId: z.string().optional(),
actionId: z.string().optional(),
builtInRouteId: z.string().optional(),
managedTargetId: z.string().optional(),
quotaBackedRoutesAllowed: z.boolean().optional(),
featureKind: z
.enum([
@@ -380,8 +380,8 @@ export interface CopilotProviderModel {
capabilities: ModelCapability[];
}
export type { CopilotModelBackendKind };
export type ModelConditions = Omit<ModelConditionsContract, 'outputType'>;
export type ModelConditions = Omit<ModelConditionsContract, 'outputType'> & {
profileId?: string;
};
export type ModelFullConditions = ModelConditionsContract;
packages/backend/server/src/plugins/copilot/resolver.ts
+60 -87
View File
@@ -31,11 +31,12 @@ import { CurrentUser } from '../../core/auth';
import { DocAction, PermissionAccess } from '../../core/permission';
import { UserType } from '../../core/user';
import type { ListSessionOptions, UpdateChatSession } from '../../models';
import { llmGetBuiltInRouteOptions } from '../../native';
import { ByokEntitlementPolicy } from './byok';
import { CompatHistoryProjector } from './compat/history-projector';
import { ConversationInboxService } from './conversation/inbox';
import { PromptService } from './prompt/service';
import { CopilotProviderFactory } from './providers/factory';
import { ModelOutputType, type StreamObject } from './providers/types';
import { CopilotEnabled } from './feature';
import type { StreamObject } from './providers/types';
import { ChatSessionService } from './session';
import { type ChatHistory, type ChatMessage, SubmittedMessage } from './types';
@@ -256,12 +257,6 @@ class CopilotHistoriesType implements Omit<ChatHistory, 'userId'> {
@Field(() => String)
promptName!: string;
@Field(() => String)
model!: string;
@Field(() => [String])
optionalModels!: string[];
@Field(() => String, {
description: 'An mark identifying which view to use to display the session',
nullable: true,
@@ -274,11 +269,6 @@ class CopilotHistoriesType implements Omit<ChatHistory, 'userId'> {
@Field(() => String, { nullable: true })
title!: string | null;
@Field(() => Number, {
description: 'The number of tokens used in the session',
})
tokens!: number;
@Field(() => [ChatMessageType])
messages!: ChatMessageType[];
@@ -303,27 +293,6 @@ class CopilotQuotaType {
used!: number;
}
@ObjectType()
class CopilotModelType {
@Field(() => String)
id!: string;
@Field(() => String)
name!: string;
}
@ObjectType()
export class CopilotModelsType {
@Field(() => String)
defaultModel!: string;
@Field(() => [CopilotModelType])
optionalModels!: CopilotModelType[];
@Field(() => [CopilotModelType])
proModels!: CopilotModelType[];
}
@ObjectType()
export class CopilotSessionType {
@Field(() => ID)
@@ -343,12 +312,33 @@ export class CopilotSessionType {
@Field(() => String)
promptName!: string;
}
@ObjectType('CopilotRouteTarget')
class CopilotRouteTargetType {
@Field(() => String)
id!: string;
@Field(() => String)
model!: string;
displayName!: string;
@Field(() => [String])
optionalModels!: string[];
@Field(() => String)
minimumTier!: string;
@Field(() => Boolean)
available!: boolean;
}
@ObjectType('CopilotRouteOptions')
class CopilotRouteOptionsType {
@Field(() => String)
routeId!: string;
@Field(() => String, { nullable: true })
defaultTargetId!: string | null;
@Field(() => [CopilotRouteTargetType])
choices!: CopilotRouteTargetType[];
}
// ================== Resolver ==================
@@ -360,20 +350,47 @@ export class CopilotType {
}
@Throttle()
@CopilotEnabled()
@Resolver(() => CopilotType)
export class CopilotResolver {
private readonly modelNames = new Map<string, string>();
constructor(
private readonly ac: PermissionAccess,
private readonly mutex: RequestMutex,
private readonly prompt: PromptService,
private readonly chatSession: ChatSessionService,
private readonly historyProjector: CompatHistoryProjector,
private readonly inbox: ConversationInboxService,
private readonly providerFactory: CopilotProviderFactory
private readonly entitlement: ByokEntitlementPolicy
) {}
@ResolveField(() => CopilotRouteOptionsType, {
nullable: true,
description: 'List native built-in route choices for a prompt',
complexity: 2,
})
async routeOptions(
@CurrentUser() user: CurrentUser,
@Args('promptName') promptName: string
): Promise<CopilotRouteOptionsType | null> {
const options = llmGetBuiltInRouteOptions(promptName);
if (!options) return null;
if (env.selfhosted) {
return { routeId: options.routeId, defaultTargetId: null, choices: [] };
}
const premium = await this.entitlement.hasAiPlan(user.id);
return {
routeId: options.routeId,
defaultTargetId: premium
? (options.premiumDefaultTargetId ?? null)
: (options.standardDefaultTargetId ?? null),
choices: options.choices.map(choice => ({
id: choice.id,
displayName: choice.displayName,
minimumTier: choice.minimumTier,
available: premium || choice.minimumTier === 'Standard',
})),
};
}
@ResolveField(() => CopilotQuotaType, {
name: 'quota',
description: 'Get the quota of the user in the workspace',
@@ -408,51 +425,6 @@ export class CopilotResolver {
return { userId: user.id, workspaceId, docId: docId || undefined };
}
@ResolveField(() => CopilotModelsType, {
description:
'List available models for a prompt, with human-readable names',
complexity: 2,
})
async models(
@Args('promptName') promptName: string
): Promise<CopilotModelsType> {
const prompt = await this.prompt.get(promptName);
if (!prompt) {
throw new NotFoundException('Prompt not found');
}
const convertModels = async (ids: string[]) => {
const models = await Promise.all(
ids.map(async id => {
const cachedName = this.modelNames.get(id);
if (cachedName) return { id, name: cachedName };
const resolved = await this.providerFactory.resolveProvider({
modelId: id,
outputType: ModelOutputType.Text,
});
const name = resolved?.provider.resolveModel(
resolved.modelId ?? id,
resolved.execution
)?.name;
if (name) {
this.modelNames.set(id, name);
return { id, name };
}
return null;
})
);
return models.filter(model => !!model) as CopilotModelType[];
};
const proModels = prompt.config?.proModels || [];
return {
defaultModel: prompt.model,
optionalModels: await convertModels(prompt.optionalModels),
proModels: await convertModels(proModels),
};
}
@ResolveField(() => CopilotSessionType, {
description: 'Get the session by id',
complexity: 2,
@@ -813,6 +785,7 @@ export class CopilotResolver {
}
@Throttle()
@CopilotEnabled()
@Resolver(() => UserType)
export class UserCopilotResolver {
constructor(private readonly ac: PermissionAccess) {}
packages/backend/server/src/plugins/copilot/runtime/action-output-projector.ts
+1 -1
View File
@@ -110,7 +110,7 @@ function isImageAction(actionId: string) {
}
function resolveProjector(actionId: string): ActionResultProjector | null {
if (actionId.startsWith('transcript.audio.')) {
if (actionId === 'transcript.audio') {
return null;
}
if (isImageAction(actionId)) {
packages/backend/server/src/plugins/copilot/runtime/action-runtime-bridge.ts
+76 -161
View File
@@ -1,15 +1,10 @@
import { Injectable, Optional } from '@nestjs/common';
import { Injectable } from '@nestjs/common';
import { Models } from '../../../models';
import type { AiActionRunStatus } from '../../../models/copilot-action-run';
import {
type NativeActionEvent,
type NativeActionRuntimeInput,
runNativeActionRecipePreparedStream,
} from '../../../native';
import { type NativeActionEvent } from '../../../native';
import type {
CopilotImageOptions,
CopilotProviderType,
CopilotStructuredOptions,
PromptMessage,
} from '../providers/types';
@@ -18,18 +13,10 @@ import {
projectActionResultToAssistantTurn,
summarizeActionResult,
} from './action-output-projector';
import {
buildStructuredResponseFromSchemaJson,
type RequiredStructuredOutputContract,
} from './contracts';
import { ExecutionPlanBuilder } from './execution-plan';
import { CapabilityRuntime } from './capability-runtime';
import { type RequiredStructuredOutputContract } from './contracts';
import { TurnPersistence } from './hosts/turn-persistence';
type ActionRuntimeBridgeNativeInput = Omit<
NativeActionRuntimeInput,
'recipeId' | 'recipeVersion'
>;
export type ActionRuntimeBridgeInput = {
userId: string;
workspaceId: string;
@@ -42,26 +29,18 @@ export type ActionRuntimeBridgeInput = {
attempt?: number;
retryOf?: string | null;
inputSnapshot?: unknown;
nativeInput?: ActionRuntimeBridgeNativeInput;
onRunCreated?: (
context: ActionRuntimeBridgeRunContext
) => Promise<void> | void;
prepareStructuredRoutes?: {
stepId?: string;
step: {
slot: string;
builtInRouteId: string;
profileId?: string;
modelId?: string;
messages: PromptMessage[];
options?: CopilotStructuredOptions;
prefer?: CopilotProviderType;
responseSchemaJson?: Record<string, unknown>;
options?: CopilotStructuredOptions | CopilotImageOptions;
responseContract?: RequiredStructuredOutputContract;
};
prepareImageRoutes?: {
stepId?: string;
modelId?: string;
messages: PromptMessage[];
options?: CopilotImageOptions;
prefer?: CopilotProviderType;
};
persistAttachment?: (attachment: unknown) => Promise<unknown> | unknown;
signal?: AbortSignal;
};
@@ -107,94 +86,41 @@ export class ActionRuntimeBridge {
constructor(
private readonly models: Models,
private readonly turnPersistence: TurnPersistence,
@Optional() private readonly plans?: ExecutionPlanBuilder
private readonly runtime: CapabilityRuntime
) {}
protected runNativeStream(
input: NativeActionRuntimeInput,
signal?: AbortSignal
) {
return runNativeActionRecipePreparedStream(input, signal);
}
private async prepareNativeInput(
input: ActionRuntimeBridgeInput
): Promise<ActionRuntimeBridgeNativeInput & { input: unknown }> {
const nativeInput = {
...input.nativeInput,
input: input.nativeInput?.input ?? {},
};
const structured = input.prepareStructuredRoutes;
const image = input.prepareImageRoutes;
if (!structured && !image) {
return nativeInput;
}
if (!this.plans) {
throw new Error('Action route preparation is not available');
}
const state =
nativeInput.input && typeof nativeInput.input === 'object'
? { ...(nativeInput.input as Record<string, unknown>) }
: {};
if (structured) {
const responseContract =
structured.responseContract ??
(buildStructuredResponseFromSchemaJson(
structured.responseSchemaJson ?? { type: 'object' }
) as RequiredStructuredOutputContract);
const plan = await this.plans.buildStructuredPlan(
{ modelId: structured.modelId },
structured.messages,
structured.options,
structured.prefer ? { prefer: structured.prefer } : undefined,
responseContract
private async execute(input: ActionRuntimeBridgeInput) {
const step = input.step;
if (step.responseContract) {
const output = await this.runtime.generateStructuredValue(
{ profileId: step.profileId, modelId: step.modelId },
step.messages,
{
...(step.options as CopilotStructuredOptions | undefined),
builtInRouteId: step.builtInRouteId,
},
step.responseContract,
undefined,
step.slot
);
const preparedRoutes = plan.nativeDispatch?.structured?.routes;
if (!preparedRoutes?.length) {
throw new Error('No native structured provider route prepared');
}
const existingPreparedRoutes =
state.preparedRoutes &&
typeof state.preparedRoutes === 'object' &&
!Array.isArray(state.preparedRoutes)
? (state.preparedRoutes as Record<string, unknown>)
: {};
state.preparedRoutes = {
...existingPreparedRoutes,
[structured.stepId ?? 'generate']: preparedRoutes,
};
return { result: output.value, attachments: [] };
}
if (image) {
const plan = await this.plans.buildImagePlan(
{ modelId: image.modelId },
image.messages,
image.options,
image.prefer ? { prefer: image.prefer } : undefined
);
const preparedRoutes = plan.nativeDispatch?.image?.routes;
if (!preparedRoutes?.length) {
throw new Error('No native image provider route prepared');
}
const existingPreparedRoutes =
state.preparedRoutes &&
typeof state.preparedRoutes === 'object' &&
!Array.isArray(state.preparedRoutes)
? (state.preparedRoutes as Record<string, unknown>)
: {};
state.preparedRoutes = {
...existingPreparedRoutes,
[image.stepId ?? 'generate-image']: preparedRoutes,
};
const images = [];
for await (const image of this.runtime.streamImageArtifacts(
{ profileId: step.profileId, modelId: step.modelId },
step.messages,
{
...(step.options as CopilotImageOptions | undefined),
builtInRouteId: step.builtInRouteId,
},
undefined,
step.slot
)) {
images.push(image);
}
return {
...nativeInput,
input: state,
};
const result = images[0];
if (!result) throw new Error('Action image generation produced no image');
return { result, attachments: [result] };
}
private async projectAssistantResult(
@@ -273,28 +199,36 @@ export class ActionRuntimeBridge {
let finalEvent: NativeActionEvent | undefined;
const attachments: unknown[] = [];
try {
const nativeInput = await this.prepareNativeInput({
...inputWithBillingUnit,
});
for await (const event of this.runNativeStream(
{
...nativeInput,
recipeId: inputWithBillingUnit.actionId,
recipeVersion: inputWithBillingUnit.actionVersion,
},
inputWithBillingUnit.signal
)) {
finalEvent = event;
let projectedEvent = event;
if (event.type === 'attachment') {
const attachment = input.persistAttachment
? await input.persistAttachment(event.attachment)
: event.attachment;
attachments.push(attachment);
projectedEvent = { ...event, attachment };
}
yield { ...projectedEvent, runId: run.id };
const actionStart: NativeActionEvent = {
type: 'action_start',
actionId: input.actionId,
actionVersion: input.actionVersion,
status: 'running',
};
yield { ...actionStart, runId: run.id };
const output = await this.execute(inputWithBillingUnit);
for (const artifact of output.attachments) {
const attachment = input.persistAttachment
? await input.persistAttachment(artifact)
: artifact;
attachments.push(attachment);
yield {
type: 'attachment',
actionId: input.actionId,
actionVersion: input.actionVersion,
status: 'running',
attachment,
runId: run.id,
};
}
finalEvent = {
type: 'action_done',
actionId: input.actionId,
actionVersion: input.actionVersion,
status: 'succeeded',
result: output.result,
};
yield { ...finalEvent, runId: run.id };
} catch (error) {
finalEvent = {
type: 'error',
@@ -351,33 +285,14 @@ export class ActionRuntimeBridge {
): 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,
step: {
...input.step,
options: {
...input.step.options,
actionId: input.step.options?.actionId ?? input.actionId,
billingUnitId: input.step.options?.billingUnitId ?? billingUnitId,
},
},
};
}
}
packages/backend/server/src/plugins/copilot/runtime/capability-runtime.ts
+316 -111
View File
@@ -1,7 +1,27 @@
/* oxlint-disable import/no-cycle -- Tool callbacks can invoke nested Copilot prompts. */
import { Injectable } from '@nestjs/common';
import { CopilotPromptInvalid } from '../../../base';
import { ValidatedStructuredValueSchema } from '../core';
import { Config } from '../../../base/config';
import { CopilotPromptInvalid } from '../../../base/error/errors.gen';
import { BackendRuntimeProvider } from '../../../core/backend-runtime';
import {
buildLlmEmbeddingRequest,
buildLlmImageRequestFromMessages,
buildLlmRerankRequest,
type LlmImageResponse,
type LlmToolCallbackRequest,
type LlmToolLoopStreamEvent,
llmValidateJsonSchema,
} from '../../../native';
import {
getByokSourceCoverage,
getCopilotFeatureAccess,
} from '../access/feature-coverage';
import { assertCopilotEnabled } from '../availability';
import { ByokEntitlementPolicy } from '../byok/policy';
import type { ByokFeatureKind } from '../byok/types';
import { ConversationPolicy } from '../conversation/policy';
import { ValidatedStructuredValueSchema } from '../core/types';
import {
type CopilotChatOptions,
type CopilotEmbeddingOptions,
@@ -9,112 +29,294 @@ import {
type CopilotProviderType,
type CopilotRerankRequest,
type CopilotStructuredOptions,
type ModelAttachmentCapability,
type ModelConditions,
type PromptMessage,
type StreamObject,
} from '../providers/types';
import {
buildToolContracts,
type RequiredStructuredOutputContract,
requireStructuredOutputContract,
} from './contracts';
import {
ExecutionPlanBuilder,
type ExecutionPlanForKind,
} from './execution-plan';
type CopilotRuntimeEvent,
CopilotRuntimeEventConsumer,
} from './copilot-runtime-event-consumer';
import { mapNativeSemanticError } from './native-errors';
import {
NativeExecutionEngine,
type NativeImageArtifact,
} from './native-execution-engine';
buildCanonicalNativeRequest,
buildCanonicalNativeStructuredRequest,
preparePromptMessagesForNativeRequest,
} from './native-request-runtime';
import { executeToolCall } from './tool/bridge';
import { NativeProviderAdapter } from './tool/native-adapter';
import { ToolRuntime } from './tool-runtime';
type ProviderFilter = {
prefer?: CopilotProviderType;
type ProviderFilter = { prefer?: CopilotProviderType };
type RuntimeOptions = NonNullable<CopilotChatOptions> & {
dimensions?: number;
responseSchemaJson?: Record<string, unknown>;
schemaHash?: string;
strict?: boolean;
profileId?: string;
};
const providerModelId = (modelId?: string) => modelId ?? 'auto';
export type NativeImageArtifact = LlmImageResponse['images'][number];
const attachmentCapability = {
kinds: ['image', 'audio', 'file'],
sourceKinds: ['url', 'data', 'bytes', 'file_handle'],
allowRemoteUrls: true,
} satisfies ModelAttachmentCapability;
@Injectable()
export class CapabilityRuntime {
constructor(
private readonly plans: ExecutionPlanBuilder,
private readonly engine: NativeExecutionEngine
private readonly backend: BackendRuntimeProvider,
private readonly entitlement: ByokEntitlementPolicy,
private readonly conversations: ConversationPolicy,
private readonly tools: ToolRuntime,
private readonly events: CopilotRuntimeEventConsumer,
private readonly config: Config
) {}
private async executePlan<TPlan, TResult>(
build: () => Promise<TPlan>,
execute: (plan: TPlan) => Promise<TResult>
) {
return await execute(await build());
private async access(options: RuntimeOptions) {
assertCopilotEnabled(this.config);
const workspaceId = options.workspace;
const featureKind = (options.featureKind ?? 'chat') as ByokFeatureKind;
const coverage = getByokSourceCoverage(featureKind);
const [serverByok, localByok, premium] = workspaceId
? await Promise.all([
coverage.server && this.entitlement.hasServerEntitlement(workspaceId),
coverage.local &&
this.entitlement.hasLocalEntitlement(workspaceId, options.user),
this.entitlement.hasAiPlan(options.user),
])
: [false, false, await this.entitlement.hasAiPlan(options.user)];
const routeAllowed =
options.quotaBackedRoutesAllowed ??
(!getCopilotFeatureAccess(featureKind).quotaMetered ||
!options.user ||
(await this.conversations.hasQuota(options.user)));
return {
routeAllowed,
managedTier: premium ? ('Premium' as const) : ('Standard' as const),
serverByok,
localByok,
};
}
private executeStreamPlan<TPlan, TChunk>(
build: () => Promise<TPlan>,
execute: (plan: TPlan) => AsyncIterableIterator<TChunk>
): AsyncIterableIterator<TChunk> {
return (async function* () {
yield* execute(await build());
})();
private eventContext(options: RuntimeOptions) {
return {
workspaceId: options.workspace,
userId: options.user,
sessionId: options.session,
taskId: options.taskId,
actionId: options.actionId,
billingUnitId: options.billingUnitId,
featureKind: (options.featureKind ?? 'chat') as ByokFeatureKind,
};
}
private hasNativeDispatch(
plan: ExecutionPlanForKind<'embedding'> | ExecutionPlanForKind<'rerank'>,
kind: 'embedding' | 'rerank'
private targetOverride(cond: ModelConditions) {
return cond.profileId && cond.modelId
? { profileId: cond.profileId, modelId: cond.modelId }
: undefined;
}
async assertRoute(
slot: string,
cond: ModelConditions,
options: CopilotChatOptions = {}
) {
return !!plan.nativeDispatch?.[kind];
try {
await this.backend.assertCopilotRoute({
slot,
builtInRouteId: options.builtInRouteId,
workspaceId: options.workspace,
userId: options.user,
localLeaseId: options.byokLeaseId,
access: await this.access(options),
managedTargetId: options.managedTargetId,
targetOverride: this.targetOverride(cond),
});
} catch (error) {
throw mapNativeSemanticError(error);
}
}
private async execute(
slot: string,
request: unknown,
cond: ModelConditions,
options: RuntimeOptions
) {
try {
const output = await this.backend.executeCopilot({
slot,
builtInRouteId: options.builtInRouteId,
workspaceId: options.workspace,
userId: options.user,
localLeaseId: options.byokLeaseId,
access: await this.access(options),
managedTargetId: options.managedTargetId,
targetOverride: this.targetOverride(cond),
request,
});
await this.events.consume(
output.events as CopilotRuntimeEvent[],
this.eventContext(options)
);
return output.result;
} catch (error) {
throw mapNativeSemanticError(error);
}
}
private async prepareChat(
messages: PromptMessage[],
options: RuntimeOptions
) {
const toolSet = await this.tools.getTools(options, '');
const { request } = await buildCanonicalNativeRequest({
model: 'route-selected',
messages,
options,
toolContracts: buildToolContracts(toolSet),
attachmentCapability,
include: options.reasoning ? ['reasoning'] : undefined,
reasoning: options.reasoning ? { effort: 'medium' } : undefined,
});
return { request: { ...request, stream: true }, toolSet };
}
private async stream(
slot: string,
cond: ModelConditions,
messages: PromptMessage[],
options: RuntimeOptions
) {
const { request, toolSet } = await this.prepareChat(messages, options);
const rawStream = this.backend.streamCopilot<
LlmToolLoopStreamEvent | CopilotRuntimeEvent
>(
{
slot,
builtInRouteId: options.builtInRouteId,
workspaceId: options.workspace,
userId: options.user,
localLeaseId: options.byokLeaseId,
access: await this.access(options),
managedTargetId: options.managedTargetId,
targetOverride: this.targetOverride(cond),
request,
},
async requestJson => {
const toolRequest = JSON.parse(requestJson) as LlmToolCallbackRequest;
return JSON.stringify(
await executeToolCall(toolSet, toolRequest, {
signal: options.signal,
messages,
})
);
},
{ maxSteps: 20, signal: options.signal }
);
const runtimeEvents = this.events;
const eventContext = this.eventContext(options);
async function* productEvents() {
for await (const event of rawStream) {
if ('route' in event) {
await runtimeEvents.consume([event], eventContext);
} else if (event.type === 'error') {
throw mapNativeSemanticError(
new Error(
typeof event.message === 'string'
? event.message
: 'native runtime stream error'
)
);
} else {
yield event;
}
}
}
return { request, stream: productEvents() };
}
async text(
cond: ModelConditions,
messages: PromptMessage[],
options?: CopilotChatOptions,
filter?: ProviderFilter
options: CopilotChatOptions = {},
_filter?: ProviderFilter
) {
return await this.executePlan(
() => this.plans.buildTextPlan(cond, messages, options, filter),
plan => this.engine.execute(plan)
const prepared = await this.stream('prompt.text', cond, messages, options);
return await new NativeProviderAdapter(() => prepared.stream).text(
prepared.request,
options.signal,
messages
);
}
async *streamText(
cond: ModelConditions,
messages: PromptMessage[],
options?: CopilotChatOptions,
filter?: ProviderFilter
options: CopilotChatOptions = {},
_filter?: ProviderFilter
): AsyncIterableIterator<string> {
yield* this.executeStreamPlan(
() => this.plans.buildStreamTextPlan(cond, messages, options, filter),
plan => this.engine.executeStream(plan)
const prepared = await this.stream('chat.default', cond, messages, options);
yield* new NativeProviderAdapter(() => prepared.stream).streamText(
prepared.request,
options.signal,
messages
);
}
async *streamObject(
cond: ModelConditions,
messages: PromptMessage[],
options?: CopilotChatOptions,
filter?: ProviderFilter
options: CopilotChatOptions = {},
_filter?: ProviderFilter
): AsyncIterableIterator<StreamObject> {
yield* this.executeStreamPlan(
() => this.plans.buildStreamObjectPlan(cond, messages, options, filter),
plan => this.engine.executeStream(plan)
const prepared = await this.stream('chat.default', cond, messages, options);
yield* new NativeProviderAdapter(() => prepared.stream).streamObject(
prepared.request,
options.signal,
messages
);
}
async generateStructured(
cond: ModelConditions,
messages: PromptMessage[],
options?: CopilotStructuredOptions,
filter?: ProviderFilter,
responseContract?: RequiredStructuredOutputContract
options: CopilotStructuredOptions = {},
_filter?: ProviderFilter,
responseContract?: RequiredStructuredOutputContract,
slot = 'prompt.structured'
) {
return await this.executePlan(
() =>
this.plans.buildStructuredPlan(
cond,
messages,
options,
filter,
responseContract
),
plan => this.engine.execute(plan)
const contract = requireStructuredOutputContract(responseContract);
if (!contract) {
throw new CopilotPromptInvalid('Structured schema contract is required');
}
const { request } = await buildCanonicalNativeStructuredRequest({
model: 'route-selected',
messages,
options,
responseContract: contract,
attachmentCapability,
});
const result = (await this.execute(slot, request, cond, options)) as {
output_json?: unknown;
output_text: string;
};
if (result.output_json === undefined) {
throw new CopilotPromptInvalid(
'Structured response is missing output_json'
);
}
return JSON.stringify(
llmValidateJsonSchema(request.schema, result.output_json)
);
}
@@ -123,87 +325,90 @@ export class CapabilityRuntime {
messages: PromptMessage[],
options: CopilotStructuredOptions,
responseContract?: RequiredStructuredOutputContract,
filter?: ProviderFilter
filter?: ProviderFilter,
slot = 'prompt.structured'
) {
const validatedResponseContract =
requireStructuredOutputContract(responseContract);
if (!options || !validatedResponseContract) {
const contract = requireStructuredOutputContract(responseContract);
if (!contract) {
throw new CopilotPromptInvalid('Structured schema contract is required');
}
const output = await this.generateStructured(
cond,
messages,
options,
filter,
validatedResponseContract
const value = JSON.parse(
await this.generateStructured(
cond,
messages,
options,
filter,
contract,
slot
)
);
const value = JSON.parse(output);
return ValidatedStructuredValueSchema.parse({
value,
schemaHash: validatedResponseContract.schemaHash,
schemaHash: contract.schemaHash,
schemaValidationVersion: 'json-schema-v1',
provider: filter?.prefer ?? 'auto',
model: providerModelId(cond.modelId),
provider: 'auto',
model: 'route-selected',
});
}
async embeddingConfigured(modelId: string) {
try {
return this.hasNativeDispatch(
await this.plans.buildEmbeddingPlan(modelId, 'ping'),
'embedding'
);
} catch {
return false;
}
async embeddingConfigured(_modelId: string) {
return this.config.copilot.enabled;
}
async embed(
modelId: string,
_modelId: string,
input: string | string[],
options?: CopilotEmbeddingOptions
options: CopilotEmbeddingOptions = {}
) {
return await this.executePlan(
() => this.plans.buildEmbeddingPlan(modelId, input, options),
plan => this.engine.execute(plan)
);
const result = (await this.execute(
'index.embedding',
buildLlmEmbeddingRequest({
model: 'route-selected',
inputs: Array.isArray(input) ? input : [input],
dimensions: options.dimensions,
}),
{},
options
)) as { embeddings: number[][] };
return result.embeddings;
}
async rerankConfigured(modelId: string) {
try {
return this.hasNativeDispatch(
await this.plans.buildRerankPlan(modelId, {
query: 'ping',
candidates: [{ text: 'ping' }],
}),
'rerank'
);
} catch {
return false;
}
async rerankConfigured(_modelId: string) {
return this.config.copilot.enabled;
}
async rerank(
modelId: string,
_modelId: string,
request: CopilotRerankRequest,
options?: CopilotChatOptions
options: CopilotChatOptions = {}
) {
return await this.executePlan(
() => this.plans.buildRerankPlan(modelId, request, options),
plan => this.engine.execute(plan)
);
const result = (await this.execute(
'search.rerank',
buildLlmRerankRequest('route-selected', request),
{},
options
)) as { scores: number[] };
return result.scores;
}
async *streamImageArtifacts(
cond: ModelConditions,
messages: PromptMessage[],
options?: CopilotImageOptions,
filter?: ProviderFilter
options: CopilotImageOptions = {},
_filter?: ProviderFilter,
slot = 'image.generate'
): AsyncIterableIterator<NativeImageArtifact> {
yield* this.executeStreamPlan(
() => this.plans.buildImagePlan(cond, messages, options, filter),
plan => this.engine.executeImageArtifacts(plan)
);
const { quality, seed } = options;
const result = (await this.execute(
slot,
buildLlmImageRequestFromMessages({
model: 'route-selected',
messages: preparePromptMessagesForNativeRequest(messages, true),
options: { quality, seed },
}),
cond,
options
)) as LlmImageResponse;
yield* result.images;
}
}
packages/backend/server/src/plugins/copilot/runtime/contracts/execution-plan-contract.ts
-104
View File
@@ -1,104 +0,0 @@
import {
type LlmBackendConfig,
llmCompileExecutionPlan,
type LlmEmbeddingRequest,
type LlmImageRequest,
type LlmProtocol,
type LlmRequest,
type LlmRerankRequest,
type LlmStructuredRequest,
} from '../../../../native';
import type {
CopilotProviderType,
ModelConditions,
PromptMessage,
} from '../../providers/types';
// Owner: runtime core mirror facade.
// The semantic source of truth is the native/Rust execution-plan contract
// behind llmCompileExecutionPlan(); this file only keeps the TypeScript shape
// needed by Node live-plan assembly until generated/native TS types replace it.
export type ExecutionRequestKind =
| 'text'
| 'streamText'
| 'streamObject'
| 'structured'
| 'embedding'
| 'rerank'
| 'image';
export type ExecutionRoute = {
providerId: string;
protocol: LlmProtocol;
model: string;
backendConfig: LlmBackendConfig;
};
export type ExecutionTransportContract =
| { kind: 'chat'; request: LlmRequest }
| { kind: 'structured'; request: LlmStructuredRequest }
| { kind: 'embedding'; request: LlmEmbeddingRequest }
| { kind: 'rerank'; request: LlmRerankRequest }
| { kind: 'image'; request: LlmImageRequest };
export type SerializableExecutionPlanRequest =
| {
kind: 'text' | 'streamText' | 'streamObject';
cond: ModelConditions;
messages: PromptMessage[];
options?: Record<string, unknown>;
}
| {
kind: 'structured';
cond: ModelConditions;
messages: PromptMessage[];
options?: Record<string, unknown>;
}
| {
kind: 'image';
cond: ModelConditions;
messages: PromptMessage[];
options?: Record<string, unknown>;
}
| {
kind: 'embedding';
cond: ModelConditions;
modelId: string;
input: string | string[];
options?: Record<string, unknown>;
}
| {
kind: 'rerank';
cond: ModelConditions;
modelId: string;
request: {
query: string;
candidates: { id?: string; text: string }[];
topK?: number;
};
options?: Record<string, unknown>;
};
export type SerializableExecutionPlan = {
routes: ExecutionRoute[];
request: SerializableExecutionPlanRequest;
transport?: ExecutionTransportContract;
routePolicy: { fallbackOrder: string[] };
runtimePolicy: {
prefer?: CopilotProviderType;
maxSteps?: number;
};
attachmentPolicy: {
materializeRemoteAttachments: boolean;
};
responsePostprocess: {
mode: ExecutionRequestKind;
};
hostContext?: {
currentMessages?: PromptMessage[];
};
};
export function parseExecutionPlan(value: unknown) {
return llmCompileExecutionPlan<SerializableExecutionPlan>(value);
}
packages/backend/server/src/plugins/copilot/runtime/contracts/index.ts
-2
View File
@@ -1,5 +1,3 @@
export * from './execution-plan-contract';
export * from './native-contract';
export * from './prompt-contract';
export * from './runtime-event-contract';
export * from './shared';
packages/backend/server/src/plugins/copilot/runtime/contracts/native-contract.ts
-97
View File
@@ -1,97 +0,0 @@
import serverNativeModule, {
type CapabilityMatchRequest,
type CapabilityMatchResponse,
type ModelRegistryMatchRequest,
type ModelRegistryMatchResponse,
type ModelRegistryResolveRequest,
type ModelRegistryResolveResponse,
type ModelRegistryVariantContract,
type ProviderDriverSpec,
type RequestedModelMatchRequest,
type RequestedModelMatchResponse,
} from '@affine/server-native';
// Owner: native/Rust contract facade.
// These types and validators intentionally proxy @affine/server-native and
// must not grow independent runtime semantics in Node.
export type {
CapabilityMatchRequest,
CapabilityMatchResponse,
ProviderDriverSpec,
RequestedModelMatchRequest,
RequestedModelMatchResponse,
};
export type CopilotModelBackendKind = ModelRegistryMatchRequest['backendKind'];
export type ModelRegistryVariant = ModelRegistryVariantContract;
export type ResolveModelRegistryVariantRequest = ModelRegistryResolveRequest;
export type ResolveModelRegistryVariantResponse = ModelRegistryResolveResponse;
export type MatchModelRegistryRequest = ModelRegistryMatchRequest;
export type MatchModelRegistryResponse = ModelRegistryMatchResponse;
function validateNativeContract<T>(name: string, value: unknown): T {
return serverNativeModule.llmValidateContract(name, value) as T;
}
export function parseCapabilityMatchRequest(value: unknown) {
return validateNativeContract<CapabilityMatchRequest>(
'capabilityMatchRequest',
value
);
}
export function parseCapabilityMatchResponse(value: unknown) {
return validateNativeContract<CapabilityMatchResponse>(
'capabilityMatchResponse',
value
);
}
export function parseResolveModelRegistryVariantRequest(value: unknown) {
return validateNativeContract<ResolveModelRegistryVariantRequest>(
'modelRegistryResolveRequest',
value
);
}
export function parseResolveModelRegistryVariantResponse(value: unknown) {
return validateNativeContract<ResolveModelRegistryVariantResponse>(
'modelRegistryResolveResponse',
value
);
}
export function parseMatchModelRegistryRequest(value: unknown) {
return validateNativeContract<MatchModelRegistryRequest>(
'modelRegistryMatchRequest',
value
);
}
export function parseMatchModelRegistryResponse(value: unknown) {
return validateNativeContract<MatchModelRegistryResponse>(
'modelRegistryMatchResponse',
value
);
}
export function parseProviderDriverSpec(value: unknown) {
return validateNativeContract<ProviderDriverSpec>(
'providerDriverSpec',
value
);
}
export function parseRequestedModelMatchRequest(value: unknown) {
return validateNativeContract<RequestedModelMatchRequest>(
'requestedModelMatchRequest',
value
);
}
export function parseRequestedModelMatchResponse(value: unknown) {
return validateNativeContract<RequestedModelMatchResponse>(
'requestedModelMatchResponse',
value
);
}
packages/backend/server/src/plugins/copilot/runtime/contracts/prompt-contract.ts
+3 -34
View File
@@ -1,15 +1,6 @@
import {
llmValidateContract,
type NativePromptCountTokensRequest,
type NativePromptCountTokensResponse,
type NativePromptMetadataRequest,
type NativePromptMetadataResponse,
type NativePromptRenderRequest,
type NativePromptRenderResponse,
type NativePromptSessionRenderRequest,
type NativePromptSessionRenderResponse,
type PromptMessageContract as NativePromptMessageContract,
type PromptStructuredResponseContract as NativePromptStructuredResponseContract,
import type {
PromptMessageContract as NativePromptMessageContract,
PromptStructuredResponseContract as NativePromptStructuredResponseContract,
} from '../../../../native';
import { normalizePromptResponseFormat } from './structured-output-contract';
@@ -32,14 +23,6 @@ type PromptMessageInput = {
params?: Record<string, unknown> | null;
responseFormat?: PromptResponseFormat | null;
};
export type PromptRenderContract = NativePromptRenderRequest;
export type PromptRenderResult = NativePromptRenderResponse;
export type PromptTokenCountContract = NativePromptCountTokensRequest;
export type PromptTokenCountResult = NativePromptCountTokensResponse;
export type PromptMetadataContract = NativePromptMetadataRequest;
export type PromptMetadataResult = NativePromptMetadataResponse;
export type PromptSessionContract = NativePromptSessionRenderRequest;
export type PromptSessionResult = NativePromptSessionRenderResponse;
export type NativePromptResponseFormatProjection = {
nativeResponseFormat?: PromptStructuredResponseContract;
};
@@ -82,17 +65,3 @@ export function projectPromptMessageForNative(
return { message: nativeMessage, nativeResponseFormat };
}
export function parsePromptRenderContract(value: unknown) {
return llmValidateContract<PromptRenderContract>(
'promptRenderContract',
value
);
}
export function parsePromptSessionContract(value: unknown) {
return llmValidateContract<PromptSessionContract>(
'promptSessionContract',
value
);
}
packages/backend/server/src/plugins/copilot/runtime/copilot-runtime-event-consumer.ts
+129
View File
@@ -0,0 +1,129 @@
import { Injectable, Logger } from '@nestjs/common';
import { metrics } from '../../../base';
import { Models } from '../../../models';
import { type ByokFeatureKind, ByokProviderSource } from '../byok/types';
export type CopilotRuntimeRouteIdentity = {
profileId: string;
source: 'server' | 'local' | 'affine_cloud';
provider: string;
model: string;
};
export type CopilotRuntimeEvent =
| { type: 'route_selected'; route: CopilotRuntimeRouteIdentity }
| {
type: 'route_failed';
route: CopilotRuntimeRouteIdentity;
errorKind: string;
}
| {
type: 'usage';
route: CopilotRuntimeRouteIdentity;
usage: {
prompt_tokens?: number;
completion_tokens?: number;
total_tokens?: number;
cached_tokens?: number;
input_tokens?: number;
output_tokens?: number;
};
};
export type CopilotRuntimeEventContext = {
workspaceId?: string;
userId?: string;
sessionId?: string;
taskId?: string;
actionId?: string;
billingUnitId?: string;
featureKind: ByokFeatureKind;
};
@Injectable()
export class CopilotRuntimeEventConsumer {
private readonly logger = new Logger(CopilotRuntimeEventConsumer.name);
constructor(private readonly models: Models) {}
async consume(
events: CopilotRuntimeEvent[],
context: CopilotRuntimeEventContext
) {
for (const event of events) {
try {
if (event.type === 'usage') {
await this.recordUsage(event, context);
} else if (event.type === 'route_failed') {
await this.recordFailure(event, context);
}
} catch (error) {
this.logger.warn(
`Failed to consume copilot runtime event: ${
error instanceof Error ? error.message : String(error)
}`
);
}
}
}
private async recordUsage(
event: Extract<CopilotRuntimeEvent, { type: 'usage' }>,
context: CopilotRuntimeEventContext
) {
if (!context.workspaceId || event.route.source === 'affine_cloud') {
return;
}
const usage = event.usage;
metrics.ai.counter('byok_usage').add(1, {
provider: event.route.provider,
source: event.route.source,
feature: context.featureKind,
});
await this.models.copilotUsage.create({
workspaceId: context.workspaceId,
userId: context.userId,
provider: event.route.provider,
providerSource:
event.route.source === 'server'
? ByokProviderSource.Server
: ByokProviderSource.Local,
featureKind: context.featureKind,
model: event.route.model,
sessionId: context.sessionId,
taskId: context.taskId,
actionId: context.actionId,
billingUnitId: context.billingUnitId,
promptTokens: usage.prompt_tokens ?? usage.input_tokens ?? 0,
completionTokens: usage.completion_tokens ?? usage.output_tokens ?? 0,
totalTokens: usage.total_tokens ?? 0,
cachedTokens: usage.cached_tokens ?? 0,
});
if (event.route.source === 'server') {
await this.models.copilotWorkspaceByokConfig.touchUsed(
context.workspaceId,
event.route.profileId
);
}
}
private async recordFailure(
event: Extract<CopilotRuntimeEvent, { type: 'route_failed' }>,
context: CopilotRuntimeEventContext
) {
metrics.ai.counter('byok_route_failure').add(1, {
provider: event.route.provider,
source: event.route.source,
feature: context.featureKind,
reason: event.errorKind,
});
if (context.workspaceId && event.route.source === 'server') {
await this.models.copilotWorkspaceByokConfig.markFailure(
context.workspaceId,
event.route.profileId,
event.errorKind
);
}
}
}

Some files were not shown because too many files have changed in this diff Show More