Jayden/AFFiNE-Mirror
0
0
Fork 0
mirror of https://github.com/toeverything/AFFiNE.git synced 2026-08-07 04:20:11 +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
.docker/selfhost/schema.json
-114
View File
@@ -1209,120 +1209,6 @@
"description": "The profile list for copilot providers.\n@default []",
"default": []
},
"providers.defaults": {
"type": "object",
"description": "The default provider ids for model output types and global fallback.\n@default {}",
"default": {}
},
"providers.openai": {
"type": "object",
"description": "The config for the openai provider.\n@default {\"apiKey\":\"\",\"baseURL\":\"https://api.openai.com/v1\"}\n@link https://github.com/openai/openai-node",
"default": {
"apiKey": "",
"baseURL": "https://api.openai.com/v1"
}
},
"providers.cloudflareWorkersAi": {
"type": "object",
"description": "The config for the Cloudflare Workers AI provider.\n@default {\"apiToken\":\"\",\"accountId\":\"\"}",
"default": {
"apiToken": "",
"accountId": ""
}
},
"providers.fal": {
"type": "object",
"description": "The config for the fal provider.\n@default {\"apiKey\":\"\"}",
"default": {
"apiKey": ""
}
},
"providers.gemini": {
"type": "object",
"description": "The config for the gemini provider.\n@default {\"apiKey\":\"\",\"baseURL\":\"https://generativelanguage.googleapis.com/v1beta\"}",
"default": {
"apiKey": "",
"baseURL": "https://generativelanguage.googleapis.com/v1beta"
}
},
"providers.geminiVertex": {
"type": "object",
"description": "The config for the google vertex provider.\n@default {}",
"properties": {
"location": {
"type": "string",
"description": "The location of the google vertex provider."
},
"project": {
"type": "string",
"description": "The project name of the google vertex provider."
},
"googleAuthOptions": {
"type": "object",
"description": "The google auth options for the google vertex provider.",
"properties": {
"credentials": {
"type": "object",
"description": "The credentials for the google vertex provider.",
"properties": {
"client_email": {
"type": "string",
"description": "The client email for the google vertex provider."
},
"private_key": {
"type": "string",
"description": "The private key for the google vertex provider."
}
}
}
}
}
},
"default": {}
},
"providers.anthropic": {
"type": "object",
"description": "The config for the anthropic provider.\n@default {\"apiKey\":\"\",\"baseURL\":\"https://api.anthropic.com/v1\"}",
"default": {
"apiKey": "",
"baseURL": "https://api.anthropic.com/v1"
}
},
"providers.anthropicVertex": {
"type": "object",
"description": "The config for the google vertex provider.\n@default {}",
"properties": {
"location": {
"type": "string",
"description": "The location of the google vertex provider."
},
"project": {
"type": "string",
"description": "The project name of the google vertex provider."
},
"googleAuthOptions": {
"type": "object",
"description": "The google auth options for the google vertex provider.",
"properties": {
"credentials": {
"type": "object",
"description": "The credentials for the google vertex provider.",
"properties": {
"client_email": {
"type": "string",
"description": "The client email for the google vertex provider."
},
"private_key": {
"type": "string",
"description": "The private key for the google vertex provider."
}
}
}
}
}
},
"default": {}
},
"unsplash": {
"type": "object",
"description": "The config for the unsplash key.\n@default {\"key\":\"\"}",
Cargo.lock
Generated
+153 -17
View File
@@ -131,8 +131,8 @@ dependencies = [
"thiserror 2.0.18",
"tokio",
"uuid",
"windows 0.61.3",
"windows-core 0.61.2",
"windows 0.62.2",
"windows-core 0.62.2",
]
[[package]]
@@ -248,7 +248,9 @@ dependencies = [
"crc32fast",
"doc_extractor",
"file-format",
"gcp_auth",
"hex",
"hkdf 0.13.0",
"hmac 0.13.0",
"homedir",
"image",
@@ -289,6 +291,7 @@ dependencies = [
"v_htmlescape",
"webpki-roots 1.0.6",
"y-octo",
"zeroize",
]
[[package]]
@@ -1242,6 +1245,7 @@ dependencies = [
"iana-time-zone",
"js-sys",
"num-traits",
"serde",
"wasm-bindgen",
"windows-link 0.2.1",
]
@@ -2919,6 +2923,33 @@ dependencies = [
"slab",
]
[[package]]
name = "gcp_auth"
version = "0.12.7"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "26d27dbcc645b60b8e7f6e2868a9d7102ece97d1bb49c1288b5321fcc67f7260"
dependencies = [
"async-trait",
"base64",
"bytes",
"chrono",
"http",
"http-body-util",
"hyper",
"hyper-rustls",
"hyper-util",
"ring",
"rustls",
"rustls-pki-types",
"serde",
"serde_json",
"thiserror 2.0.18",
"tokio",
"tracing",
"tracing-futures",
"url",
]
[[package]]
name = "generator"
version = "0.8.8"
@@ -2930,7 +2961,7 @@ dependencies = [
"libc",
"log",
"rustversion",
"windows-link 0.1.3",
"windows-link 0.2.1",
"windows-result 0.4.1",
]
@@ -3469,6 +3500,15 @@ dependencies = [
"hmac 0.12.1",
]
[[package]]
name = "hkdf"
version = "0.13.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "4aaa26c720c68b866f2c96ef5c1264b3e6f473fe5d4ce61cd44bbe913e553018"
dependencies = [
"hmac 0.13.0",
]
[[package]]
name = "hmac"
version = "0.12.1"
@@ -3614,6 +3654,7 @@ dependencies = [
"hyper",
"hyper-util",
"rustls",
"rustls-native-certs",
"tokio",
"tokio-rustls",
"tower-service",
@@ -3660,7 +3701,7 @@ dependencies = [
"js-sys",
"log",
"wasm-bindgen",
"windows-core 0.61.2",
"windows-core 0.62.2",
]
[[package]]
@@ -4712,9 +4753,9 @@ dependencies = [
[[package]]
name = "llm_adapter"
version = "0.2.11"
version = "0.2.14"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "c44c287854e9dbe2a92e14b1a41590859a47570e328ec3fbc382645f5cc06e94"
checksum = "4f4086072a8f69a2a119e187367844d542f133f1bdf96ebdf495b49b83cf4a05"
dependencies = [
"base64",
"jsonschema",
@@ -4725,13 +4766,14 @@ dependencies = [
"thiserror 2.0.18",
"ureq",
"url",
"zeroize",
]
[[package]]
name = "llm_runtime"
version = "0.2.7"
version = "0.2.9"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "85273703c62321335888c3417462b4ace530d38479bcb4c920760808a953707f"
checksum = "35d70efcecaf49ea990fb9596664cf9c146c0260e0a0ed5a643ab47673a0521a"
dependencies = [
"jsonschema",
"llm_adapter",
@@ -6069,6 +6111,26 @@ version = "0.5.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "5be167a7af36ee22fe3115051bc51f6e6c7054c9348e28deb4f49bd6f705a315"
[[package]]
name = "pin-project"
version = "1.1.13"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "2466b2336ed02bcdca6b294417127b90ec92038d1d5c4fbeac971a922e0e0924"
dependencies = [
"pin-project-internal",
]
[[package]]
name = "pin-project-internal"
version = "1.1.13"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "c96395f0a926bc13b1c17622aaddda1ecb55d49c8f1bf9777e4d877800a43f8b"
dependencies = [
"proc-macro2",
"quote",
"syn 2.0.117",
]
[[package]]
name = "pin-project-lite"
version = "0.2.17"
@@ -7826,7 +7888,7 @@ dependencies = [
"futures-util",
"generic-array",
"hex",
"hkdf",
"hkdf 0.12.4",
"hmac 0.12.1",
"itoa",
"log",
@@ -7865,7 +7927,7 @@ dependencies = [
"futures-core",
"futures-util",
"hex",
"hkdf",
"hkdf 0.12.4",
"hmac 0.12.1",
"home",
"itoa",
@@ -8887,6 +8949,16 @@ dependencies = [
"valuable",
]
[[package]]
name = "tracing-futures"
version = "0.2.5"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "97d095ae15e245a057c8e8451bab9b3ee1e1f68e9ba2b4fbc18d0ac5237835f2"
dependencies = [
"pin-project",
"tracing",
]
[[package]]
name = "tracing-log"
version = "0.2.0"
@@ -10158,7 +10230,7 @@ version = "0.1.11"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "c2a7b1c03c876122aa43f3020e6c3c3ee5c05081c9a00739faf7503aeba10d22"
dependencies = [
"windows-sys 0.48.0",
"windows-sys 0.61.2",
]
[[package]]
@@ -10183,11 +10255,23 @@ version = "0.61.3"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "9babd3a767a4c1aef6900409f85f5d53ce2544ccdfaa86dad48c91782c6d6893"
dependencies = [
"windows-collections",
"windows-collections 0.2.0",
"windows-core 0.61.2",
"windows-future",
"windows-future 0.2.1",
"windows-link 0.1.3",
"windows-numerics",
"windows-numerics 0.2.0",
]
[[package]]
name = "windows"
version = "0.62.2"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "527fadee13e0c05939a6a05d5bd6eec6cd2e3dbd648b9f8e447c6518133d8580"
dependencies = [
"windows-collections 0.3.2",
"windows-core 0.62.2",
"windows-future 0.3.2",
"windows-numerics 0.3.1",
]
[[package]]
@@ -10199,6 +10283,15 @@ dependencies = [
"windows-core 0.61.2",
]
[[package]]
name = "windows-collections"
version = "0.3.2"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "23b2d95af1a8a14a3c7367e1ed4fc9c20e0a26e79551b1454d72583c97cc6610"
dependencies = [
"windows-core 0.62.2",
]
[[package]]
name = "windows-core"
version = "0.54.0"
@@ -10222,6 +10315,19 @@ dependencies = [
"windows-strings 0.4.2",
]
[[package]]
name = "windows-core"
version = "0.62.2"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "b8e83a14d34d0623b51dce9581199302a221863196a1dde71a7663a4c2be9deb"
dependencies = [
"windows-implement",
"windows-interface",
"windows-link 0.2.1",
"windows-result 0.4.1",
"windows-strings 0.5.1",
]
[[package]]
name = "windows-future"
version = "0.2.1"
@@ -10230,7 +10336,18 @@ checksum = "fc6a41e98427b19fe4b73c550f060b59fa592d7d686537eebf9385621bfbad8e"
dependencies = [
"windows-core 0.61.2",
"windows-link 0.1.3",
"windows-threading",
"windows-threading 0.1.0",
]
[[package]]
name = "windows-future"
version = "0.3.2"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "e1d6f90251fe18a279739e78025bd6ddc52a7e22f921070ccdc67dde84c605cb"
dependencies = [
"windows-core 0.62.2",
"windows-link 0.2.1",
"windows-threading 0.2.1",
]
[[package]]
@@ -10277,6 +10394,16 @@ dependencies = [
"windows-link 0.1.3",
]
[[package]]
name = "windows-numerics"
version = "0.3.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "6e2e40844ac143cdb44aead537bbf727de9b044e107a0f1220392177d15b0f26"
dependencies = [
"windows-core 0.62.2",
"windows-link 0.2.1",
]
[[package]]
name = "windows-registry"
version = "0.6.1"
@@ -10433,6 +10560,15 @@ dependencies = [
"windows-link 0.1.3",
]
[[package]]
name = "windows-threading"
version = "0.2.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "3949bd5b99cafdf1c7ca86b43ca564028dfe27d66958f2470940f73d86d75b37"
dependencies = [
"windows-link 0.2.1",
]
[[package]]
name = "windows_aarch64_gnullvm"
version = "0.42.2"
@@ -10851,9 +10987,9 @@ dependencies = [
[[package]]
name = "zeroize"
version = "1.8.2"
version = "1.9.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "b97154e67e32c85465826e8bcc1c59429aaaf107c1e4a9e53c8d8ccd5eff88d0"
checksum = "e13c156562582aa81c60cb29407084cdb54c4164760106ab78e6c5b0858cf64e"
[[package]]
name = "zerotrie"
Cargo.toml
+4 -2
View File
@@ -38,6 +38,7 @@ resolver = "3"
dotenvy = "0.15"
file-format = { version = "0.28", features = ["reader"] }
hex = "0.4"
hkdf = "0.13"
homedir = "0.3"
image = { version = "0.25.9", default-features = false, features = [
"bmp",
@@ -101,7 +102,7 @@ resolver = "3"
url = { version = "2.5" }
uuid = "1.8"
v_htmlescape = "0.15"
windows = { version = "0.61", features = [
windows = { version = "0.62", features = [
"Win32_Devices_FunctionDiscovery",
"Win32_Foundation",
"Win32_Media_Audio",
@@ -113,8 +114,9 @@ resolver = "3"
"Win32_System_Variant",
"Win32_UI_Shell_PropertiesSystem",
] }
windows-core = { version = "0.61" }
windows-core = { version = "0.62" }
y-octo = "0.1.0"
zeroize = "1.9"
zip = "8.6"
[profile.dev.package.sqlx-macros]
packages/backend/native/Cargo.toml
+9 -1
View File
@@ -20,7 +20,9 @@ chrono = { workspace = true }
crc32fast = "1.5.0"
doc_extractor = { workspace = true }
file-format = { workspace = true }
gcp_auth = "0.12.7"
hex = { workspace = true }
hkdf = { workspace = true }
hmac = "0.13"
homedir = { workspace = true }
image = { workspace = true }
@@ -62,7 +64,12 @@ sqlx = { workspace = true, default-features = false, features = [
] }
thiserror.workspace = true
tiktoken-rs = { workspace = true }
tokio = { workspace = true, features = ["rt-multi-thread", "sync", "time"] }
tokio = { workspace = true, features = [
"net",
"rt-multi-thread",
"sync",
"time",
] }
unicode-normalization = "0.1"
unicode_skeleton = "0.1.1"
url = { workspace = true }
@@ -70,6 +77,7 @@ uuid = { workspace = true, features = ["v4"] }
v_htmlescape = { workspace = true }
webpki-roots = "1.0"
y-octo = { workspace = true, features = ["large_refs"] }
zeroize = { workspace = true }
[target.'cfg(not(target_os = "linux"))'.dependencies]
mimalloc = { workspace = true }
packages/backend/native/index.d.ts
Vendored
+238 -148
View File
@@ -8,6 +8,9 @@ export declare class BackendRuntime {
acquireCoordinationLease(key: string, owner: string, ttlMs: number): Promise<CoordinationLeaseGrant | null>
releaseCoordinationLease(key: string, owner: string, fencingToken: bigint | number): Promise<boolean>
renewCoordinationLease(key: string, owner: string, fencingToken: bigint | number, ttlMs: number): Promise<boolean>
executeCopilotStream(input: CopilotExecuteInput, maxSteps: number, callback: ((err: Error | null, arg: string) => void), toolCallback: ((err: Error | null, arg: string) => Promise<string>)): Promise<CopilotStreamHandle>
executeCopilot(input: CopilotExecuteInput): Promise<string>
assertCopilotRoute(input: CopilotRouteCheckInput): Promise<void>
/**
* Merge pending doc updates with y-octo and persist the merged snapshot.
*
@@ -51,21 +54,29 @@ export declare class BackendRuntime {
getWorkspaceInviteLink(workspaceId: string): Promise<RuntimeWorkspaceInviteLinkRecord | null>
getWorkspaceInviteLinkById(inviteId: string): Promise<RuntimeWorkspaceInviteLinkRecord | null>
revokeWorkspaceInviteLink(workspaceId: string): Promise<boolean>
createByokLocalLease(activeKey: string, leaseId: string, payload: any, ttlMs: number): Promise<RuntimeByokLocalLeaseRecord>
getByokLocalLease(leaseId: string): Promise<RuntimeByokLocalLeaseRecord | null>
cleanupExpiredRuntimeStates(limit: number): Promise<number>
refreshWorkspaceAdminStatsDirty(batchLimit: number, owner: string, leaseTtlMs: number): Promise<RuntimeWorkspaceStatsRefreshResult>
recalibrateWorkspaceAdminStats(lastSid: number, batchLimit: number, owner: string, leaseTtlMs: number): Promise<RuntimeWorkspaceStatsRecalibrationResult>
writeWorkspaceAdminStatsDailySnapshot(owner: string, leaseTtlMs: number): Promise<RuntimeWorkspaceStatsSnapshotResult>
recalibrateWorkspaceAdminStatsDaily(batchLimit: number, owner: string, leaseTtlMs: number, lockRetryTimes: number, lockRetryDelayMs: number): Promise<RuntimeWorkspaceStatsDailyRecalibrationResult>
constructor()
constructor(privateKey?: string | undefined | null)
start(): Promise<void>
stop(): Promise<void>
reloadConfig(privateKey?: string | undefined | null): Promise<void>
health(): Promise<BackendRuntimeHealth>
runMigrations(): Promise<void>
listByokProfiles(workspaceId: string): Promise<Array<ByokProfileOutput>>
createByokProfile(input: CreateByokProfileInput): Promise<ByokProfileOutput>
replaceByokProfile(input: ReplaceByokProfileInput): Promise<ByokProfileOutput>
rotateByokCredential(input: RotateByokCredentialInput): Promise<ByokProfileOutput>
probeByokProfile(input: ProbeByokProfileInput): Promise<ByokProbeResultOutput>
probeByokDraft(input: ProbeByokDraftInput): Promise<ByokProbeResultOutput>
deleteByokProfile(workspaceId: string, profileId: string): Promise<boolean>
reorderByokProfiles(input: ReorderByokProfilesInput): Promise<Array<ByokProfileOutput>>
createByokLocalLease(input: CreateByokLocalLeaseInput): Promise<ByokLocalLeaseOutput>
}
export declare class LlmStreamHandle {
export declare class CopilotStreamHandle {
abort(): void
}
@@ -107,46 +118,6 @@ export declare class Tokenizer {
count(content: string, allowedSpecial?: Array<string> | undefined | null): number
}
export interface ActionEvent {
type: ActionEventType
actionId: string
actionVersion: string
stepId?: string
status?: ActionRunStatus
attachment?: any
result?: any
errorCode?: string
errorMessage?: string
trace?: ActionTrace
}
export type ActionEventType = 'action_start'|
'step_start'|
'attachment'|
'step_end'|
'action_done'|
'error';
export type ActionRunStatus = 'created'|
'running'|
'succeeded'|
'failed'|
'aborted';
export interface ActionRuntimeInput {
recipeId: string
recipeVersion?: string
input: any
}
export interface ActionTrace {
actionId: string
actionVersion: string
status: ActionRunStatus
lightweight: Array<any>
errorCode?: string
}
export declare function activateLicense(request: LicenseKeyRequest): Promise<LicenseResponse>
/**
@@ -194,6 +165,15 @@ export interface BackendRuntimeHealth {
export declare function buildPublicRootDoc(rootDocBin: Buffer, docMetas: Array<PublicDocMetaInput>): Buffer
export interface BuiltInManagedTarget {
id: string
displayName: string
minimumTier: BuiltInManagedTargetTier
}
export type BuiltInManagedTargetTier = 'Standard'|
'Premium';
export interface BuiltInPromptRenderContract {
name: string
renderParams: Record<string, any>
@@ -203,20 +183,124 @@ export interface BuiltInPromptSessionContract {
name: string
turns: Array<PromptMessageContract>
renderParams: Record<string, any>
maxTokenSize: number
}
export interface BuiltInPromptSpec {
name: string
action?: string
model: string
optionalModels?: Array<string>
config?: any
params?: Record<string, PromptParamSpec>
builtins?: Array<PromptBuiltin>
messages: Array<PromptSpecMessage>
}
export interface BuiltInRouteOptions {
routeId: string
standardDefaultTargetId?: string
premiumDefaultTargetId?: string
choices: Array<BuiltInManagedTarget>
}
export interface ByokCapabilityInput {
input: Array<string>
output: Array<string>
features: Array<string>
attachmentKinds: Array<string>
attachmentSources: Array<string>
}
export interface ByokCatalogModelOutput {
modelId: string
displayName: string
recommended: boolean
capabilities: Array<ByokCapabilityInput>
}
export interface ByokCatalogOutput {
version: string
providers: Array<ByokCatalogProviderOutput>
}
export interface ByokCatalogProviderOutput {
provider: string
models: Array<ByokCatalogModelOutput>
}
export interface ByokEndpointInput {
kind: string
url?: string
}
export interface ByokLocalLeaseOutput {
leaseId: string
expiresAtMs: number
}
export interface ByokModelDeclarationInput {
modelId: string
enabled: boolean
capabilities: Array<ByokCapabilityInput>
}
export interface ByokModelProbeCheckOutput {
operation: string
status: ByokProbeStatusOutput
}
export interface ByokModelProbeOutput {
modelId: string
checks: Array<ByokModelProbeCheckOutput>
}
export interface ByokProbeCheckInput {
modelId: string
operation: string
}
export interface ByokProbeResultOutput {
definitionFingerprint: string
stale: boolean
connection: ByokProbeStatusOutput
models: Array<ByokModelProbeOutput>
}
export interface ByokProbeStatusOutput {
kind: string
testedAtMs?: number
errorKind?: string
}
export interface ByokProfileDefinitionInput {
version: number
endpoint: ByokEndpointInput
models: Array<ByokModelDeclarationInput>
}
export interface ByokProfileOrderInput {
profileId: string
expectedRevision: number
}
export interface ByokProfileOutput {
profileId: string
workspaceId: string
provider: string
name: string
description?: string
definition: ByokProfileDefinitionInput
enabled: boolean
sortOrder: number
revision: number
validation?: ByokValidationOutput
}
export interface ByokValidationOutput {
definitionFingerprint: string
credentialGeneration: number
connection: ByokProbeStatusOutput
models: Array<ByokModelProbeOutput>
}
export interface CanonicalChatRequestContract {
model: string
messages: Array<PromptMessageContract>
@@ -315,8 +399,74 @@ export interface CoordinationLeaseGrant {
fencingToken: bigint | number
}
export interface CopilotAccessProjection {
routeAllowed: boolean
managedTier: CopilotManagedTier
serverByok: boolean
localByok: boolean
}
export declare function copilotActionRecipe(actionId: string, actionVersion?: string | undefined | null): string
export interface CopilotExecuteInput {
slot: string
builtInRouteId?: string
workspaceId?: string
userId?: string
localLeaseId?: string
access: CopilotAccessProjection
managedTargetId?: string
targetOverride?: CopilotTargetOverrideInput
request: unknown
}
export type CopilotManagedTier = 'Standard'|
'Premium';
export interface CopilotRouteCheckInput {
slot: string
builtInRouteId?: string
workspaceId?: string
userId?: string
localLeaseId?: string
access: CopilotAccessProjection
managedTargetId?: string
targetOverride?: CopilotTargetOverrideInput
}
export interface CopilotTargetOverrideInput {
profileId: string
modelId: string
}
export declare function createAuthSessionRefreshToken(): AuthSessionRefreshToken
export interface CreateByokLocalLeaseInput {
workspaceId: string
userId: string
providers: Array<CreateByokLocalLeaseProviderInput>
}
export interface CreateByokLocalLeaseProviderInput {
provider: string
name: string
description?: string
credential: string
definition: ByokProfileDefinitionInput
enabled: boolean
}
export interface CreateByokProfileInput {
workspaceId: string
provider: string
name: string
description?: string
credential: string
definition: ByokProfileDefinitionInput
enabled: boolean
actorUserId: string
}
/**
* Converts markdown content to AFFiNE-compatible y-octo document binary.
*
@@ -409,31 +559,11 @@ export declare function llmBuildRerankRequest(request: LlmRerankRequestContract)
export declare function llmCanonicalJsonSchemaHash(schema: any): string
export declare function llmCollectPromptMetadata(request: PromptMetadataContract): PromptMetadataResult
export declare function llmCompileExecutionPlan(value: any): any
export interface LlmCoreMessage {
role: string
content: Array<any>
}
export declare function llmCountPromptTokens(request: PromptTokenCountContract): PromptTokenCountResult
export declare function llmDispatchPrepared(routesJson: string): Promise<string>
export declare function llmDispatchPreparedStream(routesJson: string, callback: ((err: Error | null, arg: string) => void)): LlmStreamHandle
export declare function llmDispatchToolLoopStream(protocol: string, backendConfigJson: string, requestJson: string, maxSteps: number, callback: ((err: Error | null, arg: string) => void), toolCallback: ((err: Error | null, arg: string) => Promise<string>)): LlmStreamHandle
export declare function llmDispatchToolLoopStreamPrepared(routesJson: string, maxSteps: number, callback: ((err: Error | null, arg: string) => void), toolCallback: ((err: Error | null, arg: string) => Promise<string>)): LlmStreamHandle
export declare function llmDispatchToolLoopStreamRouted(routesJson: string, requestJson: string, maxSteps: number, callback: ((err: Error | null, arg: string) => void), toolCallback: ((err: Error | null, arg: string) => Promise<string>)): LlmStreamHandle
export declare function llmEmbeddingDispatch(protocol: string, backendConfigJson: string, requestJson: string): Promise<string>
export declare function llmEmbeddingDispatchPrepared(routesJson: string): Promise<string>
export interface LlmEmbeddingRequestContract {
model: string
inputs: Array<string>
@@ -443,9 +573,11 @@ export interface LlmEmbeddingRequestContract {
export declare function llmGetBuiltInPromptSpec(name: string): BuiltInPromptSpec | null
export declare function llmGetContractSchema(name: string): any
export declare function llmGetBuiltInRouteOptions(name: string): BuiltInRouteOptions | null
export declare function llmImageDispatchPrepared(routesJson: string): Promise<string>
export declare function llmGetByokCatalog(): ByokCatalogOutput
export declare function llmGetContractSchema(name: string): any
export interface LlmImageInputContract {
kind: 'url' | 'data' | 'bytes'
@@ -488,7 +620,6 @@ export interface LlmImageProviderOptionsContract {
export interface LlmImageRequestBuildContract {
model: string
protocol: 'openai_chat' | 'openai_responses' | 'openai_images' | 'anthropic' | 'gemini' | 'fal_image'
messages: Array<PromptMessageContract>
options?: any
}
@@ -511,18 +642,10 @@ export declare function llmMatchModelCapabilities(payload: CapabilityMatchReques
export declare function llmMatchModelRegistry(request: ModelRegistryMatchRequest): ModelRegistryMatchResponse
export declare function llmNormalizePreparedRoutes(value: any): any
export declare function llmPlanAttachmentReference(protocol: string, backendConfigJson: string, sourceJson: string): string
export declare function llmRenderBuiltInPrompt(request: BuiltInPromptRenderContract): PromptRenderResult
export declare function llmRenderBuiltInSessionPrompt(request: BuiltInPromptSessionContract): PromptSessionResult
export declare function llmRenderPrompt(request: PromptRenderContract): PromptRenderResult
export declare function llmRenderSessionPrompt(request: PromptSessionContract): PromptSessionResult
export interface LlmRequestContract {
model: string
messages: Array<LlmCoreMessage>
@@ -537,10 +660,6 @@ export interface LlmRequestContract {
middleware?: any
}
export declare function llmRerankDispatch(protocol: string, backendConfigJson: string, requestJson: string): Promise<string>
export declare function llmRerankDispatchPrepared(routesJson: string): Promise<string>
export interface LlmRerankRequestContract {
model: string
query: string
@@ -550,14 +669,6 @@ export interface LlmRerankRequestContract {
export declare function llmResolveModelRegistryVariant(request: ModelRegistryResolveRequest): ModelRegistryResolveResponse
export declare function llmResolveRequestedModelMatch(payload: RequestedModelMatchRequest): RequestedModelMatchResponse
export declare function llmResolveRequestIntent(protocol: string, backendConfigJson: string, intentJson: string): string
export declare function llmStructuredDispatch(protocol: string, backendConfigJson: string, requestJson: string): Promise<string>
export declare function llmStructuredDispatchPrepared(routesJson: string): Promise<string>
export interface LlmStructuredRequestContract {
model: string
messages: Array<LlmCoreMessage>
@@ -695,6 +806,22 @@ export interface PortalResponse {
error?: LicenseError
}
export interface ProbeByokDraftInput {
workspaceId: string
provider: string
credential?: string
profileId?: string
expectedRevision?: number
definition: ByokProfileDefinitionInput
checks: Array<ByokProbeCheckInput>
}
export interface ProbeByokProfileInput {
workspaceId: string
profileId: string
checks: Array<ByokProbeCheckInput>
}
export declare function processImage(input: Buffer, maxEdge: number, keepExif: boolean): Promise<Buffer>
export type PromptBuiltin = 'Date'|
@@ -705,10 +832,6 @@ export type PromptBuiltin = 'Date'|
'HasSelected'|
'HasCurrentDoc';
export interface PromptCountMessage {
content: string
}
export interface PromptMessageContract {
role: 'system' | 'assistant' | 'user'
content: string
@@ -717,46 +840,16 @@ export interface PromptMessageContract {
responseFormat?: PromptStructuredResponseContract
}
export interface PromptMetadataContract {
messages: Array<PromptMessageContract>
}
export interface PromptMetadataResult {
paramKeys: Array<string>
templateParams: Record<string, any>
}
export interface PromptParamSpec {
default?: string
enumValues?: Array<string>
}
export interface PromptRenderContract {
messages: Array<PromptMessageContract>
templateParams: Record<string, any>
renderParams: Record<string, any>
}
export interface PromptRenderResult {
messages: Array<PromptMessageContract>
warnings: Array<string>
}
export interface PromptSessionContract {
prompt: PromptSessionPrompt
turns: Array<PromptMessageContract>
renderParams: Record<string, any>
maxTokenSize: number
}
export interface PromptSessionPrompt {
action?: string
model?: string
promptTokens: number
templateParams: Record<string, any>
messages: Array<PromptMessageContract>
}
export interface PromptSessionResult {
messages: Array<PromptMessageContract>
warnings: Array<string>
@@ -775,15 +868,6 @@ export interface PromptStructuredResponseContract {
strict?: boolean
}
export interface PromptTokenCountContract {
model?: string
messages: Array<PromptCountMessage>
}
export interface PromptTokenCountResult {
tokens: number
}
export interface ProviderDriverSpec {
driverId: string
providerType: string
@@ -838,16 +922,22 @@ export interface RemoteMimeTypeRequest {
timeoutMs?: number
}
export interface RequestedModelMatchRequest {
providerIds: Array<string>
optionalModels: Array<string>
requestedModelId?: string
defaultModel?: string
export interface ReorderByokProfilesInput {
workspaceId: string
profiles: Array<ByokProfileOrderInput>
actorUserId: string
}
export interface RequestedModelMatchResponse {
selectedModel?: string
matchedOptionalModel: boolean
export interface ReplaceByokProfileInput {
workspaceId: string
profileId: string
expectedRevision: number
name: string
description?: string
definition: ByokProfileDefinitionInput
credential?: string
enabled: boolean
actorUserId: string
}
export interface RerankCandidate {
@@ -896,7 +986,13 @@ export interface ResolveEntitlementInput {
export declare function resolveEntitlementV1(input: ResolveEntitlementInput): ResolvedEntitlement
export declare function runNativeActionRecipePreparedStream(input: ActionRuntimeInput, callback: ((err: Error | null, arg: string) => void)): LlmStreamHandle
export interface RotateByokCredentialInput {
workspaceId: string
profileId: string
expectedRevision: number
credential: string
actorUserId: string
}
export interface RuntimeBlobCleanupExecuteResult {
scannedCandidates: number
@@ -943,12 +1039,6 @@ export interface RuntimeBlobMetadataBackfillResult {
workspaceIds: Array<string>
}
export interface RuntimeByokLocalLeaseRecord {
leaseId: string
payload: any
expiresAtMs: number
}
export interface RuntimeDocBlobRefsResult {
scannedDocs: number
parsedDocs: number
packages/backend/native/src/llm/action/catalog.rs
+81 -274
View File
@@ -1,291 +1,98 @@
use std::collections::HashSet;
use jsonschema::Draft;
use napi::{Error, Result, Status};
use serde::Serialize;
use serde_json::{Value, json};
use super::{
super::contract_schema::{transcript_input_schema, transcript_result_schema},
ActionRecipe, ActionRecipeStep, ActionStepKind,
};
fn invalid_recipe(message: impl Into<String>) -> Error {
Error::new(Status::InvalidArg, message.into())
#[derive(Serialize)]
#[serde(rename_all = "camelCase")]
struct ActionRecipe<'a> {
action_id: &'a str,
action_version: &'a str,
slot: &'a str,
prompt_ref: &'a str,
response_contract: Value,
output_projection: &'a str,
}
pub fn built_in_recipes() -> Vec<ActionRecipe> {
vec![
action_recipe("mindmap.generate", "v1"),
action_recipe("slides.outline", "v1"),
action_recipe("image.filter.sketch", "v1"),
action_recipe("image.filter.clay", "v1"),
action_recipe("image.filter.anime", "v1"),
action_recipe("image.filter.pixel", "v1"),
transcript_recipe("transcript.audio.gemini", "v1"),
]
}
pub fn find_recipe(id: &str, version: Option<&str>) -> Result<ActionRecipe> {
let catalog = load_catalog()?;
catalog
.into_iter()
.find(|recipe| recipe.id == id && version.is_none_or(|version| recipe.version == version))
.ok_or_else(|| {
invalid_recipe(format!(
"Action recipe not found: {}{}",
id,
version.map(|version| format!("@{version}")).unwrap_or_default()
))
})
}
pub fn load_catalog() -> Result<Vec<ActionRecipe>> {
let recipes = built_in_recipes();
validate_catalog(&recipes)?;
Ok(recipes)
}
pub fn validate_catalog(recipes: &[ActionRecipe]) -> Result<()> {
let mut keys = HashSet::new();
for recipe in recipes {
validate_recipe(recipe)?;
let key = format!("{}@{}", recipe.id, recipe.version);
if !keys.insert(key.clone()) {
return Err(invalid_recipe(format!("Duplicated action recipe: {key}")));
}
#[napi_derive::napi]
pub fn copilot_action_recipe(action_id: String, action_version: Option<String>) -> Result<String> {
let version = action_version.as_deref().unwrap_or("v1");
if version != "v1" {
return Err(Error::new(Status::InvalidArg, "Action recipe not found"));
}
Ok(())
}
pub fn validate_recipe(recipe: &ActionRecipe) -> Result<()> {
if recipe.id.trim().is_empty() {
return Err(invalid_recipe("Action recipe id is required"));
}
if recipe.version.trim().is_empty() {
return Err(invalid_recipe("Action recipe version is required"));
}
if recipe.steps.is_empty() {
return Err(invalid_recipe(format!(
"Action recipe {}@{} must declare at least one step",
recipe.id, recipe.version
)));
}
compile_schema("inputSchema", &recipe.input_schema)?;
compile_schema("outputSchema", &recipe.output_schema)?;
let mut step_ids = HashSet::new();
let mut has_final = false;
for step in &recipe.steps {
if step.id.trim().is_empty() {
return Err(invalid_recipe(format!(
"Action recipe {}@{} contains a step without id",
recipe.id, recipe.version
)));
}
if !step_ids.insert(step.id.clone()) {
return Err(invalid_recipe(format!(
"Action recipe {}@{} contains duplicated step id {}",
recipe.id, recipe.version, step.id
)));
}
if step.kind == ActionStepKind::Final {
has_final = true;
}
}
if !has_final {
return Err(invalid_recipe(format!(
"Action recipe {}@{} must end with a final step",
recipe.id, recipe.version
)));
}
if recipe
.steps
.last()
.is_some_and(|step| step.kind != ActionStepKind::Final)
{
return Err(invalid_recipe(format!(
"Action recipe {}@{} must end with a final step",
recipe.id, recipe.version
)));
}
Ok(())
}
fn compile_schema(label: &str, schema: &Value) -> Result<()> {
jsonschema::options()
.with_draft(Draft::Draft7)
.build(schema)
.map(|_| ())
.map_err(|error| invalid_recipe(format!("Invalid action recipe {label}: {error}")))
}
fn action_recipe(id: &str, version: &str) -> ActionRecipe {
let steps = if id.starts_with("image.filter.") {
vec![
ActionRecipeStep {
id: "generate-image".to_string(),
kind: ActionStepKind::PromptImage,
input: Some(json!({
"preparedRoutes": { "$state": "preparedRoutes.generate-image" },
"outputKey": "artifact"
})),
state_patch: Some(json!({ "imageGenerated": true })),
let recipe = match action_id.as_str() {
"mindmap.generate" => structured(&action_id, "mindmap.generate", text_result_schema()),
"slides.outline" => structured(&action_id, "slides.outline", text_result_schema()),
"transcript.audio" => structured(
&action_id,
"Transcript audio structured",
super::super::contract_schema::transcript_result_schema(),
),
"image.filter.sketch" | "image.filter.clay" | "image.filter.anime" | "image.filter.pixel" => ActionRecipe {
action_id: &action_id,
action_version: "v1",
slot: match action_id.as_str() {
"image.filter.sketch" => "action.image.filter.sketch",
"image.filter.clay" => "action.image.filter.clay",
"image.filter.anime" => "action.image.filter.anime",
_ => "action.image.filter.pixel",
},
ActionRecipeStep {
id: "final".to_string(),
kind: ActionStepKind::Final,
input: Some(json!({
"copy": { "$state": "artifact" }
})),
state_patch: Some(json!({ "finalized": true })),
},
]
} else if id == "slides.outline" {
vec![
ActionRecipeStep {
id: "generate-structured".to_string(),
kind: ActionStepKind::PromptStructured,
input: Some(json!({
"preparedRoutes": { "$state": "preparedRoutes.generate" },
"unwrapKey": "result",
"outputKey": "generated"
})),
state_patch: Some(json!({ "generatedAt": "promptStructured" })),
},
ActionRecipeStep {
id: "validate-json".to_string(),
kind: ActionStepKind::ValidateJson,
input: Some(json!({
"value": { "$state": "generated" },
"schema": text_action_output_schema()
})),
state_patch: None,
},
ActionRecipeStep {
id: "project-outline".to_string(),
kind: ActionStepKind::Transform,
input: Some(json!({
"slidesOutlineMarkdown": { "$state": "generated" },
"outputKey": "outlineMarkdown"
})),
state_patch: Some(json!({ "projectedAt": "slidesOutlineMarkdown" })),
},
ActionRecipeStep {
id: "final".to_string(),
kind: ActionStepKind::Final,
input: Some(json!({
"copy": { "$state": "outlineMarkdown" }
})),
state_patch: Some(json!({ "finalized": true })),
},
]
} else {
vec![
ActionRecipeStep {
id: "generate-structured".to_string(),
kind: ActionStepKind::PromptStructured,
input: Some(json!({
"preparedRoutes": { "$state": "preparedRoutes.generate" },
"unwrapKey": "result",
"outputKey": "generated"
})),
state_patch: Some(json!({ "generatedAt": "promptStructured" })),
},
ActionRecipeStep {
id: "validate-json".to_string(),
kind: ActionStepKind::ValidateJson,
input: Some(json!({
"value": { "$state": "generated" },
"schema": text_action_output_schema()
})),
state_patch: None,
},
ActionRecipeStep {
id: "final".to_string(),
kind: ActionStepKind::Final,
input: Some(json!({
"copy": { "$state": "generated" }
})),
state_patch: Some(json!({ "finalized": true })),
},
]
prompt_ref: &action_id,
response_contract: Value::Null,
output_projection: "first_image",
},
_ => return Err(Error::new(Status::InvalidArg, "Action recipe not found")),
};
recipe(id, version, action_output_schema(id), steps)
serde_json::to_string(&recipe).map_err(|error| Error::new(Status::GenericFailure, error.to_string()))
}
fn transcript_recipe(id: &str, version: &str) -> ActionRecipe {
let mut recipe = recipe(
id,
version,
transcript_result_schema(),
vec![
ActionRecipeStep {
id: "transcribe".to_string(),
kind: ActionStepKind::PromptStructured,
input: Some(json!({
"preparedRoutes": { "$state": "preparedRoutes.transcribe" },
"outputKey": "transcriptResult"
})),
state_patch: Some(json!({ "transcribedAt": "promptStructured" })),
},
ActionRecipeStep {
id: "final".to_string(),
kind: ActionStepKind::Final,
input: Some(json!({
"sourceAudio": { "$state": "sourceAudio" },
"quality": { "$state": "quality" },
"infos": { "$state": "infos" },
"sliceManifest": { "$state": "sliceManifest" },
"normalizedSegments": { "$state": "transcriptResult.normalizedSegments" },
"normalizedTranscript": { "$state": "transcriptResult.normalizedTranscript" },
"summaryJson": { "$state": "transcriptResult.summaryJson" },
"providerMeta": { "$state": "transcriptResult.providerMeta" },
"version": "transcript-result-v1",
"strategy": id.strip_prefix("transcript.audio.").unwrap_or(id)
})),
state_patch: Some(json!({ "finalized": true })),
},
],
);
recipe.input_schema = transcript_input_schema();
recipe
}
fn action_output_schema(id: &str) -> Value {
if id.starts_with("image.filter.") {
json!({
"type": "object",
"properties": {
"url": { "type": "string" },
"data_base64": { "type": "string" },
"media_type": { "type": "string" }
},
"anyOf": [
{ "required": ["url"] },
{ "required": ["data_base64", "media_type"] }
],
"additionalProperties": true
})
} else {
text_action_output_schema()
}
}
fn text_action_output_schema() -> Value {
fn text_result_schema() -> Value {
json!({
"type": "string",
"minLength": 1
"type": "object",
"properties": { "result": { "type": "string", "minLength": 1 } },
"required": ["result"],
"additionalProperties": false
})
}
fn recipe(id: &str, version: &str, output_schema: Value, steps: Vec<ActionRecipeStep>) -> ActionRecipe {
fn structured<'a>(action_id: &'a str, prompt_ref: &'a str, schema: Value) -> ActionRecipe<'a> {
ActionRecipe {
id: id.to_string(),
version: version.to_string(),
input_schema: json!({}),
output_schema,
steps,
action_id,
action_version: "v1",
slot: match action_id {
"mindmap.generate" => "action.mindmap.generate",
"slides.outline" => "action.slides.outline",
"transcript.audio" => "transcript.audio",
_ => unreachable!("structured recipe action is validated by the catalog"),
},
prompt_ref,
response_contract: json!({ "schema": schema, "strict": true }),
output_projection: if action_id == "slides.outline" {
"slides_outline_markdown"
} else if action_id == "transcript.audio" {
"transcript_result"
} else {
"identity"
},
}
}
#[cfg(test)]
mod tests {
use super::copilot_action_recipe;
#[test]
fn recipes_only_expose_slot_prompt_contract_and_projection() {
for (id, slot) in [
("mindmap.generate", "action.mindmap.generate"),
("slides.outline", "action.slides.outline"),
("image.filter.sketch", "action.image.filter.sketch"),
("transcript.audio", "transcript.audio"),
] {
let recipe = copilot_action_recipe(id.to_string(), None).unwrap();
assert!(!recipe.contains("prepared"));
assert!(recipe.contains(&format!("\"slot\":\"{slot}\"")));
assert!(recipe.contains("\"promptRef\""));
assert!(recipe.contains("\"outputProjection\""));
}
}
}
packages/backend/native/src/llm/action/contract.rs
-152
View File
@@ -1,152 +1,7 @@
use napi_derive::napi;
use schemars::JsonSchema;
use serde::{Deserialize, Serialize};
use serde_json::Value;
#[derive(Clone, Debug, Deserialize, JsonSchema, PartialEq, Eq, Serialize)]
#[serde(rename_all = "camelCase")]
#[serde(deny_unknown_fields)]
pub struct ActionRecipe {
pub id: String,
pub version: String,
pub input_schema: Value,
pub output_schema: Value,
pub steps: Vec<ActionRecipeStep>,
}
#[derive(Clone, Debug, Deserialize, JsonSchema, PartialEq, Eq, Serialize)]
#[serde(rename_all = "camelCase")]
#[serde(deny_unknown_fields)]
pub struct ActionRecipeStep {
pub id: String,
pub kind: ActionStepKind,
#[serde(default)]
pub input: Option<Value>,
#[serde(default)]
pub state_patch: Option<Value>,
}
#[derive(Clone, Copy, Debug, Deserialize, JsonSchema, PartialEq, Eq, Serialize)]
#[serde(rename_all = "camelCase")]
pub enum ActionStepKind {
PromptStructured,
PromptImage,
ValidateJson,
Transform,
Final,
}
#[napi(string_enum = "snake_case")]
#[derive(Clone, Copy, Debug, Deserialize, JsonSchema, PartialEq, Eq, Serialize)]
#[serde(rename_all = "snake_case")]
pub enum ActionEventType {
ActionStart,
StepStart,
Attachment,
StepEnd,
ActionDone,
Error,
}
#[napi(object)]
#[derive(Clone, Debug, Deserialize, JsonSchema, PartialEq, Serialize)]
#[serde(rename_all = "camelCase")]
#[serde(deny_unknown_fields)]
pub struct ActionEvent {
#[serde(rename = "type")]
#[napi(js_name = "type")]
pub event_type: ActionEventType,
pub action_id: String,
pub action_version: String,
#[serde(skip_serializing_if = "Option::is_none")]
pub step_id: Option<String>,
#[serde(skip_serializing_if = "Option::is_none")]
pub status: Option<ActionRunStatus>,
#[serde(skip_serializing_if = "Option::is_none")]
pub attachment: Option<Value>,
#[serde(skip_serializing_if = "Option::is_none")]
pub result: Option<Value>,
#[serde(skip_serializing_if = "Option::is_none")]
pub error_code: Option<String>,
#[serde(skip_serializing_if = "Option::is_none")]
pub error_message: Option<String>,
#[serde(skip_serializing_if = "Option::is_none")]
pub trace: Option<ActionTrace>,
}
#[napi(string_enum = "snake_case")]
#[derive(Clone, Copy, Debug, Deserialize, JsonSchema, PartialEq, Eq, Serialize)]
#[serde(rename_all = "snake_case")]
pub enum ActionRunStatus {
Created,
Running,
Succeeded,
Failed,
Aborted,
}
#[napi(object)]
#[derive(Clone, Debug, Deserialize, JsonSchema, PartialEq, Serialize)]
#[serde(rename_all = "camelCase")]
#[serde(deny_unknown_fields)]
pub struct ActionRuntimeInput {
pub recipe_id: String,
#[serde(default)]
pub recipe_version: Option<String>,
#[serde(default)]
pub input: Value,
}
#[derive(Clone, Debug, Deserialize, JsonSchema, PartialEq, Serialize)]
#[serde(rename_all = "camelCase")]
#[serde(deny_unknown_fields)]
pub struct ActionRuntimeOutput {
pub result: Value,
pub status: ActionRunStatus,
#[serde(skip_serializing_if = "Option::is_none")]
pub error_code: Option<String>,
pub state: Value,
pub steps: Vec<ActionStepRuntimeState>,
pub trace: ActionTrace,
pub events: Vec<ActionEvent>,
}
#[derive(Clone, Debug, Deserialize, JsonSchema, PartialEq, Serialize)]
#[serde(rename_all = "camelCase")]
#[serde(deny_unknown_fields)]
pub struct ActionStepRuntimeState {
pub id: String,
pub input: Value,
#[serde(skip_serializing_if = "Option::is_none")]
pub output: Option<Value>,
#[serde(skip_serializing_if = "Option::is_none")]
pub state_patch: Option<Value>,
#[serde(skip_serializing_if = "Option::is_none")]
pub error: Option<ActionStepError>,
}
#[derive(Clone, Debug, Deserialize, JsonSchema, PartialEq, Eq, Serialize)]
#[serde(rename_all = "camelCase")]
#[serde(deny_unknown_fields)]
pub struct ActionStepError {
pub code: String,
pub message: String,
}
#[napi(object)]
#[derive(Clone, Debug, Deserialize, JsonSchema, PartialEq, Serialize)]
#[serde(rename_all = "camelCase")]
#[serde(deny_unknown_fields)]
pub struct ActionTrace {
pub action_id: String,
pub action_version: String,
pub status: ActionRunStatus,
#[serde(default)]
pub lightweight: Vec<Value>,
#[serde(skip_serializing_if = "Option::is_none")]
pub error_code: Option<String>,
}
#[derive(Clone, Debug, Deserialize, JsonSchema, PartialEq, Serialize)]
#[serde(rename_all = "camelCase")]
#[serde(deny_unknown_fields)]
@@ -159,8 +14,6 @@ pub struct TranscriptInputContract {
pub infos: Option<Vec<TranscriptAudioInfo>>,
#[serde(skip_serializing_if = "Option::is_none")]
pub slice_manifest: Option<Vec<TranscriptSliceManifestItem>>,
#[serde(skip_serializing_if = "Option::is_none")]
pub prepared_routes: Option<Value>,
}
#[derive(Clone, Debug, Deserialize, JsonSchema, PartialEq, Eq, Serialize)]
@@ -232,8 +85,6 @@ pub struct TranscriptGeneratedResult {
pub normalized_transcript: String,
#[schemars(required)]
pub summary_json: Option<MeetingSummary>,
#[schemars(required)]
pub provider_meta: Option<Value>,
}
#[derive(Clone, Debug, Deserialize, JsonSchema, PartialEq, Serialize)]
@@ -253,8 +104,5 @@ pub struct TranscriptResult {
pub normalized_transcript: String,
#[schemars(required)]
pub summary_json: Option<MeetingSummary>,
#[schemars(required)]
pub provider_meta: Option<Value>,
pub version: String,
pub strategy: String,
}
packages/backend/native/src/llm/action/mod.rs
+1 -95
View File
@@ -1,99 +1,5 @@
mod catalog;
mod contract;
mod runtime;
mod slides_outline;
use std::sync::{Arc, atomic::AtomicBool, mpsc};
#[cfg(test)]
use catalog::{load_catalog, validate_catalog, validate_recipe};
use contract::{
ActionEvent, ActionEventType, ActionRecipe, ActionRecipeStep, ActionRunStatus, ActionRuntimeInput,
ActionRuntimeOutput, ActionStepError, ActionStepKind, ActionStepRuntimeState, ActionTrace,
};
pub use catalog::copilot_action_recipe;
pub(crate) use contract::{TranscriptGeneratedResult, TranscriptInputContract, TranscriptResult};
use napi::{
Result,
threadsafe_function::{ThreadsafeFunction, ThreadsafeFunctionCallMode},
};
#[cfg(test)]
use runtime::{ACTION_ABORTED_ERROR_CODE, run_action_recipe_for_test, run_action_recipe_for_test_with_control};
use runtime::{ActionRuntimeControl, run_action_recipe_prepared_with_control};
use crate::llm::{LlmStreamHandle, STREAM_END_MARKER};
#[napi(catch_unwind)]
pub fn run_native_action_recipe_prepared_stream(
input: ActionRuntimeInput,
callback: ThreadsafeFunction<String, ()>,
) -> Result<LlmStreamHandle> {
let action_id = input.recipe_id.clone();
let action_version = input.recipe_version.clone().unwrap_or_default();
let aborted = Arc::new(AtomicBool::new(false));
let aborted_in_worker = aborted.clone();
let (event_sender, event_receiver) = mpsc::channel::<ActionEvent>();
let error_sender = event_sender.clone();
std::thread::spawn(move || {
if let Err(error) = run_action_recipe_prepared_with_control(
input,
ActionRuntimeControl {
abort_signal: Some(aborted_in_worker.clone()),
event_sender: Some(event_sender),
#[cfg(test)]
abort_after_events: None,
#[cfg(test)]
mock_output: None,
},
) {
let _ = error_sender.send(ActionEvent {
event_type: ActionEventType::Error,
action_id,
action_version,
step_id: None,
status: Some(ActionRunStatus::Failed),
attachment: None,
result: None,
error_code: Some("action_runtime_error".to_string()),
error_message: Some(error.reason.clone()),
trace: None,
});
}
});
std::thread::spawn(move || {
for event in event_receiver {
match serde_json::to_string(&event) {
Ok(event) => {
let _ = callback.call(Ok(event), ThreadsafeFunctionCallMode::NonBlocking);
}
Err(error) => {
let _ = callback.call(
Ok(
serde_json::json!({
"type": "error",
"actionId": event.action_id,
"actionVersion": event.action_version,
"errorCode": "action_event_encode_failed",
"errorMessage": error.to_string()
})
.to_string(),
),
ThreadsafeFunctionCallMode::NonBlocking,
);
break;
}
}
}
let _ = callback.call(
Ok(STREAM_END_MARKER.to_string()),
ThreadsafeFunctionCallMode::NonBlocking,
);
});
Ok(LlmStreamHandle { aborted })
}
#[cfg(test)]
mod tests;
packages/backend/native/src/llm/action/runtime.rs
-564
View File
@@ -1,564 +0,0 @@
use std::{
cell::Cell,
sync::{
Arc, Mutex,
atomic::{AtomicBool, Ordering},
mpsc::Sender,
},
time::Instant,
};
use llm_runtime::{
RecipeDefinition, RecipeRuntimeEvent, RecipeRuntimeOutput, RecipeRuntimeStatus, RecipeStepExecution,
RecipeStepExecutor, StepExecutionError, execute_transform_step, execute_validate_json_step, resolve_state_ref,
run_recipe_runtime, validate_json_schema,
};
use napi::{Error, Result, Status};
use serde_json::{Map, Value, json};
use super::{
ActionEvent, ActionEventType, ActionRecipe, ActionRunStatus, ActionRuntimeInput, ActionRuntimeOutput,
ActionStepError, ActionStepKind, ActionStepRuntimeState, ActionTrace, catalog::find_recipe,
slides_outline::project_slides_outline_markdown,
};
use crate::llm::{
LlmPreparedImageDispatchRoutePayload, dispatch_prepared_image_route_payloads, dispatch_prepared_structured_routes,
};
pub const ACTION_ABORTED_ERROR_CODE: &str = "action_aborted";
pub const ACTION_INVALID_STEP_ERROR_CODE: &str = "action_invalid_step";
#[derive(Clone, Debug, Default)]
pub struct ActionRuntimeControl {
pub abort_signal: Option<Arc<AtomicBool>>,
pub event_sender: Option<Sender<ActionEvent>>,
#[cfg(test)]
pub abort_after_events: Option<usize>,
#[cfg(test)]
pub mock_output: Option<Value>,
}
#[derive(Clone, Debug)]
pub struct ActionRuntimeState {
pub status: ActionRunStatus,
pub result: Value,
pub action_state: Value,
pub steps: Vec<ActionStepRuntimeState>,
pub events: Vec<ActionEvent>,
pub trace: ActionTrace,
pub error_code: Option<String>,
}
fn invalid_input(message: impl Into<String>) -> Error {
Error::new(Status::InvalidArg, message.into())
}
pub fn run_action_recipe_prepared_with_control(
input: ActionRuntimeInput,
control: ActionRuntimeControl,
) -> Result<ActionRuntimeOutput> {
let recipe = find_recipe(&input.recipe_id, input.recipe_version.as_deref())?;
validate_value("input", &recipe.input_schema, &input.input)?;
run_recipe(recipe, input, control)
}
#[cfg(test)]
pub(crate) fn run_action_recipe_for_test(
recipe: ActionRecipe,
input: ActionRuntimeInput,
) -> Result<ActionRuntimeOutput> {
validate_value("input", &recipe.input_schema, &input.input)?;
run_recipe(recipe, input, ActionRuntimeControl::default())
}
#[cfg(test)]
pub(crate) fn run_action_recipe_for_test_with_control(
recipe: ActionRecipe,
input: ActionRuntimeInput,
control: ActionRuntimeControl,
) -> Result<ActionRuntimeOutput> {
validate_value("input", &recipe.input_schema, &input.input)?;
run_recipe(recipe, input, control)
}
fn run_recipe(
recipe: ActionRecipe,
input: ActionRuntimeInput,
control: ActionRuntimeControl,
) -> Result<ActionRuntimeOutput> {
let mut runtime = Runtime::new(recipe, input, control);
runtime.run()
}
struct Runtime {
recipe: ActionRecipe,
state: ActionRuntimeState,
started_at: Instant,
control: ActionRuntimeControl,
}
impl Runtime {
fn new(recipe: ActionRecipe, input: ActionRuntimeInput, control: ActionRuntimeControl) -> Self {
let trace = ActionTrace {
action_id: recipe.id.clone(),
action_version: recipe.version.clone(),
status: ActionRunStatus::Created,
lightweight: Vec::new(),
error_code: None,
};
Self {
recipe,
state: ActionRuntimeState {
status: ActionRunStatus::Created,
result: input.input.clone(),
action_state: input.input,
steps: Vec::new(),
events: Vec::new(),
trace,
error_code: None,
},
started_at: Instant::now(),
control,
}
}
fn run(&mut self) -> Result<ActionRuntimeOutput> {
let recipe = self.recipe_definition();
let action_id = self.recipe.id.clone();
let action_version = self.recipe.version.clone();
let output_schema = self.recipe.output_schema.clone();
let step_patches = self
.recipe
.steps
.iter()
.map(|step| (step.id.clone(), step.state_patch.clone()))
.collect::<std::collections::HashMap<_, _>>();
let attachments = Arc::new(Mutex::new(Vec::new()));
let mut executor = AffineActionStepExecutor::new(&self.control, attachments.clone());
let mut events = Vec::new();
let mut lightweight = Vec::new();
let event_sender = self.control.event_sender.clone();
let abort_signal = self.control.abort_signal.clone();
let event_count = Cell::new(0usize);
#[cfg(test)]
let abort_after_events = self.control.abort_after_events;
let mut record = |event: ActionEvent| {
lightweight.push(json!({
"type": event.event_type,
"stepId": event.step_id,
"status": event.status
}));
if let Some(sender) = &event_sender {
let _ = sender.send(event.clone());
}
events.push(event);
event_count.set(events.len());
};
let runtime_output = run_recipe_runtime(
recipe,
self.state.action_state.clone(),
&mut executor,
|event| {
for action_event in map_recipe_event(&action_id, &action_version, event, &attachments) {
record(action_event);
}
},
|| {
abort_signal
.as_ref()
.is_some_and(|signal| signal.load(Ordering::SeqCst))
|| {
#[cfg(test)]
{
abort_after_events.is_some_and(|max_events| event_count.get() >= max_events)
}
#[cfg(not(test))]
{
false
}
}
},
);
if matches!(runtime_output.status, RecipeRuntimeStatus::Succeeded) {
validate_value("output", &output_schema, &runtime_output.result)?;
}
self.state = self.action_state_from_runtime_output(runtime_output, events, lightweight, step_patches);
self.finalize_trace();
if let Some(event) = self
.state
.events
.iter_mut()
.rev()
.find(|event| matches!(event.event_type, ActionEventType::ActionDone))
{
event.trace = Some(self.state.trace.clone());
}
Ok(self.output())
}
fn recipe_definition(&self) -> RecipeDefinition {
RecipeDefinition {
id: self.recipe.id.clone(),
version: self.recipe.version.clone(),
steps: self
.recipe
.steps
.iter()
.map(|step| RecipeStepExecution {
id: step.id.clone(),
kind: action_step_kind_name(step.kind).to_string(),
input: step.input.clone(),
state_patch: step.state_patch.clone(),
})
.collect(),
}
}
fn action_state_from_runtime_output(
&self,
output: RecipeRuntimeOutput,
events: Vec<ActionEvent>,
lightweight: Vec<Value>,
step_patches: std::collections::HashMap<String, Option<Value>>,
) -> ActionRuntimeState {
let status = recipe_status_to_action_status(&output.status);
let error_code = output
.trace
.error_code
.as_deref()
.map(map_recipe_error_code)
.map(ToString::to_string);
ActionRuntimeState {
status,
result: output.result,
action_state: output.state,
steps: output
.steps
.into_iter()
.map(|step| ActionStepRuntimeState {
id: step.id.clone(),
input: step.input.unwrap_or(Value::Null),
output: step.output,
state_patch: step_patches.get(&step.id).cloned().flatten(),
error: step.error.map(ActionStepError::from),
})
.collect(),
events,
trace: ActionTrace {
action_id: self.recipe.id.clone(),
action_version: self.recipe.version.clone(),
status,
lightweight,
error_code: error_code.clone(),
},
error_code,
}
}
fn output(&mut self) -> ActionRuntimeOutput {
self.finalize_trace();
ActionRuntimeOutput {
result: self.state.result.clone(),
status: self.state.status,
error_code: self.state.error_code.clone(),
state: self.state.action_state.clone(),
steps: self.state.steps.clone(),
trace: self.state.trace.clone(),
events: self.state.events.clone(),
}
}
fn finalize_trace(&mut self) {
self.state.trace.status = self.state.status;
if self
.state
.trace
.lightweight
.last()
.and_then(|event| event.get("type"))
.is_some_and(|event_type| event_type == "action_trace")
{
return;
}
self.state.trace.lightweight.push(json!({
"type": "action_trace",
"actionId": self.recipe.id.clone(),
"actionVersion": self.recipe.version.clone(),
"status": self.state.status,
"durationMs": self.started_at.elapsed().as_millis()
}));
}
}
fn recipe_status_to_action_status(status: &RecipeRuntimeStatus) -> ActionRunStatus {
match status {
RecipeRuntimeStatus::Created => ActionRunStatus::Created,
RecipeRuntimeStatus::Running => ActionRunStatus::Running,
RecipeRuntimeStatus::Succeeded => ActionRunStatus::Succeeded,
RecipeRuntimeStatus::Failed => ActionRunStatus::Failed,
RecipeRuntimeStatus::Aborted => ActionRunStatus::Aborted,
}
}
fn map_recipe_error_code(code: &str) -> &str {
match code {
"aborted" => ACTION_ABORTED_ERROR_CODE,
"invalid_step" | "invalid_schema" | "invalid_value" => ACTION_INVALID_STEP_ERROR_CODE,
other => other,
}
}
fn map_recipe_event(
action_id: &str,
action_version: &str,
event: &RecipeRuntimeEvent,
attachments: &Arc<Mutex<Vec<Value>>>,
) -> Vec<ActionEvent> {
let status = recipe_status_to_action_status(&event.status);
let mut events = Vec::new();
if event.event_type == "step_end" {
let mut pending = attachments.lock().expect("attachment queue lock");
events.extend(pending.drain(..).map(|attachment| ActionEvent {
event_type: ActionEventType::Attachment,
action_id: action_id.to_string(),
action_version: action_version.to_string(),
step_id: None,
status: Some(ActionRunStatus::Running),
attachment: Some(attachment),
result: None,
error_code: None,
error_message: None,
trace: None,
}));
}
let event_type = match event.event_type.as_str() {
"recipe_start" => ActionEventType::ActionStart,
"step_start" => ActionEventType::StepStart,
"step_end" => ActionEventType::StepEnd,
"recipe_done" => ActionEventType::ActionDone,
"error" => ActionEventType::Error,
_ => return events,
};
let error = event.error.as_ref();
events.push(ActionEvent {
event_type,
action_id: action_id.to_string(),
action_version: action_version.to_string(),
step_id: event.step_id.clone(),
status: Some(status),
attachment: None,
result: event.result.clone(),
error_code: error.map(|error| map_recipe_error_code(&error.code).to_string()),
error_message: error.map(|error| error.message.clone()),
trace: None,
});
events
}
impl From<StepExecutionError> for ActionStepError {
fn from(error: StepExecutionError) -> Self {
let code = if error.code == "invalid_step" || error.code == "invalid_schema" || error.code == "invalid_value" {
ACTION_INVALID_STEP_ERROR_CODE.to_string()
} else {
error.code
};
Self {
code,
message: error.message,
}
}
}
fn action_step_kind_name(kind: ActionStepKind) -> &'static str {
match kind {
ActionStepKind::PromptStructured => "promptStructured",
ActionStepKind::PromptImage => "promptImage",
ActionStepKind::ValidateJson => "validateJson",
ActionStepKind::Transform => "transform",
ActionStepKind::Final => "final",
}
}
struct AffineActionStepExecutor<'a> {
#[cfg(test)]
control: &'a ActionRuntimeControl,
#[cfg(not(test))]
_marker: std::marker::PhantomData<&'a ()>,
attachments: Arc<Mutex<Vec<Value>>>,
}
impl<'a> AffineActionStepExecutor<'a> {
fn new(_control: &'a ActionRuntimeControl, attachments: Arc<Mutex<Vec<Value>>>) -> Self {
Self {
#[cfg(test)]
control: _control,
#[cfg(not(test))]
_marker: std::marker::PhantomData,
attachments,
}
}
fn test_mock_output(&self, _step_id: &str) -> Option<&Value> {
#[cfg(test)]
{
self
.control
.mock_output
.as_ref()
.and_then(|mock_output| mock_output.get(_step_id))
.filter(|value| !value.is_null())
}
#[cfg(not(test))]
{
None
}
}
fn prompt_structured_step(
&self,
step: &RecipeStepExecution,
input: Option<Value>,
) -> std::result::Result<Value, StepExecutionError> {
let value = if let Some(routes) = input
.as_ref()
.and_then(|input| input.get("preparedRoutes"))
.filter(|routes| !routes.is_null())
{
let (_provider_id, response) =
dispatch_prepared_structured_routes(&serde_json::to_string(routes).map_err(|error| {
StepExecutionError::new(
"invalid_step",
format!("Invalid promptStructured prepared routes: {error}"),
)
})?)
.map_err(|error| StepExecutionError::new("invalid_step", error.reason.clone()))?;
response.output_json.unwrap_or(Value::Null)
} else if let Some(mock_output) = self.test_mock_output(&step.id) {
mock_output.clone()
} else {
return Err(StepExecutionError::new(
"invalid_step",
"promptStructured requires preparedRoutes",
));
};
Ok(
input
.as_ref()
.and_then(|input| input.get("unwrapKey"))
.and_then(Value::as_str)
.and_then(|key| value.get(key).cloned())
.unwrap_or(value),
)
}
fn prompt_image_step(
&mut self,
step: &RecipeStepExecution,
input: Option<Value>,
) -> std::result::Result<Value, StepExecutionError> {
let attachment = if let Some(routes) = input
.as_ref()
.and_then(|input| input.get("preparedRoutes"))
.filter(|routes| !routes.is_null())
{
let payload =
serde_json::from_value::<Vec<LlmPreparedImageDispatchRoutePayload>>(routes.clone()).map_err(|error| {
StepExecutionError::new("invalid_step", format!("Invalid promptImage prepared routes: {error}"))
})?;
let (_provider_id, response) = dispatch_prepared_image_route_payloads(payload)
.map_err(|error| StepExecutionError::new("invalid_step", error.reason.clone()))?;
image_response_attachment(response.provider_metadata, response.images)
.ok_or_else(|| StepExecutionError::new("invalid_step", "promptImage native dispatch produced no image"))?
} else if let Some(mock_output) = self.test_mock_output(&step.id) {
mock_output.clone()
} else {
return Err(StepExecutionError::new(
"invalid_step",
"promptImage requires preparedRoutes",
));
};
self
.attachments
.lock()
.expect("attachment queue lock")
.push(attachment.clone());
Ok(attachment)
}
fn transform_step(&self, input: Option<Value>, state: &Value) -> std::result::Result<Value, StepExecutionError> {
if let Some(value) = execute_transform_step(input.clone(), state)? {
return Ok(value);
}
let Some(input) = input else {
return Ok(state.clone());
};
if let Some(slides_outline) = input.get("slidesOutlineMarkdown") {
let value = resolve_state_ref(slides_outline, state);
return project_slides_outline_markdown(&value)
.map(Value::String)
.map_err(|message| StepExecutionError::new("invalid_step", message));
}
Ok(input)
}
}
impl RecipeStepExecutor for AffineActionStepExecutor<'_> {
fn execute_step(
&mut self,
step: &RecipeStepExecution,
input: Option<Value>,
state: &Value,
) -> std::result::Result<Value, StepExecutionError> {
match step.kind.as_str() {
"promptStructured" => self.prompt_structured_step(step, input),
"promptImage" => self.prompt_image_step(step, input),
"validateJson" => execute_validate_json_step(input.or_else(|| Some(state.clone()))),
"transform" | "final" => self.transform_step(input, state),
other => Err(StepExecutionError::new(
"invalid_step",
format!("Unsupported action step kind: {other}"),
)),
}
}
}
fn image_response_attachment(provider_metadata: Value, images: Vec<llm_adapter::core::ImageArtifact>) -> Option<Value> {
let image = images.into_iter().next()?;
let mut attachment = Map::new();
if let Some(url) = image.url {
attachment.insert("url".to_string(), Value::String(url));
}
if let Some(data_base64) = image.data_base64 {
attachment.insert("data_base64".to_string(), Value::String(data_base64));
}
attachment.insert("media_type".to_string(), Value::String(image.media_type));
if let Some(width) = image.width {
attachment.insert("width".to_string(), json!(width));
}
if let Some(height) = image.height {
attachment.insert("height".to_string(), json!(height));
}
if !image.provider_metadata.is_null() {
attachment.insert("providerMetadata".to_string(), image.provider_metadata);
} else if !provider_metadata.is_null() {
attachment.insert("providerMetadata".to_string(), provider_metadata);
}
if !attachment.contains_key("url") && !attachment.contains_key("data_base64") {
return None;
}
Some(Value::Object(attachment))
}
fn validate_value(label: &str, schema: &Value, value: &Value) -> Result<()> {
validate_json_schema(label, schema, value).map_err(|error| invalid_input(error.message))
}
packages/backend/native/src/llm/action/slides_outline.rs
-240
View File
@@ -1,240 +0,0 @@
use serde_json::{Map, Value};
pub(super) fn project_slides_outline_markdown(value: &Value) -> Result<String, String> {
let text = match value {
Value::String(text) => text.as_str(),
Value::Object(object) => {
if let Some(Value::String(text)) = object.get("result") {
text
} else if let Some(Value::String(text)) = object.get("content") {
text
} else if let Some(Value::String(text)) = object.get("text") {
text
} else {
return Err("slidesOutlineMarkdown requires a string result".to_string());
}
}
_ => return Err("slidesOutlineMarkdown requires a string result".to_string()),
};
if is_markdown_list(text) {
return Ok(text.to_string());
}
let mut projected = Vec::new();
for line in text.lines().filter(|line| !line.trim().is_empty()) {
let item = serde_json::from_str::<Value>(line)
.map_err(|_| "slidesOutlineMarkdown requires markdown or NDJSON object lines".to_string())?;
if !item.is_object() {
return Err("slidesOutlineMarkdown requires markdown or NDJSON object lines".to_string());
}
projected.push(render_slide_item(&item)?);
}
if projected.is_empty() {
Err("slidesOutlineMarkdown requires markdown or NDJSON object lines".to_string())
} else {
Ok(projected.join("\n"))
}
}
fn is_markdown_list(text: &str) -> bool {
let mut saw_line = false;
for line in text.lines().map(str::trim_start).filter(|line| !line.trim().is_empty()) {
saw_line = true;
if !(line.starts_with("- ") || line.starts_with("* ") || line.starts_with("+ ")) {
return false;
}
}
saw_line
}
fn render_legacy_slide_item(item: &Value) -> Option<String> {
let kind = item.get("type").and_then(Value::as_str)?;
let content = item.get("content").and_then(value_to_optional_string)?;
if content.is_empty() {
return None;
}
match kind {
"name" => Some(format!("- {content}")),
"title" => Some(format!(" - {content}")),
"content" => {
if content.contains('\n') {
Some(
content
.lines()
.map(|line| format!(" - {line}"))
.collect::<Vec<_>>()
.join("\n"),
)
} else {
Some(format!(" - {content}"))
}
}
_ => None,
}
}
fn render_slide_item(item: &Value) -> Result<String, String> {
if let Some(markdown) = render_legacy_slide_item(item) {
return Ok(markdown);
}
if item.get("content").and_then(Value::as_object).is_some() {
return render_structured_slide_item(item);
}
if item.get("content").and_then(Value::as_str).is_some() {
return render_labeled_string_slide_item(item);
}
Err("slidesOutlineMarkdown item is not a recognized slide outline object".to_string())
}
fn render_labeled_string_slide_item(item: &Value) -> Result<String, String> {
let content = item
.get("content")
.and_then(Value::as_str)
.ok_or_else(|| "slidesOutlineMarkdown labeled item requires string content".to_string())?;
if content.trim().is_empty() {
return Err("slidesOutlineMarkdown labeled item requires string content".to_string());
}
let labels = parse_labeled_segments(content);
let title = labels
.get("title")
.cloned()
.filter(|value| !value.is_empty())
.ok_or_else(|| "slidesOutlineMarkdown labeled item requires Title".to_string())?;
let keywords = labels
.get("image keywords")
.cloned()
.or_else(|| labels.get("keywords").cloned())
.filter(|value| !value.is_empty())
.ok_or_else(|| "slidesOutlineMarkdown labeled item requires Image Keywords".to_string())?;
let description = labels
.get("description")
.cloned()
.or_else(|| labels.get("content").cloned())
.filter(|value| !value.is_empty())
.ok_or_else(|| "slidesOutlineMarkdown labeled item requires Description".to_string())?;
Ok(
[
format!("- {title}"),
format!(" - {title}"),
format!(" - {keywords}"),
format!(" - {description}"),
]
.join("\n"),
)
}
fn render_structured_slide_item(item: &Value) -> Result<String, String> {
let item_object = item
.as_object()
.ok_or_else(|| "slidesOutlineMarkdown structured item requires object content".to_string())?;
let content = item
.get("content")
.and_then(Value::as_object)
.ok_or_else(|| "slidesOutlineMarkdown structured item requires object content".to_string())?;
let title = string_prop(content, &["title", "name", "page_name", "pageName"])
.or_else(|| string_prop(item_object, &["title", "name", "page_name", "pageName", "page"]))
.filter(|value| !value.is_empty())
.ok_or_else(|| "slidesOutlineMarkdown requires slide title".to_string())?;
let sections = content.get("sections").and_then(Value::as_array);
let rendered_sections = if let Some(sections) = sections.filter(|sections| !sections.is_empty()) {
sections
.iter()
.enumerate()
.map(|(index, section)| render_slide_section(section, index + 1))
.collect::<Result<Vec<_>, _>>()?
.into_iter()
.flatten()
.collect::<Vec<_>>()
} else {
render_slide_object(content)?
};
Ok(
std::iter::once(format!("- {title}"))
.chain(rendered_sections)
.collect::<Vec<_>>()
.join("\n"),
)
}
fn parse_labeled_segments(text: &str) -> std::collections::HashMap<String, String> {
text
.split(';')
.filter_map(|segment| {
let (key, value) = segment.split_once(':')?;
let key = key.trim().to_ascii_lowercase();
let value = value.trim().to_string();
if key.is_empty() || value.is_empty() {
None
} else {
Some((key, value))
}
})
.collect()
}
fn render_slide_section(section: &Value, index: usize) -> Result<Vec<String>, String> {
let Some(object) = section.as_object() else {
return Err(format!("slidesOutlineMarkdown section {index} requires object content"));
};
render_slide_object(object)
}
fn render_slide_object(object: &Map<String, Value>) -> Result<Vec<String>, String> {
let title = required_string_prop(
object,
&["title", "name", "section", "page_name", "pageName"],
"slide section title",
)?;
let keywords = string_prop(
object,
&["image_keywords", "imageKeywords", "keywords", "image_keywords_optional"],
)
.filter(|value| !value.is_empty())
.unwrap_or_else(|| title.clone());
let content = required_string_prop(
object,
&["content", "description", "summary", "text"],
"slide section content",
)?;
Ok(vec![
format!(" - {title}"),
format!(" - {keywords}"),
format!(" - {content}"),
])
}
fn string_prop(object: &Map<String, Value>, keys: &[&str]) -> Option<String> {
keys
.iter()
.find_map(|key| object.get(*key).and_then(value_to_optional_string))
}
fn required_string_prop(object: &Map<String, Value>, keys: &[&str], name: &str) -> Result<String, String> {
string_prop(object, keys)
.filter(|value| !value.is_empty())
.ok_or_else(|| format!("slidesOutlineMarkdown requires {name}"))
}
fn value_to_optional_string(value: &Value) -> Option<String> {
match value {
Value::String(text) => Some(text.clone()),
Value::Number(number) => Some(number.to_string()),
Value::Array(items) => {
let joined = items
.iter()
.filter_map(value_to_optional_string)
.filter(|value| !value.is_empty())
.collect::<Vec<_>>()
.join(", ");
Some(joined)
}
_ => None,
}
}
packages/backend/native/src/llm/action/tests.rs
-854
View File
@@ -1,854 +0,0 @@
use napi::Status;
use serde_json::json;
use super::{
ACTION_ABORTED_ERROR_CODE, ActionEventType, ActionRecipe, ActionRecipeStep, ActionRunStatus, ActionRuntimeControl,
ActionRuntimeInput, ActionStepKind, load_catalog, run_action_recipe_for_test,
run_action_recipe_for_test_with_control, run_action_recipe_prepared_with_control, validate_catalog, validate_recipe,
};
#[test]
fn validates_built_in_recipe_catalog() {
let catalog = load_catalog().unwrap();
let mindmap = catalog.iter().find(|recipe| recipe.id == "mindmap.generate").unwrap();
assert!(
mindmap
.steps
.iter()
.any(|step| step.kind == ActionStepKind::PromptStructured)
);
assert!(
mindmap
.steps
.iter()
.any(|step| step.kind == ActionStepKind::ValidateJson)
);
let slides = catalog.iter().find(|recipe| recipe.id == "slides.outline").unwrap();
assert!(
slides
.steps
.iter()
.any(|step| step.id == "project-outline" && step.kind == ActionStepKind::Transform)
);
assert!(catalog.iter().any(|recipe| recipe.id == "transcript.audio.gemini"));
assert!(!catalog.iter().any(|recipe| recipe.id == "transcript.audio.local-asr"));
}
#[test]
fn built_in_transcript_action_final_result_is_schema_checked() {
let output = run_action_recipe_prepared_with_control(
ActionRuntimeInput {
recipe_id: "transcript.audio.gemini".to_string(),
recipe_version: Some("v1".to_string()),
input: json!({
"sourceAudio": { "blobId": "blob-1", "mimeType": "audio/opus" },
"quality": null,
"infos": [{ "url": "https://example.com/audio.opus", "mimeType": "audio/opus", "index": 0 }],
"sliceManifest": [{
"index": 0,
"fileName": "audio.opus",
"mimeType": "audio/opus",
"startSec": 12,
"durationSec": 30,
"byteSize": 42
}],
}),
},
mock_control(json!({
"transcribe": {
"normalizedTranscript": "00:00:01 A: Hello",
"summaryJson": {
"title": "Sync",
"durationMinutes": 1,
"attendees": ["A"],
"keyPoints": ["Hello"],
"actionItems": [],
"decisions": [],
"openQuestions": [],
"blockers": []
},
"providerMeta": { "provider": "gemini", "model": "gemini-3.5-flash-lite" }
}
})),
)
.unwrap();
assert_eq!(output.status, ActionRunStatus::Succeeded);
assert_eq!(output.result["version"], json!("transcript-result-v1"));
assert_eq!(output.result["strategy"], json!("gemini"));
assert_eq!(output.result["normalizedSegments"], json!(null));
assert_eq!(output.result["sourceAudio"]["blobId"], json!("blob-1"));
assert_eq!(
output.result["infos"][0]["url"],
json!("https://example.com/audio.opus")
);
assert_eq!(output.result["sliceManifest"][0]["startSec"], json!(12));
}
#[test]
fn built_in_transcript_action_rejects_malformed_summary() {
let error = run_action_recipe_prepared_with_control(
ActionRuntimeInput {
recipe_id: "transcript.audio.gemini".to_string(),
recipe_version: Some("v1".to_string()),
input: json!({}),
},
mock_control(json!({
"transcribe": {
"normalizedTranscript": "00:00:01 A: Hello",
"summaryJson": { "title": "Sync" },
"providerMeta": { "provider": "gemini", "model": "gemini-3.5-flash-lite" }
}
})),
)
.unwrap_err();
assert!(error.reason.contains("does not match JSON schema"));
}
#[test]
fn built_in_action_final_result_comes_from_prompt_output_state() {
let output = run_action_recipe_prepared_with_control(
ActionRuntimeInput {
recipe_id: "mindmap.generate".to_string(),
recipe_version: Some("v1".to_string()),
input: json!({}),
},
mock_control(json!({
"generate-structured": {
"result": "- Root"
}
})),
)
.unwrap();
assert_eq!(output.status, ActionRunStatus::Succeeded);
assert_eq!(output.result, json!("- Root"));
assert_eq!(output.state["generated"], json!("- Root"));
}
#[test]
fn built_in_action_unwraps_structured_text_result() {
let output = run_action_recipe_prepared_with_control(
ActionRuntimeInput {
recipe_id: "mindmap.generate".to_string(),
recipe_version: Some("v1".to_string()),
input: json!({}),
},
mock_control(json!({
"generate-structured": {
"result": "- Root"
}
})),
)
.unwrap();
assert_eq!(output.status, ActionRunStatus::Succeeded);
assert_eq!(output.result, json!("- Root"));
assert_eq!(output.state["generated"], json!("- Root"));
}
#[test]
fn built_in_slides_outline_projects_final_result_to_markdown() {
let outline = [
serde_json::to_string(&json!({
"page": "Cover",
"type": "cover",
"content": {
"title": "Apple Inc.",
"description": "Company overview",
"image_keywords": ["Apple logo", "Apple Park"]
}
}))
.unwrap(),
serde_json::to_string(&json!({
"page": 2,
"type": "content",
"content": {
"title": "Products",
"sections": [{
"title": "iPhone",
"keywords": ["smartphone", "iOS"],
"content": "Flagship product line"
}]
}
}))
.unwrap(),
serde_json::to_string(&json!({
"page": 3,
"type": "cover",
"content": "Page Name: Closing; Title: Outlook; Description: Future strategy; Image Keywords: roadmap, devices"
}))
.unwrap(),
]
.join("\n");
let output = run_action_recipe_prepared_with_control(
ActionRuntimeInput {
recipe_id: "slides.outline".to_string(),
recipe_version: Some("v1".to_string()),
input: json!({}),
},
mock_control(json!({
"generate-structured": {
"result": outline
}
})),
)
.unwrap();
assert_eq!(output.status, ActionRunStatus::Succeeded);
assert_eq!(
output.result,
json!(
[
"- Apple Inc.",
" - Apple Inc.",
" - Apple logo, Apple Park",
" - Company overview",
"- Products",
" - iPhone",
" - smartphone, iOS",
" - Flagship product line",
"- Outlook",
" - Outlook",
" - roadmap, devices",
" - Future strategy"
]
.join("\n")
)
);
assert_eq!(
output
.steps
.iter()
.find(|step| step.id == "project-outline")
.and_then(|step| step.output.as_ref()),
Some(&output.result)
);
}
#[test]
fn slides_outline_transform_keeps_legacy_markdown_shape() {
let outline = [
serde_json::to_string(&json!({ "page": 1, "type": "name", "content": "Launch deck" })).unwrap(),
serde_json::to_string(&json!({ "page": 1, "type": "title", "content": "Context" })).unwrap(),
serde_json::to_string(&json!({ "page": 1, "type": "content", "content": "Problem\nOpportunity" })).unwrap(),
]
.join("\n");
let recipe = test_recipe(vec![
ActionRecipeStep {
id: "project-outline".to_string(),
kind: ActionStepKind::Transform,
input: Some(json!({
"slidesOutlineMarkdown": { "$state": "outline" },
"outputKey": "outlineMarkdown"
})),
state_patch: None,
},
ActionRecipeStep {
id: "final".to_string(),
kind: ActionStepKind::Final,
input: Some(json!({ "copy": { "$state": "outlineMarkdown" } })),
state_patch: None,
},
]);
let output = run_action_recipe_for_test(
recipe,
runtime_input(json!({
"outline": outline
})),
)
.unwrap();
assert_eq!(output.status, ActionRunStatus::Succeeded);
assert_eq!(
output.result,
json!(["- Launch deck", " - Context", " - Problem", " - Opportunity"].join("\n"))
);
}
#[test]
fn slides_outline_transform_rejects_unrecognized_text() {
let recipe = test_recipe(vec![
ActionRecipeStep {
id: "project-outline".to_string(),
kind: ActionStepKind::Transform,
input: Some(json!({
"slidesOutlineMarkdown": { "$state": "outline" },
"outputKey": "outlineMarkdown"
})),
state_patch: None,
},
ActionRecipeStep {
id: "final".to_string(),
kind: ActionStepKind::Final,
input: Some(json!({ "copy": { "$state": "outlineMarkdown" } })),
state_patch: None,
},
]);
let output = run_action_recipe_for_test(
recipe,
runtime_input(json!({
"outline": "not valid ndjson"
})),
)
.unwrap();
assert_eq!(output.status, ActionRunStatus::Failed);
assert_eq!(output.error_code, Some("action_invalid_step".to_string()));
assert_eq!(
output.events.last().and_then(|event| event.error_message.as_deref()),
Some("slidesOutlineMarkdown requires markdown or NDJSON object lines")
);
}
#[test]
fn slides_outline_transform_accepts_cover_without_image_keywords() {
let outline = serde_json::to_string(&json!({
"page": 1,
"type": "cover",
"content": {
"title": "Launch deck",
"description": "Overview"
}
}))
.unwrap();
let recipe = test_recipe(vec![
ActionRecipeStep {
id: "project-outline".to_string(),
kind: ActionStepKind::Transform,
input: Some(json!({
"slidesOutlineMarkdown": { "$state": "outline" },
"outputKey": "outlineMarkdown"
})),
state_patch: None,
},
ActionRecipeStep {
id: "final".to_string(),
kind: ActionStepKind::Final,
input: Some(json!({ "copy": { "$state": "outlineMarkdown" } })),
state_patch: None,
},
]);
let output = run_action_recipe_for_test(
recipe,
runtime_input(json!({
"outline": outline
})),
)
.unwrap();
assert_eq!(
output.result,
json!(
[
"- Launch deck",
" - Launch deck",
" - Launch deck",
" - Overview"
]
.join("\n")
)
);
}
#[test]
fn slides_outline_transform_accepts_page_name_from_item() {
let outline = serde_json::to_string(&json!({
"page": 2,
"type": "content",
"page_name": "Workspace Benefits",
"content": {
"sections": [
{
"section": "Unified writing",
"keywords": ["docs", "canvas"],
"text": "AFFiNE combines documents and whiteboards."
}
]
}
}))
.unwrap();
let recipe = test_recipe(vec![
ActionRecipeStep {
id: "project-outline".to_string(),
kind: ActionStepKind::Transform,
input: Some(json!({
"slidesOutlineMarkdown": { "$state": "outline" },
"outputKey": "outlineMarkdown"
})),
state_patch: None,
},
ActionRecipeStep {
id: "final".to_string(),
kind: ActionStepKind::Final,
input: Some(json!({ "copy": { "$state": "outlineMarkdown" } })),
state_patch: None,
},
]);
let output = run_action_recipe_for_test(
recipe,
runtime_input(json!({
"outline": outline
})),
)
.unwrap();
assert_eq!(output.status, ActionRunStatus::Succeeded);
assert_eq!(
output.result,
json!(
[
"- Workspace Benefits",
" - Unified writing",
" - docs, canvas",
" - AFFiNE combines documents and whiteboards."
]
.join("\n")
)
);
}
#[test]
fn serializes_action_events_for_server_contract() {
let output = run_action_recipe_prepared_with_control(
ActionRuntimeInput {
recipe_id: "mindmap.generate".to_string(),
recipe_version: Some("v1".to_string()),
input: json!({}),
},
mock_control(json!({
"generate-structured": {
"result": "- Root"
}
})),
)
.unwrap();
let first = serde_json::to_value(output.events.first().unwrap()).unwrap();
let last = serde_json::to_value(output.events.last().unwrap()).unwrap();
assert_eq!(first["type"], json!("action_start"));
assert_eq!(last["type"], json!("action_done"));
assert_eq!(last["status"], json!("succeeded"));
assert_eq!(last["trace"]["status"], json!("succeeded"));
}
#[test]
fn built_in_action_fails_without_routes_or_mock_output() {
let output = run_action_recipe_prepared_with_control(
ActionRuntimeInput {
recipe_id: "mindmap.generate".to_string(),
recipe_version: Some("v1".to_string()),
input: json!({}),
},
ActionRuntimeControl::default(),
)
.unwrap();
assert_eq!(output.status, ActionRunStatus::Failed);
assert!(
output
.events
.last()
.and_then(|event| event.error_message.as_deref())
.unwrap_or_default()
.contains("promptStructured requires")
);
}
#[test]
fn built_in_image_action_uses_prompt_image_step_output() {
let output = run_action_recipe_prepared_with_control(
ActionRuntimeInput {
recipe_id: "image.filter.sketch".to_string(),
recipe_version: Some("v1".to_string()),
input: json!({}),
},
mock_control(json!({
"generate-image": {
"url": "https://example.com/artifact-1.png"
}
})),
)
.unwrap();
assert_eq!(output.status, ActionRunStatus::Succeeded);
assert_eq!(output.result, json!({ "url": "https://example.com/artifact-1.png" }));
assert_eq!(
output.state.pointer("/artifact/url"),
Some(&json!("https://example.com/artifact-1.png"))
);
}
#[test]
fn built_in_image_action_accepts_inline_artifact_output() {
let output = run_action_recipe_prepared_with_control(
ActionRuntimeInput {
recipe_id: "image.filter.sketch".to_string(),
recipe_version: Some("v1".to_string()),
input: json!({}),
},
mock_control(json!({
"generate-image": {
"data_base64": "aW1n",
"media_type": "image/webp"
}
})),
)
.unwrap();
assert_eq!(output.status, ActionRunStatus::Succeeded);
assert_eq!(
output.result,
json!({
"data_base64": "aW1n",
"media_type": "image/webp"
})
);
assert_eq!(output.state.pointer("/artifact/data_base64"), Some(&json!("aW1n")));
}
#[test]
fn rejects_invalid_recipe_without_final_step() {
let recipe = ActionRecipe {
id: "invalid.recipe".to_string(),
version: "v1".to_string(),
input_schema: json!({}),
output_schema: json!({}),
steps: vec![ActionRecipeStep {
id: "start".to_string(),
kind: ActionStepKind::ValidateJson,
input: None,
state_patch: None,
}],
};
let error = validate_recipe(&recipe).unwrap_err();
assert_eq!(error.status, Status::InvalidArg);
assert!(error.reason.contains("must end with a final step"));
}
#[test]
fn rejects_duplicated_recipe_identity() {
let recipe = ActionRecipe {
id: "duplicated.recipe".to_string(),
version: "v1".to_string(),
input_schema: json!({}),
output_schema: json!({}),
steps: vec![ActionRecipeStep {
id: "final".to_string(),
kind: ActionStepKind::Final,
input: None,
state_patch: None,
}],
};
let error = validate_catalog(&[recipe.clone(), recipe]).unwrap_err();
assert_eq!(error.status, Status::InvalidArg);
assert!(error.reason.contains("Duplicated action recipe"));
}
#[test]
fn rejects_recipe_where_final_step_is_not_last() {
let recipe = ActionRecipe {
id: "invalid.recipe".to_string(),
version: "v1".to_string(),
input_schema: json!({}),
output_schema: json!({}),
steps: vec![
ActionRecipeStep {
id: "final".to_string(),
kind: ActionStepKind::Final,
input: None,
state_patch: None,
},
ActionRecipeStep {
id: "after-final".to_string(),
kind: ActionStepKind::Transform,
input: None,
state_patch: None,
},
],
};
let error = validate_recipe(&recipe).unwrap_err();
assert_eq!(error.status, Status::InvalidArg);
assert!(error.reason.contains("must end with a final step"));
}
#[test]
fn validates_json_and_prompt_projection_steps() {
let recipe = test_recipe(vec![
ActionRecipeStep {
id: "prompt-structured".to_string(),
kind: ActionStepKind::PromptStructured,
input: Some(json!({})),
state_patch: None,
},
ActionRecipeStep {
id: "prompt-image".to_string(),
kind: ActionStepKind::PromptImage,
input: Some(json!({})),
state_patch: None,
},
ActionRecipeStep {
id: "validate-json".to_string(),
kind: ActionStepKind::ValidateJson,
input: Some(json!({
"schema": { "type": "object", "required": ["title"] },
"value": { "title": "Hello" }
})),
state_patch: None,
},
ActionRecipeStep {
id: "final".to_string(),
kind: ActionStepKind::Final,
input: Some(json!({ "copy": { "done": true } })),
state_patch: None,
},
]);
let output = run_action_recipe_for_test_with_control(
recipe,
runtime_input(json!({})),
mock_control(json!({
"prompt-structured": { "title": "Hello" },
"prompt-image": { "url": "https://example.com/artifact-1.png" }
})),
)
.unwrap();
assert_eq!(
output
.events
.iter()
.map(|event| event.event_type)
.filter(|event_type| matches!(event_type, ActionEventType::Attachment))
.collect::<Vec<_>>(),
vec![ActionEventType::Attachment]
);
assert_eq!(output.steps[2].output, Some(json!(true)));
}
#[test]
fn rejects_prompt_steps_without_prepared_routes_or_explicit_boundary() {
let recipe = test_recipe(vec![
ActionRecipeStep {
id: "prompt".to_string(),
kind: ActionStepKind::PromptStructured,
input: Some(json!({})),
state_patch: None,
},
ActionRecipeStep {
id: "final".to_string(),
kind: ActionStepKind::Final,
input: None,
state_patch: None,
},
]);
let output = run_action_recipe_for_test(recipe, runtime_input(json!({}))).unwrap();
assert_eq!(output.status, ActionRunStatus::Failed);
assert_eq!(output.error_code, Some("action_invalid_step".to_string()));
assert!(
output
.events
.last()
.and_then(|event| event.error_message.as_deref())
.unwrap_or_default()
.contains("requires")
);
}
#[test]
fn rejects_prompt_image_without_prepared_routes() {
let recipe = test_recipe(vec![
ActionRecipeStep {
id: "prompt-image".to_string(),
kind: ActionStepKind::PromptImage,
input: Some(json!({})),
state_patch: None,
},
ActionRecipeStep {
id: "final".to_string(),
kind: ActionStepKind::Final,
input: None,
state_patch: None,
},
]);
let output = run_action_recipe_for_test(recipe, runtime_input(json!({}))).unwrap();
assert_eq!(output.status, ActionRunStatus::Failed);
assert!(
output
.events
.last()
.and_then(|event| event.error_message.as_deref())
.unwrap_or_default()
.contains("preparedRoutes")
);
}
#[test]
fn validate_json_distinguishes_invalid_schema_from_invalid_value() {
let invalid_value = run_action_recipe_for_test(
test_recipe(vec![
ActionRecipeStep {
id: "validate-json".to_string(),
kind: ActionStepKind::ValidateJson,
input: Some(json!({
"schema": { "type": "object", "required": ["title"] },
"value": {}
})),
state_patch: None,
},
ActionRecipeStep {
id: "final".to_string(),
kind: ActionStepKind::Final,
input: Some(json!({ "copy": {} })),
state_patch: None,
},
]),
runtime_input(json!({})),
)
.unwrap();
assert_eq!(invalid_value.status, ActionRunStatus::Succeeded);
assert_eq!(invalid_value.steps[0].output, Some(json!(false)));
let invalid_schema = run_action_recipe_for_test(
test_recipe(vec![
ActionRecipeStep {
id: "validate-json".to_string(),
kind: ActionStepKind::ValidateJson,
input: Some(json!({
"schema": { "type": 1 },
"value": {}
})),
state_patch: None,
},
ActionRecipeStep {
id: "final".to_string(),
kind: ActionStepKind::Final,
input: None,
state_patch: None,
},
]),
runtime_input(json!({})),
)
.unwrap();
assert_eq!(invalid_schema.status, ActionRunStatus::Failed);
}
#[test]
fn emits_ordered_action_events_and_final_result() {
let output = run_action_recipe_for_test(
test_recipe(vec![ActionRecipeStep {
id: "final".to_string(),
kind: ActionStepKind::Final,
input: Some(json!({ "copy": {} })),
state_patch: Some(json!({ "finalized": true })),
}]),
ActionRuntimeInput {
recipe_id: "test.recipe".to_string(),
recipe_version: Some("v1".to_string()),
input: json!({ "content": "hello" }),
},
)
.unwrap();
assert_eq!(output.status, ActionRunStatus::Succeeded);
assert_eq!(output.result, json!({}));
assert_eq!(output.error_code, None);
assert_eq!(output.state, json!({ "content": "hello", "finalized": true }));
assert_eq!(output.steps.len(), 1);
assert_eq!(output.steps[0].id, "final");
assert_eq!(output.steps[0].output, Some(json!({})));
assert_eq!(output.steps[0].state_patch, Some(json!({ "finalized": true })));
assert_eq!(output.steps[0].error, None);
assert_eq!(
output.events.iter().map(|event| event.event_type).collect::<Vec<_>>(),
vec![
ActionEventType::ActionStart,
ActionEventType::StepStart,
ActionEventType::StepEnd,
ActionEventType::ActionDone,
]
);
}
fn runtime_input(input: serde_json::Value) -> ActionRuntimeInput {
ActionRuntimeInput {
recipe_id: "test.recipe".to_string(),
recipe_version: Some("v1".to_string()),
input,
}
}
fn mock_control(mock_output: serde_json::Value) -> ActionRuntimeControl {
ActionRuntimeControl {
abort_signal: None,
event_sender: None,
abort_after_events: None,
mock_output: Some(mock_output),
}
}
fn test_recipe(steps: Vec<ActionRecipeStep>) -> ActionRecipe {
ActionRecipe {
id: "test.recipe".to_string(),
version: "v1".to_string(),
input_schema: json!({}),
output_schema: json!({}),
steps,
}
}
#[test]
fn generates_lightweight_trace() {
let output = run_action_recipe_for_test(
test_recipe(vec![ActionRecipeStep {
id: "final".to_string(),
kind: ActionStepKind::Final,
input: Some(json!({ "copy": {} })),
state_patch: None,
}]),
ActionRuntimeInput {
recipe_id: "test.recipe".to_string(),
recipe_version: Some("v1".to_string()),
input: json!({}),
},
)
.unwrap();
assert_eq!(output.trace.status, ActionRunStatus::Succeeded);
assert!(!output.trace.lightweight.is_empty());
}
#[test]
fn abort_control_stops_runtime() {
let output = run_action_recipe_prepared_with_control(
ActionRuntimeInput {
recipe_id: "image.filter.sketch".to_string(),
recipe_version: Some("v1".to_string()),
input: json!({}),
},
ActionRuntimeControl {
abort_signal: None,
event_sender: None,
abort_after_events: Some(1),
mock_output: None,
},
)
.unwrap();
assert_eq!(output.status, ActionRunStatus::Aborted);
assert_eq!(output.error_code, Some(ACTION_ABORTED_ERROR_CODE.to_string()));
assert_eq!(
output.events.last().map(|event| event.event_type),
Some(ActionEventType::Error)
);
}
packages/backend/native/src/llm/assets/prompts/built-in.json
+94 -74
View File
@@ -2,8 +2,10 @@
{
"name": "Transcript audio",
"action": "Transcript audio",
"model": "gemini-3.5-flash-lite",
"optionalModels": ["gemini-3.5-flash-lite", "gemini-3.6-flash"],
"managedRoute": {
"targets": ["gemini-3.5-flash-lite"],
"premiumTargets": ["gemini-3.6-flash"]
},
"config": {
"requireContent": false,
"requireAttachment": true,
@@ -14,8 +16,10 @@
{
"name": "Transcript audio structured",
"action": "Transcript audio structured",
"model": "gemini-3.5-flash-lite",
"optionalModels": ["gemini-3.5-flash-lite", "gemini-3.6-flash"],
"managedRoute": {
"targets": ["gemini-3.5-flash-lite"],
"premiumTargets": ["gemini-3.6-flash"]
},
"config": {
"requireContent": false,
"requireAttachment": true,
@@ -35,7 +39,7 @@
{
"name": "Generate a caption",
"action": "Generate a caption",
"model": "gemini-3.5-flash-lite",
"managedRoute": { "targets": ["gemini-3.5-flash-lite"] },
"config": {
"requireContent": false,
"requireAttachment": true
@@ -50,7 +54,7 @@
{
"name": "Conversation Summary",
"action": "Conversation Summary",
"model": "gpt-5.6-luna",
"managedRoute": { "targets": ["gpt-5.6-luna"] },
"config": {
"requireContent": false
},
@@ -68,7 +72,7 @@
{
"name": "Summary",
"action": "Summary",
"model": "gpt-5.6-luna",
"managedRoute": { "targets": ["gpt-5.6-luna"] },
"messages": [
{
"role": "system",
@@ -83,7 +87,7 @@
{
"name": "Summary as title",
"action": "Summary as title",
"model": "gpt-5.6-luna",
"managedRoute": { "targets": ["gpt-5.6-luna"] },
"messages": [
{
"role": "system",
@@ -98,7 +102,7 @@
{
"name": "Summary the webpage",
"action": "Summary the webpage",
"model": "gpt-5.6-luna",
"managedRoute": { "targets": ["gpt-5.6-luna"] },
"messages": [
{
"role": "user",
@@ -109,7 +113,7 @@
{
"name": "Explain this",
"action": "Explain this",
"model": "gpt-5.6-luna",
"managedRoute": { "targets": ["gpt-5.6-luna"] },
"builtins": ["language"],
"messages": [
{
@@ -125,7 +129,7 @@
{
"name": "Explain this image",
"action": "Explain this image",
"model": "gemini-3.5-flash-lite",
"managedRoute": { "targets": ["gemini-3.5-flash-lite"] },
"config": {
"requireContent": false,
"requireAttachment": true
@@ -144,7 +148,7 @@
{
"name": "Explain this code",
"action": "Explain this code",
"model": "gemini-3.5-flash-lite",
"managedRoute": { "targets": ["gemini-3.5-flash-lite"] },
"messages": [
{
"role": "system",
@@ -159,7 +163,7 @@
{
"name": "Translate to",
"action": "Translate",
"model": "gemini-3.5-flash-lite",
"managedRoute": { "targets": ["gemini-3.5-flash-lite"] },
"params": {
"language": {
"default": "English",
@@ -192,7 +196,7 @@
{
"name": "Summarize the meeting structured",
"action": "Summarize the meeting structured",
"model": "gpt-5.6-luna",
"managedRoute": { "targets": ["gpt-5.6-luna"] },
"messages": [
{
"role": "system",
@@ -207,7 +211,7 @@
{
"name": "Summarize the meeting",
"action": "Summarize the meeting",
"model": "gpt-5.6-luna",
"managedRoute": { "targets": ["gpt-5.6-luna"] },
"messages": [
{
"role": "system",
@@ -222,7 +226,7 @@
{
"name": "Find action for summary",
"action": "Find action for summary",
"model": "gpt-5.6-luna",
"managedRoute": { "targets": ["gpt-5.6-luna"] },
"messages": [
{
"role": "system",
@@ -237,7 +241,7 @@
{
"name": "Write an article about this",
"action": "Write an article about this",
"model": "gpt-5.6-luna",
"managedRoute": { "targets": ["gpt-5.6-luna"] },
"messages": [
{
"role": "system",
@@ -252,7 +256,7 @@
{
"name": "Write a twitter about this",
"action": "Write a twitter about this",
"model": "gemini-3.5-flash-lite",
"managedRoute": { "targets": ["gemini-3.5-flash-lite"] },
"messages": [
{
"role": "system",
@@ -267,7 +271,7 @@
{
"name": "Write a poem about this",
"action": "Write a poem about this",
"model": "gemini-3.5-flash-lite",
"managedRoute": { "targets": ["gemini-3.5-flash-lite"] },
"messages": [
{
"role": "system",
@@ -282,7 +286,7 @@
{
"name": "Write a blog post about this",
"action": "Write a blog post about this",
"model": "gpt-5.6-luna",
"managedRoute": { "targets": ["gpt-5.6-luna"] },
"messages": [
{
"role": "system",
@@ -297,7 +301,7 @@
{
"name": "Write outline",
"action": "Write outline",
"model": "gemini-3.5-flash-lite",
"managedRoute": { "targets": ["gemini-3.5-flash-lite"] },
"messages": [
{
"role": "system",
@@ -312,7 +316,7 @@
{
"name": "Change tone to",
"action": "Change tone",
"model": "gemini-3.5-flash-lite",
"managedRoute": { "targets": ["gemini-3.5-flash-lite"] },
"params": {
"tone": {
"default": "professional",
@@ -333,7 +337,7 @@
{
"name": "Brainstorm ideas about this",
"action": "Brainstorm ideas about this",
"model": "gemini-3.5-flash-lite",
"managedRoute": { "targets": ["gemini-3.5-flash-lite"] },
"messages": [
{
"role": "system",
@@ -348,7 +352,7 @@
{
"name": "Brainstorm mindmap",
"action": "Brainstorm mindmap",
"model": "gpt-5.6-luna",
"managedRoute": { "targets": ["gpt-5.6-luna"] },
"messages": [
{
"role": "system",
@@ -363,7 +367,7 @@
{
"name": "Expand mind map",
"action": "Expand mind map",
"model": "gpt-5.6-luna",
"managedRoute": { "targets": ["gpt-5.6-luna"] },
"messages": [
{
"role": "system",
@@ -382,7 +386,7 @@
{
"name": "Improve writing for it",
"action": "Improve writing for it",
"model": "gemini-3.5-flash-lite",
"managedRoute": { "targets": ["gemini-3.5-flash-lite"] },
"messages": [
{
"role": "system",
@@ -397,7 +401,7 @@
{
"name": "Improve grammar for it",
"action": "Improve grammar for it",
"model": "gpt-5.6-luna",
"managedRoute": { "targets": ["gpt-5.6-luna"] },
"messages": [
{
"role": "system",
@@ -412,7 +416,7 @@
{
"name": "Fix spelling for it",
"action": "Fix spelling for it",
"model": "gemini-3.5-flash-lite",
"managedRoute": { "targets": ["gemini-3.5-flash-lite"] },
"messages": [
{
"role": "system",
@@ -427,7 +431,7 @@
{
"name": "Find action items from it",
"action": "Find action items from it",
"model": "gpt-5.6-luna",
"managedRoute": { "targets": ["gpt-5.6-luna"] },
"messages": [
{
"role": "system",
@@ -442,7 +446,7 @@
{
"name": "Check code error",
"action": "Check code error",
"model": "gpt-5.6-luna",
"managedRoute": { "targets": ["gpt-5.6-luna"] },
"messages": [
{
"role": "system",
@@ -457,7 +461,7 @@
{
"name": "Create a presentation",
"action": "Create a presentation",
"model": "gpt-5.6-luna",
"managedRoute": { "targets": ["gpt-5.6-luna"] },
"messages": [
{
"role": "system",
@@ -472,7 +476,7 @@
{
"name": "Create headings",
"action": "Create headings",
"model": "gemini-3.5-flash-lite",
"managedRoute": { "targets": ["gemini-3.5-flash-lite"] },
"messages": [
{
"role": "system",
@@ -487,7 +491,7 @@
{
"name": "Make it real",
"action": "Make it real",
"model": "claude-sonnet-4-6",
"managedRoute": { "targets": ["claude-sonnet-4-6"] },
"messages": [
{
"role": "system",
@@ -502,7 +506,7 @@
{
"name": "Make it real with text",
"action": "Make it real with text",
"model": "claude-sonnet-4-6",
"managedRoute": { "targets": ["claude-sonnet-4-6"] },
"messages": [
{
"role": "system",
@@ -517,7 +521,7 @@
{
"name": "Make it longer",
"action": "Make it longer",
"model": "gemini-3.5-flash-lite",
"managedRoute": { "targets": ["gemini-3.5-flash-lite"] },
"messages": [
{
"role": "system",
@@ -532,7 +536,7 @@
{
"name": "Make it shorter",
"action": "Make it shorter",
"model": "gemini-3.5-flash-lite",
"managedRoute": { "targets": ["gemini-3.5-flash-lite"] },
"messages": [
{
"role": "system",
@@ -547,7 +551,7 @@
{
"name": "Continue writing",
"action": "Continue writing",
"model": "gemini-3.5-flash-lite",
"managedRoute": { "targets": ["gemini-3.5-flash-lite"] },
"messages": [
{
"role": "system",
@@ -562,7 +566,7 @@
{
"name": "Section Edit",
"action": "Section Edit",
"model": "claude-sonnet-4-6",
"managedRoute": { "targets": ["claude-sonnet-4-6"] },
"messages": [
{
"role": "system",
@@ -577,7 +581,7 @@
{
"name": "Generate image",
"action": "image",
"model": "gpt-image-1",
"managedRoute": { "targets": ["gpt-image-1"] },
"messages": [
{
"role": "user",
@@ -588,7 +592,7 @@
{
"name": "Convert to Clay style",
"action": "Convert to Clay style",
"model": "gpt-image-1",
"managedRoute": { "targets": ["gpt-image-1"] },
"messages": [
{
"role": "user",
@@ -599,7 +603,7 @@
{
"name": "Convert to Sketch style",
"action": "Convert to Sketch style",
"model": "gpt-image-1",
"managedRoute": { "targets": ["gpt-image-1"] },
"messages": [
{
"role": "user",
@@ -610,7 +614,7 @@
{
"name": "Convert to Anime style",
"action": "Convert to Anime style",
"model": "gpt-image-1",
"managedRoute": { "targets": ["gpt-image-1"] },
"messages": [
{
"role": "user",
@@ -621,7 +625,7 @@
{
"name": "Convert to Pixel style",
"action": "Convert to Pixel style",
"model": "gpt-image-1",
"managedRoute": { "targets": ["gpt-image-1"] },
"messages": [
{
"role": "user",
@@ -632,7 +636,7 @@
{
"name": "Convert to sticker",
"action": "Convert to sticker",
"model": "gpt-image-1",
"managedRoute": { "targets": ["gpt-image-1"] },
"messages": [
{
"role": "user",
@@ -643,7 +647,7 @@
{
"name": "Upscale image",
"action": "Upscale image",
"model": "gpt-image-1",
"managedRoute": { "targets": ["gpt-image-1"] },
"messages": [
{
"role": "user",
@@ -654,7 +658,7 @@
{
"name": "Remove background",
"action": "Remove background",
"model": "gpt-image-1",
"managedRoute": { "targets": ["gpt-image-1"] },
"messages": [
{
"role": "user",
@@ -665,7 +669,7 @@
{
"name": "debug:action:fal-teed",
"action": "fal-teed",
"model": "workflowutils/teed",
"managedRoute": { "targets": ["workflowutils/teed"] },
"messages": [
{
"role": "user",
@@ -675,7 +679,7 @@
},
{
"name": "Code Artifact",
"model": "claude-sonnet-4-6",
"managedRoute": { "targets": ["claude-sonnet-4-6"] },
"messages": [
{
"role": "system",
@@ -689,13 +693,36 @@
},
{
"name": "Chat With AFFiNE AI",
"model": "gpt-5.6-luna",
"optionalModels": [
"gpt-5.6-luna",
"gpt-5.6-terra",
"gemini-3.6-flash",
"claude-sonnet-4-6"
],
"managedRoute": {
"targets": ["gpt-5.6-luna"],
"premiumTargets": ["gpt-5.6-luna"],
"selectableTargets": [
{
"id": "luna",
"modelId": "gpt-5.6-luna",
"displayName": "GPT 5.6 Luna",
"minimumTier": "standard"
},
{
"id": "terra",
"modelId": "gpt-5.6-terra",
"displayName": "GPT 5.6 Terra",
"minimumTier": "premium"
},
{
"id": "gemini",
"modelId": "gemini-3.6-flash",
"displayName": "Gemini 3.6 Flash",
"minimumTier": "premium"
},
{
"id": "claude",
"modelId": "claude-sonnet-4-6",
"displayName": "Claude Sonnet 4.6",
"minimumTier": "premium"
}
]
},
"config": {
"tools": [
"docRead",
@@ -708,8 +735,7 @@
"docCompose",
"codeArtifact",
"blobRead"
],
"proModels": ["gpt-5.6-terra", "gemini-3.6-flash", "claude-sonnet-4-6"]
]
},
"builtins": [
"date",
@@ -734,7 +760,7 @@
{
"name": "mindmap.generate",
"action": "mindmap.generate",
"model": "gpt-5.6-luna",
"managedRoute": { "targets": ["gpt-5.6-luna"] },
"config": {
"frequencyPenalty": 0.5,
"presencePenalty": 0.5,
@@ -759,7 +785,7 @@
{
"name": "slides.outline",
"action": "slides.outline",
"model": "gpt-5.6-luna",
"managedRoute": { "targets": ["gpt-5.6-luna"] },
"messages": [
{
"role": "system",
@@ -778,7 +804,7 @@
{
"name": "image.filter.sketch",
"action": "image.filter.sketch",
"model": "lora/image-to-image",
"managedRoute": { "targets": ["lora/image-to-image"] },
"config": {
"modelName": "stabilityai/stable-diffusion-xl-base-1.0",
"loras": [
@@ -798,7 +824,7 @@
{
"name": "image.filter.clay",
"action": "image.filter.clay",
"model": "lora/image-to-image",
"managedRoute": { "targets": ["lora/image-to-image"] },
"config": {
"modelName": "stabilityai/stable-diffusion-xl-base-1.0",
"loras": [
@@ -818,7 +844,7 @@
{
"name": "image.filter.anime",
"action": "image.filter.anime",
"model": "lora/image-to-image",
"managedRoute": { "targets": ["lora/image-to-image"] },
"config": {
"modelName": "stabilityai/stable-diffusion-xl-base-1.0",
"loras": [
@@ -838,7 +864,7 @@
{
"name": "image.filter.pixel",
"action": "image.filter.pixel",
"model": "lora/image-to-image",
"managedRoute": { "targets": ["lora/image-to-image"] },
"config": {
"modelName": "stabilityai/stable-diffusion-xl-base-1.0",
"loras": [
@@ -858,13 +884,12 @@
{
"name": "workflow:presentation",
"action": "workflow:presentation",
"model": "slides.outline",
"messages": []
},
{
"name": "workflow:presentation:step1",
"action": "workflow:presentation:step1",
"model": "gpt-5.6-luna",
"managedRoute": { "targets": ["gpt-5.6-luna"] },
"config": {
"temperature": 0.7
},
@@ -882,7 +907,7 @@
{
"name": "workflow:presentation:step2",
"action": "workflow:presentation:step2",
"model": "gpt-5.6-luna",
"managedRoute": { "targets": ["gpt-5.6-luna"] },
"messages": [
{
"role": "system",
@@ -901,7 +926,7 @@
{
"name": "workflow:presentation:step4",
"action": "workflow:presentation:step4",
"model": "gpt-5.6-luna",
"managedRoute": { "targets": ["gpt-5.6-luna"] },
"messages": [
{
"role": "system",
@@ -920,13 +945,12 @@
{
"name": "workflow:brainstorm",
"action": "workflow:brainstorm",
"model": "mindmap.generate",
"messages": []
},
{
"name": "workflow:brainstorm:step1",
"action": "workflow:brainstorm:step1",
"model": "gpt-5.6-luna",
"managedRoute": { "targets": ["gpt-5.6-luna"] },
"config": {
"temperature": 0.7
},
@@ -944,7 +968,7 @@
{
"name": "workflow:brainstorm:step2",
"action": "workflow:brainstorm:step2",
"model": "gpt-5.6-luna",
"managedRoute": { "targets": ["gpt-5.6-luna"] },
"config": {
"frequencyPenalty": 0.5,
"presencePenalty": 0.5,
@@ -969,25 +993,21 @@
{
"name": "workflow:image-sketch",
"action": "workflow:image-sketch",
"model": "image.filter.sketch",
"messages": []
},
{
"name": "workflow:image-clay",
"action": "workflow:image-clay",
"model": "image.filter.clay",
"messages": []
},
{
"name": "workflow:image-anime",
"action": "workflow:image-anime",
"model": "image.filter.anime",
"messages": []
},
{
"name": "workflow:image-pixel",
"action": "workflow:image-pixel",
"model": "image.filter.pixel",
"messages": []
}
]
packages/backend/native/src/llm/byok/catalog.rs
+110
View File
@@ -0,0 +1,110 @@
use std::collections::BTreeMap;
use llm_adapter::capability::provider_default_capability_upper_bound;
use serde::Serialize;
use sha2::{Digest, Sha256};
use super::{ByokCapabilityInput, contract::capability_input};
#[derive(Clone, Serialize)]
#[serde(rename_all = "camelCase")]
#[napi_derive::napi(object)]
pub struct ByokCatalogModelOutput {
pub model_id: String,
pub display_name: String,
pub recommended: bool,
pub capabilities: Vec<ByokCapabilityInput>,
}
#[derive(Clone, Serialize)]
#[serde(rename_all = "camelCase")]
#[napi_derive::napi(object)]
pub struct ByokCatalogProviderOutput {
pub provider: String,
pub models: Vec<ByokCatalogModelOutput>,
}
#[derive(Clone, Serialize)]
#[serde(rename_all = "camelCase")]
#[napi_derive::napi(object)]
pub struct ByokCatalogOutput {
pub version: String,
pub providers: Vec<ByokCatalogProviderOutput>,
}
pub fn byok_catalog() -> ByokCatalogOutput {
let variants = llm_adapter::core::default_model_registry_variants();
let mut providers = ["openai", "anthropic", "gemini", "fal"]
.into_iter()
.map(|provider| (provider, BTreeMap::new()))
.collect::<BTreeMap<_, BTreeMap<String, ByokCatalogModelOutput>>>();
for variant in variants {
let Some(provider) = provider_for_backend(&variant.backend_kind) else {
continue;
};
let Some(capabilities) = provider_default_capability_upper_bound(provider, &variant.raw_model_id) else {
continue;
};
providers
.entry(provider)
.or_default()
.entry(variant.raw_model_id.clone())
.or_insert_with(|| ByokCatalogModelOutput {
model_id: variant.raw_model_id.clone(),
display_name: variant.display_name.unwrap_or_else(|| variant.raw_model_id.clone()),
recommended: variant
.capabilities
.iter()
.any(|capability| capability.default_for_output_type == Some(true)),
capabilities: capabilities.into_iter().map(capability_input).collect(),
});
}
let providers = providers
.into_iter()
.map(|(provider, models)| ByokCatalogProviderOutput {
provider: provider.to_string(),
models: models.into_values().collect(),
})
.collect::<Vec<_>>();
let encoded = serde_json::to_vec(&providers).expect("BYOK catalog must serialize");
let version = Sha256::digest(encoded)
.iter()
.take(8)
.map(|byte| format!("{byte:02x}"))
.collect();
ByokCatalogOutput { version, providers }
}
fn provider_for_backend(backend: &str) -> Option<&'static str> {
match backend {
"openai_chat" | "openai_responses" => Some("openai"),
"anthropic" => Some("anthropic"),
"gemini_api" => Some("gemini"),
"fal" => Some("fal"),
_ => None,
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn catalog_contains_explicit_provider_default_declarations() {
let catalog = byok_catalog();
assert!(!catalog.version.is_empty());
for provider in &catalog.providers {
assert!(
!provider.models.is_empty(),
"{} has no catalog models",
provider.provider
);
for model in &provider.models {
assert!(!model.model_id.is_empty());
assert!(!model.capabilities.is_empty());
}
}
}
}
packages/backend/native/src/llm/byok/contract.rs
+579
View File
@@ -0,0 +1,579 @@
use std::collections::HashSet;
use llm_adapter::{
capability::{
AttachmentKind, AttachmentSource, DeclaredModelCapability, ModelFeature, ModelInput, ModelOutput,
provider_default_capability_upper_bound, validate_capability_upper_bound, validate_declared_capability,
},
target::canonicalize_endpoint,
};
use serde::{Deserialize, Serialize};
use thiserror::Error;
#[derive(Clone, Deserialize, Serialize)]
#[serde(rename_all = "camelCase", deny_unknown_fields)]
#[napi_derive::napi(object)]
pub struct ByokCapabilityInput {
pub input: Vec<String>,
pub output: Vec<String>,
pub features: Vec<String>,
pub attachment_kinds: Vec<String>,
pub attachment_sources: Vec<String>,
}
#[derive(Clone, Deserialize, Serialize)]
#[serde(rename_all = "camelCase", deny_unknown_fields)]
#[napi_derive::napi(object)]
pub struct ByokModelDeclarationInput {
pub model_id: String,
pub enabled: bool,
pub capabilities: Vec<ByokCapabilityInput>,
}
#[derive(Clone, Deserialize, Serialize)]
#[serde(rename_all = "camelCase", deny_unknown_fields)]
#[napi_derive::napi(object)]
pub struct ByokEndpointInput {
pub kind: String,
pub url: Option<String>,
}
#[derive(Clone, Deserialize, Serialize)]
#[serde(rename_all = "camelCase", deny_unknown_fields)]
#[napi_derive::napi(object)]
pub struct ByokProfileDefinitionInput {
pub version: u32,
pub endpoint: ByokEndpointInput,
pub models: Vec<ByokModelDeclarationInput>,
}
#[derive(Clone)]
#[napi_derive::napi(object)]
pub struct CreateByokProfileInput {
pub workspace_id: String,
pub provider: String,
pub name: String,
pub description: Option<String>,
pub credential: String,
pub definition: ByokProfileDefinitionInput,
pub enabled: bool,
pub actor_user_id: String,
}
#[derive(Clone, Deserialize, Serialize)]
#[serde(rename_all = "camelCase", deny_unknown_fields)]
#[napi_derive::napi(object)]
pub struct ReplaceByokProfileInput {
pub workspace_id: String,
pub profile_id: String,
pub expected_revision: i32,
pub name: String,
pub description: Option<String>,
pub definition: ByokProfileDefinitionInput,
pub credential: Option<String>,
pub enabled: bool,
pub actor_user_id: String,
}
#[derive(Clone)]
#[napi_derive::napi(object)]
pub struct RotateByokCredentialInput {
pub workspace_id: String,
pub profile_id: String,
pub expected_revision: i32,
pub credential: String,
pub actor_user_id: String,
}
#[derive(Clone, Deserialize, Serialize)]
#[serde(rename_all = "camelCase", deny_unknown_fields)]
#[napi_derive::napi(object)]
pub struct ByokProfileOrderInput {
pub profile_id: String,
pub expected_revision: i32,
}
#[derive(Clone, Deserialize, Serialize)]
#[serde(rename_all = "camelCase", deny_unknown_fields)]
#[napi_derive::napi(object)]
pub struct ReorderByokProfilesInput {
pub workspace_id: String,
pub profiles: Vec<ByokProfileOrderInput>,
pub actor_user_id: String,
}
#[derive(Clone, Deserialize, Serialize)]
#[serde(rename_all = "camelCase", deny_unknown_fields)]
#[napi_derive::napi(object)]
pub struct ProbeByokProfileInput {
pub workspace_id: String,
pub profile_id: String,
pub checks: Vec<ByokProbeCheckInput>,
}
#[derive(Clone, Deserialize, Serialize)]
#[serde(rename_all = "camelCase", deny_unknown_fields)]
#[napi_derive::napi(object)]
pub struct ProbeByokDraftInput {
pub workspace_id: String,
pub provider: String,
pub credential: Option<String>,
pub profile_id: Option<String>,
pub expected_revision: Option<i32>,
pub definition: ByokProfileDefinitionInput,
pub checks: Vec<ByokProbeCheckInput>,
}
#[derive(Clone, Deserialize, Serialize)]
#[serde(rename_all = "camelCase", deny_unknown_fields)]
#[napi_derive::napi(object)]
pub struct ByokProbeCheckInput {
pub model_id: String,
pub operation: String,
}
#[derive(Clone)]
#[napi_derive::napi(object)]
pub struct CreateByokLocalLeaseProviderInput {
pub provider: String,
pub name: String,
pub description: Option<String>,
pub credential: String,
pub definition: ByokProfileDefinitionInput,
pub enabled: bool,
}
#[derive(Clone)]
#[napi_derive::napi(object)]
pub struct CreateByokLocalLeaseInput {
pub workspace_id: String,
pub user_id: String,
pub providers: Vec<CreateByokLocalLeaseProviderInput>,
}
#[derive(Clone, Deserialize, Serialize)]
#[serde(rename_all = "camelCase", deny_unknown_fields)]
#[napi_derive::napi(object)]
pub struct ByokLocalLeaseOutput {
pub lease_id: String,
pub expires_at_ms: i64,
}
#[derive(Clone, Deserialize, Serialize)]
#[serde(rename_all = "camelCase", deny_unknown_fields)]
#[napi_derive::napi(object)]
pub struct ByokProfileOutput {
pub profile_id: String,
pub workspace_id: String,
pub provider: String,
pub name: String,
pub description: Option<String>,
pub definition: ByokProfileDefinitionInput,
pub enabled: bool,
pub sort_order: i32,
pub revision: i32,
pub validation: Option<ByokValidationOutput>,
}
#[derive(Clone, Deserialize, Serialize)]
#[serde(rename_all = "camelCase", deny_unknown_fields)]
#[napi_derive::napi(object)]
pub struct ByokProbeStatusOutput {
pub kind: String,
pub tested_at_ms: Option<i64>,
pub error_kind: Option<String>,
}
#[derive(Clone, Deserialize, Serialize)]
#[serde(rename_all = "camelCase", deny_unknown_fields)]
#[napi_derive::napi(object)]
pub struct ByokModelProbeOutput {
pub model_id: String,
pub checks: Vec<ByokModelProbeCheckOutput>,
}
#[derive(Clone, Deserialize, Serialize)]
#[serde(rename_all = "camelCase", deny_unknown_fields)]
#[napi_derive::napi(object)]
pub struct ByokModelProbeCheckOutput {
pub operation: String,
pub status: ByokProbeStatusOutput,
}
#[derive(Clone, Deserialize, Serialize)]
#[serde(rename_all = "camelCase", deny_unknown_fields)]
#[napi_derive::napi(object)]
pub struct ByokValidationOutput {
pub definition_fingerprint: String,
pub credential_generation: i32,
pub connection: ByokProbeStatusOutput,
pub models: Vec<ByokModelProbeOutput>,
}
#[derive(Clone, Deserialize, Serialize)]
#[serde(rename_all = "camelCase", deny_unknown_fields)]
#[napi_derive::napi(object)]
pub struct ByokProbeResultOutput {
pub definition_fingerprint: String,
pub stale: bool,
pub connection: ByokProbeStatusOutput,
pub models: Vec<ByokModelProbeOutput>,
}
#[derive(Clone, PartialEq, Eq, Deserialize, Serialize)]
#[serde(tag = "kind", rename_all = "snake_case", deny_unknown_fields)]
pub(crate) enum ByokEndpoint {
ProviderDefault,
Custom { url: String },
}
#[derive(Clone, PartialEq, Eq, Deserialize, Serialize)]
#[serde(rename_all = "camelCase", deny_unknown_fields)]
pub(crate) struct ByokModelDeclaration {
pub(crate) model_id: String,
pub(crate) enabled: bool,
pub(crate) capabilities: Vec<DeclaredModelCapability>,
}
#[derive(Clone, PartialEq, Eq, Deserialize, Serialize)]
#[serde(rename_all = "camelCase", deny_unknown_fields)]
pub(crate) struct ByokProfileDefinition {
pub(crate) version: u32,
pub(crate) endpoint: ByokEndpoint,
pub(crate) models: Vec<ByokModelDeclaration>,
}
#[derive(Debug, Error)]
pub(crate) enum ByokContractError {
#[error("unsupported BYOK definition version")]
Version,
#[error("unsupported BYOK provider")]
Provider,
#[error("{0} is required")]
Required(&'static str),
#[error("duplicate {0}")]
Duplicate(&'static str),
#[error("invalid BYOK endpoint")]
Endpoint,
#[error("invalid model capability: {0}")]
Capability(String),
#[error("declared capability exceeds provider or model upper bound")]
CapabilityUpperBound,
}
impl ByokProfileDefinition {
pub(crate) fn endpoint_identity(&self) -> &str {
match &self.endpoint {
ByokEndpoint::ProviderDefault => "default",
ByokEndpoint::Custom { url } => url,
}
}
}
pub(crate) fn validate_definition(
provider: &str,
input: ByokProfileDefinitionInput,
) -> Result<ByokProfileDefinition, ByokContractError> {
if input.version != 1 {
return Err(ByokContractError::Version);
}
if !matches!(provider, "openai" | "anthropic" | "gemini" | "fal") {
return Err(ByokContractError::Provider);
}
let endpoint = match (input.endpoint.kind.as_str(), input.endpoint.url) {
("provider_default", None) => ByokEndpoint::ProviderDefault,
("custom", Some(url)) if !url.trim().is_empty() => ByokEndpoint::Custom {
url: canonicalize_endpoint(&url).map_err(|_| ByokContractError::Endpoint)?,
},
_ => return Err(ByokContractError::Endpoint),
};
if input.models.is_empty() {
return Err(ByokContractError::Required("models"));
}
let mut ids = HashSet::new();
let mut models = Vec::with_capacity(input.models.len());
for model in input.models {
let model_id = model.model_id.trim().to_string();
if model_id.is_empty() || model_id.len() > 512 {
return Err(ByokContractError::Required("modelId"));
}
if !ids.insert(model_id.clone()) {
return Err(ByokContractError::Duplicate("modelId"));
}
if model.capabilities.is_empty() {
return Err(ByokContractError::Required("capabilities"));
}
let capabilities = model
.capabilities
.into_iter()
.map(parse_capability)
.collect::<Result<Vec<_>, _>>()?;
validate_upper_bound(provider, &endpoint, &model_id, &capabilities)?;
models.push(ByokModelDeclaration {
model_id,
enabled: model.enabled,
capabilities,
});
}
Ok(ByokProfileDefinition {
version: 1,
endpoint,
models,
})
}
fn parse_capability(input: ByokCapabilityInput) -> Result<DeclaredModelCapability, ByokContractError> {
let capability = DeclaredModelCapability {
input: parse_values(input.input, |value| match value {
"text" => Some(ModelInput::Text),
"image" => Some(ModelInput::Image),
"audio" => Some(ModelInput::Audio),
"file" => Some(ModelInput::File),
_ => None,
})?,
output: parse_values(input.output, |value| match value {
"text" => Some(ModelOutput::Text),
"object" => Some(ModelOutput::Object),
"structured" => Some(ModelOutput::Structured),
"embedding" => Some(ModelOutput::Embedding),
"rerank" => Some(ModelOutput::Rerank),
"image" => Some(ModelOutput::Image),
_ => None,
})?,
features: parse_values(input.features, |value| match value {
"tool_calling" => Some(ModelFeature::ToolCalling),
"reasoning" => Some(ModelFeature::Reasoning),
"web_search" => Some(ModelFeature::WebSearch),
_ => None,
})?,
attachment_kinds: parse_values(input.attachment_kinds, |value| match value {
"image" => Some(AttachmentKind::Image),
"audio" => Some(AttachmentKind::Audio),
"file" => Some(AttachmentKind::File),
_ => None,
})?,
attachment_sources: parse_values(input.attachment_sources, |value| match value {
"url" => Some(AttachmentSource::Url),
"data" => Some(AttachmentSource::Data),
"bytes" => Some(AttachmentSource::Bytes),
"file_handle" => Some(AttachmentSource::FileHandle),
_ => None,
})?,
};
validate_declared_capability(&capability).map_err(|error| ByokContractError::Capability(error.to_string()))?;
Ok(capability)
}
fn parse_values<T>(values: Vec<String>, parse: impl Fn(&str) -> Option<T>) -> Result<Vec<T>, ByokContractError> {
values
.into_iter()
.map(|value| parse(&value).ok_or_else(|| ByokContractError::Capability(format!("unknown enum {value}"))))
.collect()
}
fn validate_upper_bound(
provider: &str,
endpoint: &ByokEndpoint,
model_id: &str,
capabilities: &[DeclaredModelCapability],
) -> Result<(), ByokContractError> {
if provider == "fal"
&& capabilities.iter().any(|capability| {
capability.output.iter().any(|output| *output != ModelOutput::Image)
|| capability
.input
.iter()
.any(|input| !matches!(input, ModelInput::Text | ModelInput::Image))
})
{
return Err(ByokContractError::CapabilityUpperBound);
}
if matches!(endpoint, ByokEndpoint::Custom { .. }) {
return Ok(());
}
let upper_bound =
provider_default_capability_upper_bound(provider, model_id).ok_or(ByokContractError::CapabilityUpperBound)?;
for capability in capabilities {
validate_capability_upper_bound(capability, &upper_bound).map_err(|_| ByokContractError::CapabilityUpperBound)?;
}
Ok(())
}
fn input_name(value: &ModelInput) -> &'static str {
match value {
ModelInput::Text => "text",
ModelInput::Image => "image",
ModelInput::Audio => "audio",
ModelInput::File => "file",
}
}
fn output_name(value: &ModelOutput) -> &'static str {
match value {
ModelOutput::Text => "text",
ModelOutput::Object => "object",
ModelOutput::Structured => "structured",
ModelOutput::Embedding => "embedding",
ModelOutput::Rerank => "rerank",
ModelOutput::Image => "image",
}
}
fn attachment_kind_name(value: &AttachmentKind) -> &'static str {
match value {
AttachmentKind::Image => "image",
AttachmentKind::Audio => "audio",
AttachmentKind::File => "file",
}
}
fn attachment_source_name(value: &AttachmentSource) -> &'static str {
match value {
AttachmentSource::Url => "url",
AttachmentSource::Data => "data",
AttachmentSource::Bytes => "bytes",
AttachmentSource::FileHandle => "file_handle",
}
}
impl From<ByokProfileDefinition> for ByokProfileDefinitionInput {
fn from(definition: ByokProfileDefinition) -> Self {
Self {
version: definition.version,
endpoint: match definition.endpoint {
ByokEndpoint::ProviderDefault => ByokEndpointInput {
kind: "provider_default".to_string(),
url: None,
},
ByokEndpoint::Custom { url } => ByokEndpointInput {
kind: "custom".to_string(),
url: Some(url),
},
},
models: definition
.models
.into_iter()
.map(|model| ByokModelDeclarationInput {
model_id: model.model_id,
enabled: model.enabled,
capabilities: model.capabilities.into_iter().map(capability_input).collect(),
})
.collect(),
}
}
}
pub(super) fn capability_input(capability: DeclaredModelCapability) -> ByokCapabilityInput {
ByokCapabilityInput {
input: capability.input.iter().map(input_name).map(str::to_string).collect(),
output: capability.output.iter().map(output_name).map(str::to_string).collect(),
features: capability
.features
.iter()
.map(|value| match value {
ModelFeature::ToolCalling => "tool_calling",
ModelFeature::Reasoning => "reasoning",
ModelFeature::WebSearch => "web_search",
})
.map(str::to_string)
.collect(),
attachment_kinds: capability
.attachment_kinds
.iter()
.map(attachment_kind_name)
.map(str::to_string)
.collect(),
attachment_sources: capability
.attachment_sources
.iter()
.map(attachment_source_name)
.map(str::to_string)
.collect(),
}
}
#[cfg(test)]
mod tests {
use super::*;
fn definition(model_id: &str, capabilities: Vec<ByokCapabilityInput>) -> ByokProfileDefinitionInput {
ByokProfileDefinitionInput {
version: 1,
endpoint: ByokEndpointInput {
kind: "custom".to_string(),
url: Some("https://example.com/v1/".to_string()),
},
models: vec![ByokModelDeclarationInput {
model_id: model_id.to_string(),
enabled: true,
capabilities,
}],
}
}
fn text_capability() -> ByokCapabilityInput {
ByokCapabilityInput {
input: vec!["text".to_string()],
output: vec!["text".to_string()],
features: vec![],
attachment_kinds: vec![],
attachment_sources: vec![],
}
}
#[test]
fn keeps_opaque_model_id_and_explicit_empty_features() {
let validated =
validate_definition("openai", definition(" vendor/model:latest ", vec![text_capability()])).unwrap();
assert_eq!(validated.models[0].model_id, "vendor/model:latest");
assert!(validated.models[0].capabilities[0].features.is_empty());
assert_eq!(validated.endpoint_identity(), "https://example.com/v1");
}
#[test]
fn rejects_invalid_shapes_table() {
let mut cases = Vec::new();
cases.push(ByokProfileDefinitionInput {
models: vec![],
..definition("model", vec![text_capability()])
});
cases.push(definition("", vec![text_capability()]));
cases.push(definition("model", vec![]));
let mut empty_input = text_capability();
empty_input.input.clear();
cases.push(definition("model", vec![empty_input]));
let mut duplicate = text_capability();
duplicate.output.push("text".to_string());
cases.push(definition("model", vec![duplicate]));
assert!(
cases
.into_iter()
.all(|case| validate_definition("openai", case).is_err())
);
}
#[test]
fn rejects_endpoint_tag_mismatches() {
for endpoint in [
ByokEndpointInput {
kind: "provider_default".to_string(),
url: Some("https://example.com".to_string()),
},
ByokEndpointInput {
kind: "custom".to_string(),
url: None,
},
ByokEndpointInput {
kind: "custom".to_string(),
url: Some(" ".to_string()),
},
] {
let mut input = definition("model", vec![text_capability()]);
input.endpoint = endpoint;
assert!(validate_definition("openai", input).is_err());
}
}
}
packages/backend/native/src/llm/byok/envelope.rs
+171
View File
@@ -0,0 +1,171 @@
use aes_gcm::{
Aes256Gcm, Nonce,
aead::{Aead, KeyInit, Payload},
};
use base64::{Engine, engine::general_purpose::URL_SAFE_NO_PAD};
use hkdf::Hkdf;
use rand::RngCore;
use sha2::Sha256;
use thiserror::Error;
use zeroize::{Zeroize, Zeroizing};
const PREFIX: &str = "byok:v1:";
const INFO: &[u8] = b"AFFiNE/Copilot/BYOK/v1";
const NONCE_LEN: usize = 12;
const TAG_LEN: usize = 16;
pub(crate) struct CredentialEnvelopeKey(Zeroizing<[u8; 32]>);
pub(crate) struct SensitiveCredential(Zeroizing<Vec<u8>>);
impl SensitiveCredential {
pub(crate) fn new(value: impl Into<Vec<u8>>) -> Self {
Self(Zeroizing::new(value.into()))
}
pub(crate) fn expose(&self) -> &[u8] {
self.0.as_slice()
}
}
#[derive(Debug, Error)]
pub(crate) enum CredentialEnvelopeError {
#[error("credential_unavailable")]
Unavailable,
}
impl CredentialEnvelopeKey {
pub(crate) fn derive(root_secret: &[u8]) -> Result<Self, CredentialEnvelopeError> {
if root_secret.is_empty() {
return Err(CredentialEnvelopeError::Unavailable);
}
let mut key = Zeroizing::new([0_u8; 32]);
Hkdf::<Sha256>::new(None, root_secret)
.expand(INFO, key.as_mut())
.map_err(|_| CredentialEnvelopeError::Unavailable)?;
Ok(Self(key))
}
pub(crate) fn encrypt(
&self,
credential: &SensitiveCredential,
aad: &[u8],
) -> Result<String, CredentialEnvelopeError> {
let cipher = Aes256Gcm::new_from_slice(self.0.as_slice()).map_err(|_| CredentialEnvelopeError::Unavailable)?;
let mut nonce = [0_u8; NONCE_LEN];
rand::rng().fill_bytes(&mut nonce);
let ciphertext = cipher
.encrypt(
Nonce::from_slice(&nonce),
Payload {
msg: credential.expose(),
aad,
},
)
.map_err(|_| CredentialEnvelopeError::Unavailable)?;
let mut body = Vec::with_capacity(NONCE_LEN + ciphertext.len());
body.extend_from_slice(&nonce);
body.extend_from_slice(&ciphertext);
let encoded = URL_SAFE_NO_PAD.encode(&body);
body.zeroize();
Ok(format!("{PREFIX}{encoded}"))
}
pub(crate) fn decrypt(&self, envelope: &str, aad: &[u8]) -> Result<SensitiveCredential, CredentialEnvelopeError> {
let encoded = envelope
.strip_prefix(PREFIX)
.ok_or(CredentialEnvelopeError::Unavailable)?;
let mut body = URL_SAFE_NO_PAD
.decode(encoded)
.map_err(|_| CredentialEnvelopeError::Unavailable)?;
if body.len() < NONCE_LEN + TAG_LEN {
body.zeroize();
return Err(CredentialEnvelopeError::Unavailable);
}
let (nonce, ciphertext) = body.split_at(NONCE_LEN);
let cipher = Aes256Gcm::new_from_slice(self.0.as_slice()).map_err(|_| CredentialEnvelopeError::Unavailable)?;
let result = cipher
.decrypt(Nonce::from_slice(nonce), Payload { msg: ciphertext, aad })
.map(SensitiveCredential::new)
.map_err(|_| CredentialEnvelopeError::Unavailable);
body.zeroize();
result
}
}
pub(crate) fn server_aad(workspace_id: &str, profile_id: &str, provider: &str, endpoint_identity: &str) -> Vec<u8> {
["server", workspace_id, profile_id, provider, endpoint_identity]
.join("\0")
.into_bytes()
}
pub(crate) fn local_aad(
workspace_id: &str,
user_id: &str,
lease_id: &str,
index: usize,
provider: &str,
endpoint_identity: &str,
) -> Vec<u8> {
[
"local".to_string(),
workspace_id.to_string(),
user_id.to_string(),
lease_id.to_string(),
index.to_string(),
provider.to_string(),
endpoint_identity.to_string(),
]
.join("\0")
.into_bytes()
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn round_trips_with_random_nonce_and_rejects_tampering() {
let key = CredentialEnvelopeKey::derive(b"stable-root").unwrap();
let credential = SensitiveCredential::new(b"secret".to_vec());
let aad = server_aad("workspace", "profile", "openai", "default");
let first = key.encrypt(&credential, &aad).unwrap();
let second = key.encrypt(&credential, &aad).unwrap();
assert_ne!(first, second);
assert_eq!(key.decrypt(&first, &aad).unwrap().expose(), b"secret");
let mut tampered = first.into_bytes();
let last = tampered.len() - 1;
tampered[last] = if tampered[last] == b'A' { b'B' } else { b'A' };
assert!(key.decrypt(std::str::from_utf8(&tampered).unwrap(), &aad).is_err());
}
#[test]
fn rejects_wrong_context_key_version_and_legacy_ciphertext() {
let key = CredentialEnvelopeKey::derive(b"stable-root").unwrap();
let credential = SensitiveCredential::new(b"secret".to_vec());
let aad = server_aad("workspace", "profile", "openai", "default");
let encrypted = key.encrypt(&credential, &aad).unwrap();
assert!(
key
.decrypt(&encrypted, &server_aad("other", "profile", "openai", "default"))
.is_err()
);
assert!(
CredentialEnvelopeKey::derive(b"other")
.unwrap()
.decrypt(&encrypted, &aad)
.is_err()
);
assert!(
key
.decrypt(&encrypted.replacen("byok:v1:", "byok:v2:", 1), &aad)
.is_err()
);
assert!(key.decrypt("bGVnYWN5", &aad).is_err());
}
#[test]
fn requires_stable_root_secret() {
assert!(CredentialEnvelopeKey::derive(b"").is_err());
}
}
packages/backend/native/src/llm/byok/mod.rs
+16
View File
@@ -0,0 +1,16 @@
mod catalog;
mod contract;
mod envelope;
mod validation;
pub use catalog::{ByokCatalogModelOutput, ByokCatalogOutput, ByokCatalogProviderOutput, byok_catalog};
pub use contract::{
ByokCapabilityInput, ByokEndpointInput, ByokLocalLeaseOutput, ByokModelDeclarationInput, ByokModelProbeCheckOutput,
ByokModelProbeOutput, ByokProbeCheckInput, ByokProbeResultOutput, ByokProbeStatusOutput, ByokProfileDefinitionInput,
ByokProfileOutput, ByokValidationOutput, CreateByokLocalLeaseInput, CreateByokLocalLeaseProviderInput,
CreateByokProfileInput, ProbeByokDraftInput, ProbeByokProfileInput, ReorderByokProfilesInput,
ReplaceByokProfileInput, RotateByokCredentialInput,
};
pub(crate) use contract::{ByokEndpoint, ByokModelDeclaration, ByokProfileDefinition, validate_definition};
pub(crate) use envelope::{CredentialEnvelopeKey, SensitiveCredential, local_aad, server_aad};
pub(crate) use validation::{definition_fingerprint, reconcile_validation};
packages/backend/native/src/llm/byok/validation.rs
+93
View File
@@ -0,0 +1,93 @@
use sha2::{Digest, Sha256};
use super::{ByokProfileDefinition, ByokValidationOutput};
pub(crate) fn definition_fingerprint(definition: &ByokProfileDefinition) -> String {
let encoded = serde_json::to_vec(definition).expect("validated BYOK definition must serialize");
Sha256::digest(encoded)
.iter()
.map(|byte| format!("{byte:02x}"))
.collect()
}
pub(crate) fn reconcile_validation(
validation: Option<ByokValidationOutput>,
old_definition: &ByokProfileDefinition,
definition: &ByokProfileDefinition,
credential_generation: i32,
credential_changed: bool,
) -> Option<ByokValidationOutput> {
let mut validation = validation?;
if credential_changed || old_definition.endpoint != definition.endpoint {
return None;
}
validation.models.retain(|evidence| {
let old = old_definition
.models
.iter()
.find(|model| model.model_id == evidence.model_id);
let new = definition
.models
.iter()
.find(|model| model.model_id == evidence.model_id);
old.is_some() && old == new
});
validation.definition_fingerprint = definition_fingerprint(definition);
validation.credential_generation = credential_generation;
Some(validation)
}
#[cfg(test)]
mod tests {
use llm_adapter::capability::{DeclaredModelCapability, ModelInput, ModelOutput};
use super::*;
use crate::llm::byok::{ByokEndpoint, ByokModelDeclaration, ByokModelProbeOutput, ByokProbeStatusOutput};
fn definition(models: &[&str]) -> ByokProfileDefinition {
ByokProfileDefinition {
version: 1,
endpoint: ByokEndpoint::ProviderDefault,
models: models
.iter()
.map(|model| ByokModelDeclaration {
model_id: (*model).to_string(),
enabled: true,
capabilities: vec![DeclaredModelCapability {
input: vec![ModelInput::Text],
output: vec![ModelOutput::Text],
features: vec![],
attachment_kinds: vec![],
attachment_sources: vec![],
}],
})
.collect(),
}
}
#[test]
fn keeps_only_unchanged_model_evidence() {
let old = definition(&["a", "b"]);
let new = definition(&["b", "c"]);
let validation = ByokValidationOutput {
definition_fingerprint: definition_fingerprint(&old),
credential_generation: 1,
connection: ByokProbeStatusOutput {
kind: "verified".to_string(),
tested_at_ms: Some(1),
error_kind: None,
},
models: ["a", "b"]
.into_iter()
.map(|model_id| ByokModelProbeOutput {
model_id: model_id.to_string(),
checks: vec![],
})
.collect(),
};
let next = reconcile_validation(Some(validation), &old, &new, 1, false).unwrap();
assert_eq!(next.models.len(), 1);
assert_eq!(next.models[0].model_id, "b");
assert_eq!(next.definition_fingerprint, definition_fingerprint(&new));
}
}
packages/backend/native/src/llm/contract_schema.rs
+4 -150
View File
@@ -7,8 +7,7 @@ use super::{
action::{TranscriptGeneratedResult, TranscriptInputContract, TranscriptResult},
core::contracts::{
CapabilityMatchRequest, CapabilityMatchResponse, ModelConditionsContract, ModelRegistryMatchRequest,
ModelRegistryMatchResponse, ModelRegistryResolveRequest, ModelRegistryResolveResponse, PromptRenderContract,
PromptSessionContract, ProviderDriverSpec, RequestedModelMatchRequest, RequestedModelMatchResponse,
ModelRegistryMatchResponse, ModelRegistryResolveRequest, ModelRegistryResolveResponse, ProviderDriverSpec,
},
};
@@ -78,7 +77,7 @@ fn mark_definition_property_nullable(schema: &mut Value, definition: &str, prope
pub(crate) fn transcript_input_schema() -> Value {
let mut schema = generated_schema_for::<TranscriptInputContract>();
for property in ["sourceAudio", "quality", "infos", "sliceManifest", "preparedRoutes"] {
for property in ["sourceAudio", "quality", "infos", "sliceManifest"] {
mark_property_nullable(&mut schema, property);
}
mark_definition_property_nullable(&mut schema, "TranscriptAudioInfo", "index");
@@ -88,7 +87,7 @@ pub(crate) fn transcript_input_schema() -> Value {
pub(crate) fn transcript_generated_result_schema() -> Value {
let mut schema = generated_schema_for::<TranscriptGeneratedResult>();
for property in ["normalizedSegments", "summaryJson", "providerMeta"] {
for property in ["normalizedSegments", "summaryJson"] {
mark_property_nullable(&mut schema, property);
}
mark_definition_property_nullable(&mut schema, "MeetingSummaryActionItem", "owner");
@@ -105,7 +104,6 @@ pub(crate) fn transcript_result_schema() -> Value {
"sliceManifest",
"normalizedSegments",
"summaryJson",
"providerMeta",
] {
mark_property_nullable(&mut schema, property);
}
@@ -118,12 +116,6 @@ pub(crate) fn transcript_result_schema() -> Value {
fn schema_by_name(name: &str) -> Option<Value> {
match name {
// runtime-owned temporary native facade
"executionPlan" => Some(generated_schema_for::<llm_runtime::SerializableExecutionPlan>()),
// adapter-owned temporary native facade
"preparedRoutes" => Some(generated_schema_for::<
Vec<llm_adapter::router::SerializablePreparedRoute>,
>()),
// AFFiNE-native-owned N-API projection over adapter model registry/matcher
"capabilityMatchRequest" => Some(generated_schema_for::<CapabilityMatchRequest>()),
"capabilityMatchResponse" => Some(generated_schema_for::<CapabilityMatchResponse>()),
@@ -133,11 +125,6 @@ fn schema_by_name(name: &str) -> Option<Value> {
"modelRegistryResolveRequest" => Some(generated_schema_for::<ModelRegistryResolveRequest>()),
"modelRegistryResolveResponse" => Some(generated_schema_for::<ModelRegistryResolveResponse>()),
"providerDriverSpec" => Some(generated_schema_for::<ProviderDriverSpec>()),
// AFFiNE-native-owned prompt facade over adapter prompt DTOs/catalog
"promptRenderContract" => Some(generated_schema_for::<PromptRenderContract>()),
"promptSessionContract" => Some(generated_schema_for::<PromptSessionContract>()),
"requestedModelMatchRequest" => Some(generated_schema_for::<RequestedModelMatchRequest>()),
"requestedModelMatchResponse" => Some(generated_schema_for::<RequestedModelMatchResponse>()),
// runtime-owned
"toolCallbackRequest" => Some(generated_schema_for::<llm_runtime::ToolCallbackRequest>()),
"toolCallbackResponse" => Some(generated_schema_for::<llm_runtime::ToolCallbackResponse>()),
@@ -176,23 +163,6 @@ pub fn llm_validate_contract(name: String, value: Value) -> Result<Value> {
)))
}
#[napi(catch_unwind)]
pub fn llm_compile_execution_plan(value: Value) -> Result<Value> {
let value = llm_validate_contract("executionPlan".to_string(), value)?;
llm_runtime::compile_execution_plan_value(value.clone()).map_err(|error| invalid_contract(error.to_string()))?;
Ok(value)
}
#[napi(catch_unwind)]
pub fn llm_normalize_prepared_routes(value: Value) -> Result<Value> {
let value = llm_adapter::router::normalize_prepared_routes(value).map_err(|error| {
invalid_contract(format!(
"LLM prepared routes value does not match adapter contract: {error}"
))
})?;
llm_validate_contract("preparedRoutes".to_string(), value)
}
#[cfg(test)]
mod tests {
use serde_json::json;
@@ -220,8 +190,7 @@ mod tests {
"decisions": [],
"openQuestions": [],
"blockers": []
},
"providerMeta": { "provider": "gemini" }
}
});
assert!(llm_validate_contract("transcriptGeneratedResult".to_string(), value).is_ok());
}
@@ -234,7 +203,6 @@ mod tests {
"normalizedSegments": null,
"normalizedTranscript": "",
"summaryJson": null,
"providerMeta": null,
"extra": true
}),
)
@@ -242,24 +210,6 @@ mod tests {
assert!(error.reason.contains("does not match schema"));
}
#[test]
fn compiles_execution_plan_contract() {
let value = json!({
"routes": [{
"providerId": "openai-main",
"protocol": "openai_chat",
"model": "gpt-5-mini",
"backendConfig": { "base_url": "https://api.openai.com/v1", "auth_token": "token" }
}],
"request": { "kind": "text", "cond": { "modelId": "gpt-5-mini" }, "messages": [] },
"routePolicy": { "fallbackOrder": ["openai-main"] },
"runtimePolicy": {},
"attachmentPolicy": { "materializeRemoteAttachments": true },
"responsePostprocess": { "mode": "text" }
});
assert!(super::llm_compile_execution_plan(value).is_ok());
}
#[test]
fn validates_runtime_tool_callback_contracts() {
assert!(
@@ -288,100 +238,4 @@ mod tests {
.unwrap_err();
assert!(error.reason.contains("does not match schema"));
}
#[test]
fn validates_prompt_contracts_from_native_types() {
assert!(
llm_validate_contract(
"promptRenderContract".to_string(),
json!({
"messages": [{ "role": "user", "content": "hello" }],
"templateParams": {},
"renderParams": {}
}),
)
.is_ok()
);
assert!(
llm_validate_contract(
"promptSessionContract".to_string(),
json!({
"prompt": {
"promptTokens": 1,
"templateParams": {},
"messages": [{ "role": "system", "content": "hello" }]
},
"turns": [],
"renderParams": {},
"maxTokenSize": 1000
}),
)
.is_ok()
);
}
#[test]
fn validates_adapter_prepared_route_contract() {
assert!(
super::llm_normalize_prepared_routes(json!([
{
"provider_id": "openai-main",
"protocol": "openai_chat",
"model": "gpt-5-mini",
"config": {
"base_url": "https://api.openai.com/v1",
"auth_token": "token"
},
"request": {
"model": "gpt-5-mini",
"messages": []
}
}
]))
.is_ok()
);
let error = super::llm_normalize_prepared_routes(json!([
{
"provider_id": "openai-main",
"protocol": "openai_chat",
"model": "gpt-5-mini",
"config": { "base_url": "https://api.openai.com/v1" },
"request": {}
}
]))
.unwrap_err();
assert!(error.reason.contains("adapter contract"));
}
#[test]
fn execution_plan_rejects_host_only_state() {
let value = json!({
"routes": [],
"request": {
"kind": "text",
"cond": { "modelId": "gpt-5-mini" },
"messages": [],
"options": { "signal": {} }
},
"routePolicy": { "fallbackOrder": [] },
"runtimePolicy": {},
"attachmentPolicy": { "materializeRemoteAttachments": true },
"responsePostprocess": { "mode": "text" }
});
let error = super::llm_compile_execution_plan(value).unwrap_err();
assert!(error.reason.contains("request.options.signal"));
let value = json!({
"routes": [],
"request": { "kind": "text", "cond": { "modelId": "gpt-5-mini" }, "messages": [] },
"routePolicy": { "fallbackOrder": [] },
"runtimePolicy": {},
"attachmentPolicy": { "materializeRemoteAttachments": true },
"responsePostprocess": { "mode": "text" },
"hostContext": { "signal": {} }
});
let error = super::llm_compile_execution_plan(value).unwrap_err();
assert!(error.reason.contains("does not match schema"));
}
}
packages/backend/native/src/llm/core/capability.rs
+2 -22
View File
@@ -1,8 +1,6 @@
use napi::Result;
use crate::llm::core::contracts::{
CapabilityMatchRequest, CapabilityMatchResponse, RequestedModelMatchRequest, RequestedModelMatchResponse,
};
use crate::llm::core::contracts::{CapabilityMatchRequest, CapabilityMatchResponse};
#[napi(catch_unwind)]
pub fn llm_match_model_capabilities(payload: CapabilityMatchRequest) -> Result<CapabilityMatchResponse> {
@@ -14,25 +12,7 @@ pub fn llm_match_model_capabilities(payload: CapabilityMatchRequest) -> Result<C
.map_err(crate::llm::map_json_error)?;
Ok(CapabilityMatchResponse {
model_id: llm_adapter::core::select_model_id(&models, &cond).map_err(crate::llm::host::invalid_arg)?,
})
}
#[napi(catch_unwind)]
pub fn llm_resolve_requested_model_match(payload: RequestedModelMatchRequest) -> Result<RequestedModelMatchResponse> {
let matched_optional_model = llm_adapter::core::matches_requested_model_list(
&payload.provider_ids,
&payload.optional_models,
payload.requested_model_id.as_deref(),
);
Ok(RequestedModelMatchResponse {
selected_model: if matched_optional_model {
payload.requested_model_id
} else {
payload.default_model
},
matched_optional_model,
model_id: llm_adapter::core::select_model_id(&models, &cond).map_err(crate::llm::invalid_arg)?,
})
}
packages/backend/native/src/llm/core/contracts/mod.rs
+1 -131
View File
@@ -8,17 +8,6 @@ use schemars::JsonSchema;
use serde::{Deserialize, Serialize};
use serde_json::Value;
#[napi(object)]
#[derive(Debug, Clone, Deserialize, Serialize, PartialEq, JsonSchema)]
#[serde(rename_all = "camelCase")]
pub struct PromptRenderContract {
pub messages: Vec<PromptMessageContract>,
#[napi(ts_type = "Record<string, any>")]
pub template_params: Value,
#[napi(ts_type = "Record<string, any>")]
pub render_params: Value,
}
#[napi(object)]
#[derive(Debug, Clone, Deserialize, Serialize, PartialEq, JsonSchema)]
pub struct PromptRenderResult {
@@ -35,66 +24,13 @@ pub struct BuiltInPromptRenderContract {
pub render_params: Value,
}
#[napi(object)]
#[derive(Debug, Clone, Deserialize, Serialize, PartialEq, JsonSchema)]
pub struct PromptTokenCountContract {
#[serde(skip_serializing_if = "Option::is_none")]
pub model: Option<String>,
pub messages: Vec<PromptCountMessage>,
}
#[napi(object)]
#[derive(Debug, Clone, Deserialize, Serialize, PartialEq, JsonSchema)]
pub struct PromptTokenCountResult {
pub tokens: u32,
}
#[napi(object)]
#[derive(Debug, Clone, Deserialize, Serialize, PartialEq, JsonSchema)]
pub struct PromptCountMessage {
pub content: String,
}
#[napi(object)]
#[derive(Debug, Clone, Deserialize, Serialize, PartialEq)]
pub struct PromptMetadataContract {
pub messages: Vec<PromptMessageContract>,
}
#[napi(object)]
#[derive(Debug, Clone, Deserialize, Serialize, PartialEq)]
#[serde(rename_all = "camelCase")]
pub struct PromptMetadataResult {
pub param_keys: Vec<String>,
#[napi(ts_type = "Record<string, any>")]
pub template_params: Value,
}
#[napi(object)]
#[derive(Debug, Clone, Deserialize, Serialize, PartialEq, JsonSchema)]
#[serde(rename_all = "camelCase")]
pub struct PromptSessionContract {
pub prompt: PromptSessionPrompt,
pub turns: Vec<PromptMessageContract>,
#[napi(ts_type = "Record<string, any>")]
pub render_params: Value,
pub max_token_size: u32,
}
#[napi(object)]
#[derive(Debug, Clone, Deserialize, Serialize, PartialEq, JsonSchema)]
#[serde(rename_all = "camelCase")]
pub struct PromptSessionPrompt {
#[serde(skip_serializing_if = "Option::is_none")]
pub action: Option<String>,
#[serde(skip_serializing_if = "Option::is_none")]
pub model: Option<String>,
pub prompt_tokens: u32,
#[napi(ts_type = "Record<string, any>")]
pub template_params: Value,
pub messages: Vec<PromptMessageContract>,
}
#[napi(object)]
#[derive(Debug, Clone, Deserialize, Serialize, PartialEq)]
#[serde(rename_all = "camelCase")]
@@ -112,7 +48,6 @@ pub struct BuiltInPromptSessionContract {
pub turns: Vec<PromptMessageContract>,
#[napi(ts_type = "Record<string, any>")]
pub render_params: Value,
pub max_token_size: u32,
}
#[napi(object)]
@@ -289,29 +224,6 @@ pub struct CapabilityMatchResponse {
pub model_id: Option<String>,
}
#[napi(object)]
#[derive(Debug, Clone, Deserialize, JsonSchema, Serialize, PartialEq)]
#[serde(rename_all = "camelCase")]
#[serde(deny_unknown_fields)]
pub struct RequestedModelMatchRequest {
pub provider_ids: Vec<String>,
pub optional_models: Vec<String>,
#[serde(skip_serializing_if = "Option::is_none")]
pub requested_model_id: Option<String>,
#[serde(skip_serializing_if = "Option::is_none")]
pub default_model: Option<String>,
}
#[napi(object)]
#[derive(Debug, Clone, Deserialize, JsonSchema, Serialize, PartialEq)]
#[serde(rename_all = "camelCase")]
#[serde(deny_unknown_fields)]
pub struct RequestedModelMatchResponse {
#[serde(skip_serializing_if = "Option::is_none")]
pub selected_model: Option<String>,
pub matched_optional_model: bool,
}
#[napi(object)]
#[derive(Debug, Clone, Deserialize, JsonSchema, Serialize, PartialEq)]
#[serde(rename_all = "camelCase")]
@@ -636,8 +548,6 @@ pub struct LlmImageRequestContract {
#[serde(rename_all = "camelCase")]
pub struct LlmImageRequestBuildContract {
pub model: String,
#[napi(ts_type = "'openai_chat' | 'openai_responses' | 'openai_images' | 'anthropic' | 'gemini' | 'fal_image'")]
pub protocol: String,
pub messages: Vec<PromptMessageContract>,
#[serde(skip_serializing_if = "Option::is_none")]
pub options: Option<Value>,
@@ -647,47 +557,7 @@ pub struct LlmImageRequestBuildContract {
mod tests {
use serde_json::json;
use super::{CapabilityMatchRequest, PromptRenderContract, PromptSessionContract, ProviderDriverSpec};
#[test]
fn should_roundtrip_prompt_contracts() {
let render_value = json!({
"messages": [{
"role": "system",
"content": "summarize",
"responseFormat": {
"type": "json_schema",
"responseSchemaJson": {
"type": "object",
"properties": {
"summary": { "type": "string" }
},
"required": ["summary"]
},
"schemaHash": "abc123"
}
}],
"templateParams": { "tone": "short" },
"renderParams": { "topic": "docs" }
});
let session_value = json!({
"prompt": {
"model": "gpt-5-mini",
"promptTokens": 12,
"templateParams": {},
"messages": [{ "role": "system", "content": "summarize" }]
},
"turns": [{ "role": "user", "content": "hello" }],
"renderParams": { "tone": "short" },
"maxTokenSize": 1024
});
let render_contract: PromptRenderContract = serde_json::from_value(render_value.clone()).unwrap();
let session_contract: PromptSessionContract = serde_json::from_value(session_value.clone()).unwrap();
assert_eq!(serde_json::to_value(render_contract).unwrap(), render_value);
assert_eq!(serde_json::to_value(session_contract).unwrap(), session_value);
}
use super::{CapabilityMatchRequest, ProviderDriverSpec};
#[test]
fn should_roundtrip_tool_and_runtime_contracts() {
packages/backend/native/src/llm/core/model_registry.rs
+2 -2
View File
@@ -21,7 +21,7 @@ pub fn llm_resolve_model_registry_variant(
request.backend_kind.as_deref(),
request.model_id.as_str(),
)
.map_err(crate::llm::host::invalid_arg)?
.map_err(crate::llm::invalid_arg)?
{
Some((variant, matched_by)) => ModelRegistryResolveResponse {
variant: Some(to_contract_variant(variant)?),
@@ -44,7 +44,7 @@ pub fn llm_match_model_registry(request: ModelRegistryMatchRequest) -> Result<Mo
.map_err(crate::llm::map_json_error)?;
let response = ModelRegistryMatchResponse {
variant: llm_adapter::core::select_model_registry_variant(&variants, request.backend_kind.as_str(), &cond)
.map_err(crate::llm::host::invalid_arg)?
.map_err(crate::llm::invalid_arg)?
.map(to_contract_variant)
.transpose()?,
};
packages/backend/native/src/llm/core/prompt/mod.rs
+15 -365
View File
@@ -1,16 +1,12 @@
use napi::{Error, Result, Status};
use serde_json::{Map, Value};
use crate::{
llm::{
core::contracts::{
BuiltInPromptRenderContract, BuiltInPromptSessionContract, PromptMessageContract, PromptMetadataContract,
PromptMetadataResult, PromptRenderContract, PromptRenderResult, PromptSessionContract, PromptSessionPrompt,
PromptSessionResult, PromptTokenCountContract, PromptTokenCountResult,
},
prompt_catalog::{BuiltInPrompt, BuiltInPromptSpec, built_in_prompt, built_in_prompt_spec, built_in_prompt_specs},
use crate::llm::{
core::contracts::{
BuiltInPromptRenderContract, BuiltInPromptSessionContract, PromptMessageContract, PromptMetadataResult,
PromptRenderResult, PromptSessionResult,
},
tiktoken::{Tokenizer, from_model_name},
prompt_catalog::{BuiltInPrompt, BuiltInPromptSpec, built_in_prompt, built_in_prompt_spec, built_in_prompt_specs},
};
mod metadata;
@@ -51,57 +47,6 @@ fn built_in_prompt_metadata(prompt: &BuiltInPrompt) -> Result<PromptMetadataResu
.map_err(|error| invalid_arg(format!("Failed to collect built-in prompt metadata: {error}")))
}
fn count_prompt_tokens(model: Option<&str>, messages: &[PromptMessageContract]) -> u32 {
let content = messages
.iter()
.map(|message| message.content.as_str())
.collect::<String>();
prompt_tokenizer(model)
.map(|tokenizer| tokenizer.count(content, None))
.unwrap_or(0)
}
fn prompt_tokenizer(model: Option<&str>) -> Option<Tokenizer> {
let model = model?;
if model.starts_with("gpt") {
return from_model_name(model.to_string());
}
if model.starts_with("dall") {
return None;
}
from_model_name("gpt-4".to_string())
}
#[napi(catch_unwind)]
pub fn llm_render_prompt(request: PromptRenderContract) -> Result<PromptRenderResult> {
let response = render_prompt_response(
&request.messages,
&value_to_map(request.template_params, "templateParams")?,
&value_to_map(request.render_params, "renderParams")?,
)
.map_err(|error| invalid_arg(format!("Failed to render prompt: {error}")))?;
Ok(response)
}
#[napi(catch_unwind)]
pub fn llm_count_prompt_tokens(request: PromptTokenCountContract) -> Result<PromptTokenCountResult> {
let content = request
.messages
.iter()
.map(|message| message.content.as_str())
.collect::<String>();
let tokens = request
.model
.as_deref()
.and_then(|model| prompt_tokenizer(Some(model)))
.map(|tokenizer| tokenizer.count(content, None))
.unwrap_or(0);
Ok(PromptTokenCountResult { tokens })
}
#[napi(catch_unwind)]
pub fn llm_render_built_in_prompt(request: BuiltInPromptRenderContract) -> Result<PromptRenderResult> {
let prompt = built_in_prompt(&request.name)
@@ -118,46 +63,22 @@ pub fn llm_render_built_in_prompt(request: BuiltInPromptRenderContract) -> Resul
Ok(response)
}
#[napi(catch_unwind)]
pub fn llm_collect_prompt_metadata(request: PromptMetadataContract) -> Result<PromptMetadataResult> {
let response = collect_prompt_metadata(&request.messages)
.map_err(|error| invalid_arg(format!("Failed to collect prompt metadata: {error}")))?;
Ok(response)
}
#[napi(catch_unwind)]
pub fn llm_render_session_prompt(request: PromptSessionContract) -> Result<PromptSessionResult> {
let template_params = value_to_map(request.prompt.template_params.clone(), "prompt.templateParams")?;
let render_params = value_to_map(request.render_params.clone(), "renderParams")?;
let response = render_session_prompt(&request, &template_params, &render_params)
.map_err(|error| invalid_arg(format!("Failed to render session prompt: {error}")))?;
Ok(response)
}
#[napi(catch_unwind)]
pub fn llm_render_built_in_session_prompt(request: BuiltInPromptSessionContract) -> Result<PromptSessionResult> {
let prompt = built_in_prompt(&request.name)
.ok_or_else(|| invalid_arg(format!("Built-in prompt not found: {}", request.name)))?;
let messages = built_in_prompt_messages(prompt);
let metadata = built_in_prompt_metadata(prompt)?;
let session_contract = PromptSessionContract {
prompt: PromptSessionPrompt {
action: prompt.action.clone(),
model: Some(prompt.model.clone()),
prompt_tokens: count_prompt_tokens(Some(prompt.model.as_str()), &messages),
template_params: metadata.template_params,
messages,
},
turns: request.turns,
render_params: request.render_params,
max_token_size: request.max_token_size,
};
let template_params = value_to_map(session_contract.prompt.template_params.clone(), "prompt.templateParams")?;
let render_params = value_to_map(session_contract.render_params.clone(), "renderParams")?;
let response = render_session_prompt(&session_contract, &template_params, &render_params)
.map_err(|error| invalid_arg(format!("Failed to render built-in session prompt: {error}")))?;
let template_params = value_to_map(metadata.template_params, "prompt.templateParams")?;
let render_params = value_to_map(request.render_params, "renderParams")?;
let response = render_session_prompt(
&messages,
prompt.action.as_deref(),
&request.turns,
&template_params,
&render_params,
)
.map_err(|error| invalid_arg(format!("Failed to render built-in session prompt: {error}")))?;
Ok(response)
}
@@ -171,274 +92,3 @@ pub fn llm_list_built_in_prompt_specs() -> Result<Vec<BuiltInPromptSpec>> {
pub fn llm_get_built_in_prompt_spec(name: String) -> Result<Option<BuiltInPromptSpec>> {
Ok(built_in_prompt_spec(&name).cloned())
}
#[cfg(test)]
mod tests {
use llm_adapter::core::prompt_template::{is_truthy_number, parse_template, render_tokens};
use serde_json::json;
use super::{llm_collect_prompt_metadata, llm_count_prompt_tokens, llm_render_prompt, llm_render_session_prompt};
use crate::llm::core::contracts::{
PromptMetadataContract, PromptRenderContract, PromptSessionContract, PromptTokenCountContract,
};
#[test]
fn should_render_sections_and_current_item() {
let tokens = parse_template("{{#links}}- {{.}}\n{{/links}}").unwrap();
let rendered = render_tokens(
&tokens,
&[&json!({
"links": ["https://affine.pro", "https://github.com/toeverything/affine"]
})],
);
assert_eq!(
rendered,
"- https://affine.pro\n- https://github.com/toeverything/affine\n"
);
}
#[test]
fn should_render_prompt_with_normalized_params_and_attachments() {
let response = llm_render_prompt(
serde_json::from_value::<PromptRenderContract>(json!({
"messages": [
{
"role": "system",
"content": "tone={{tone}}"
},
{
"role": "user",
"content": "{{content}}"
}
],
"templateParams": { "tone": ["formal", "casual"] },
"renderParams": {
"attachments": ["https://affine.pro/example.jpg"],
"content": "hello world"
}
}))
.unwrap(),
)
.unwrap();
let response = serde_json::to_value(response).unwrap();
assert_eq!(
response,
json!({
"messages": [
{
"role": "system",
"content": "tone=formal",
"params": {
"attachments": ["https://affine.pro/example.jpg"],
"content": "hello world",
"tone": "formal"
}
},
{
"role": "user",
"content": "hello world",
"attachments": ["https://affine.pro/example.jpg"],
"params": {
"attachments": ["https://affine.pro/example.jpg"],
"content": "hello world",
"tone": "formal"
}
}
],
"warnings": ["Missing param value: tone, use default options: formal"]
}),
);
}
#[test]
fn should_render_host_builtins_and_js_like_variable_strings() {
let response = llm_render_prompt(
serde_json::from_value::<PromptRenderContract>(json!({
"messages": [
{
"role": "system",
"content": "{{affine::language}}|{{tags}}|{{obj}}|{{#links}}- {{.}}\n{{/links}}"
}
],
"templateParams": {},
"renderParams": {
"language": "French",
"affine::language": "ignored",
"links": ["https://affine.pro", "https://github.com/toeverything/affine"],
"obj": { "hello": "world" },
"tags": ["a", "b"]
}
}))
.unwrap(),
)
.unwrap();
let response = serde_json::to_value(response).unwrap();
assert_eq!(
response,
json!({
"messages": [
{
"role": "system",
"content": "French|a,b|[object Object]|- https://affine.pro\n- https://github.com/toeverything/affine\n",
"params": {
"language": "French",
"affine::language": "ignored",
"links": ["https://affine.pro", "https://github.com/toeverything/affine"],
"obj": { "hello": "world" },
"tags": ["a", "b"]
}
}
],
"warnings": []
}),
);
}
#[test]
fn should_count_prompt_tokens_for_unknown_models_as_zero() {
let response = llm_count_prompt_tokens(
serde_json::from_value::<PromptTokenCountContract>(json!({
"model": null,
"messages": [{ "content": "hello" }]
}))
.unwrap(),
)
.unwrap();
let response = serde_json::to_value(response).unwrap();
assert_eq!(response, json!({ "tokens": 0 }));
}
#[test]
fn should_count_prompt_tokens_for_non_gpt_models_with_fallback_tokenizer() {
let response = llm_count_prompt_tokens(
serde_json::from_value::<PromptTokenCountContract>(json!({
"model": "claude-3-5-sonnet",
"messages": [{ "content": "hello" }]
}))
.unwrap(),
)
.unwrap();
assert!(response.tokens > 0);
}
#[test]
fn should_follow_js_truthiness_for_numbers() {
assert!(!is_truthy_number(&serde_json::Number::from(0)));
assert!(is_truthy_number(&serde_json::Number::from(1)));
assert!(is_truthy_number(&serde_json::Number::from_f64(0.5).unwrap()));
}
#[test]
fn should_render_session_prompt_by_merging_latest_user_content() {
let response = llm_render_session_prompt(
serde_json::from_value::<PromptSessionContract>(json!({
"prompt": {
"model": "test",
"promptTokens": 0,
"templateParams": {},
"messages": [
{ "role": "system", "content": "answer briefly" },
{ "role": "user", "content": "{{content}}" }
]
},
"turns": [
{ "role": "user", "content": "hello", "attachments": ["https://affine.pro/hello.png"] }
],
"renderParams": {},
"maxTokenSize": 1000
}))
.unwrap(),
)
.unwrap();
let response = serde_json::to_value(response).unwrap();
assert_eq!(
response,
json!({
"messages": [
{ "role": "system", "content": "answer briefly", "params": { "content": "hello" } },
{
"role": "user",
"content": "hello",
"attachments": ["https://affine.pro/hello.png"],
"params": { "content": "hello" }
}
],
"warnings": [],
"promptMessagePositions": [0, 1]
}),
);
}
#[test]
fn should_render_session_prompt_by_picking_recent_turns_under_budget() {
let response = llm_render_session_prompt(
serde_json::from_value::<PromptSessionContract>(json!({
"prompt": {
"model": "test",
"promptTokens": 0,
"templateParams": {},
"messages": [
{ "role": "system", "content": "hello {{word}}" }
]
},
"turns": [
{ "role": "user", "content": "older turn" }
],
"renderParams": { "word": "world" },
"maxTokenSize": 0
}))
.unwrap(),
)
.unwrap();
let response = serde_json::to_value(response).unwrap();
assert_eq!(
response,
json!({
"messages": [
{ "role": "system", "content": "hello world", "params": { "word": "world" } }
],
"warnings": [],
"promptMessagePositions": [0]
}),
);
}
#[test]
fn should_collect_prompt_metadata_from_templates_and_params() {
let response = llm_collect_prompt_metadata(
serde_json::from_value::<PromptMetadataContract>(json!({
"messages": [
{
"role": "system",
"content": "tone={{tone}}"
},
{
"role": "user",
"content": "{{content}}",
"params": { "tone": ["formal", "casual"] }
}
]
}))
.unwrap(),
)
.unwrap();
let response = serde_json::to_value(response).unwrap();
assert_eq!(
response,
json!({
"paramKeys": ["tone", "content"],
"templateParams": {
"tone": ["formal", "casual"]
}
}),
);
}
}
packages/backend/native/src/llm/core/prompt/render.rs
+65
View File
@@ -156,3 +156,68 @@ fn render_prompt_message(
Ok(next)
}
#[cfg(test)]
mod tests {
use serde_json::json;
use super::*;
fn message(role: &str, content: &str) -> PromptMessageContract {
serde_json::from_value(json!({ "role": role, "content": content })).unwrap()
}
#[test]
fn renders_lists_and_normalizes_missing_or_invalid_params_to_defaults() {
let messages = vec![
message("system", "translate {{src}} to {{dest}}: {{content}}"),
message("user", "links:\n{{#links}}- {{.}}\n{{/links}}"),
];
let template_params = serde_json::from_value(json!({
"src": ["eng"],
"dest": ["chs", "jpn"]
}))
.unwrap();
let params = serde_json::from_value(json!({
"src": "invalid",
"content": "hello",
"links": ["https://affine.pro", "https://github.com/toeverything/AFFiNE"]
}))
.unwrap();
let rendered = render_prompt_response(&messages, &template_params, &params).unwrap();
assert_eq!(rendered.messages[0].content, "translate eng to chs: hello");
assert_eq!(
rendered.messages[1].content,
"links:\n- https://affine.pro\n- https://github.com/toeverything/AFFiNE\n"
);
assert_eq!(rendered.warnings.len(), 2);
}
#[test]
fn appends_input_attachments_only_to_user_messages() {
let messages = vec![message("system", "system"), message("user", "{{content}}")];
let params = serde_json::from_value(json!({
"content": "summarize",
"attachments": [{
"kind": "file_handle",
"fileHandle": "file-1",
"mimeType": "application/pdf"
}]
}))
.unwrap();
let rendered = render_prompt_response(&messages, &Map::new(), &params).unwrap();
assert!(rendered.messages[0].attachments.is_none());
assert_eq!(
rendered.messages[1].attachments,
Some(vec![json!({
"kind": "file_handle",
"fileHandle": "file-1",
"mimeType": "application/pdf"
})])
);
}
}
packages/backend/native/src/llm/core/prompt/session.rs
+184 -39
View File
@@ -2,22 +2,45 @@ use llm_adapter::core::prompt_template::{parse_template, template_uses_key};
use serde_json::{Map, Value};
use super::{
super::contracts::{PromptMessageContract, PromptSessionContract, PromptSessionResult},
super::contracts::{PromptMessageContract, PromptSessionResult},
render::render_prompt_response,
};
use crate::tiktoken::{Tokenizer, from_model_name};
const DEFAULT_HISTORY_INPUT_BYTES: usize = 128 * 1024;
const MESSAGE_FRAMING_BYTES: usize = 16;
pub(super) fn render_session_prompt(
request: &PromptSessionContract,
prompt_messages: &[PromptMessageContract],
action: Option<&str>,
turns: &[PromptMessageContract],
template_params: &Map<String, Value>,
params: &Map<String, Value>,
) -> std::result::Result<PromptSessionResult, String> {
let tokenizer = session_tokenizer(request.prompt.model.as_deref());
let mut selected_turns = take_session_turns(request, tokenizer.as_ref())?;
let latest_turn = selected_turns.pop();
render_session_prompt_with_budget(
prompt_messages,
action,
turns,
template_params,
params,
DEFAULT_HISTORY_INPUT_BYTES,
)
}
if prompt_uses_content(&request.prompt.messages)?
&& !selected_turns.iter().any(message_is_assistant)
fn render_session_prompt_with_budget(
prompt_messages: &[PromptMessageContract],
action: Option<&str>,
turns: &[PromptMessageContract],
template_params: &Map<String, Value>,
params: &Map<String, Value>,
history_input_bytes: usize,
) -> std::result::Result<PromptSessionResult, String> {
let (prior_turns, latest_turn) = turns
.split_last()
.map(|(latest, prior)| (prior, Some(latest.clone())))
.unwrap_or((&[], None));
if prompt_uses_content(prompt_messages)?
&& !prior_turns.iter().any(message_is_assistant)
&& let Some(last_message) = latest_turn
.as_ref()
.filter(|message| message_role(message) == Some("user"))
@@ -29,16 +52,17 @@ pub(super) fn render_session_prompt(
}
merged_params.insert("content".to_string(), Value::String(last_message.content.clone()));
let rendered = render_prompt_response(&request.prompt.messages, template_params, &merged_params)?;
let rendered = render_prompt_response(prompt_messages, template_params, &merged_params)?;
let mut messages = rendered.messages;
let Some(first_user_message_index) = messages
.iter()
.position(|message| message_role(message) == Some("user"))
else {
ensure_messages_fit(&messages, &[], history_input_bytes)?;
return Ok(PromptSessionResult {
messages,
warnings: rendered.warnings,
prompt_message_positions: (0..request.prompt.messages.len()).map(|index| index as u32).collect(),
prompt_message_positions: (0..prompt_messages.len()).map(|index| index as u32).collect(),
});
};
@@ -57,9 +81,10 @@ pub(super) fn render_session_prompt(
messages[first_user_message_index].attachments = Some(merged_attachments);
}
let selected_turns = select_history_turns(&messages, prior_turns, history_input_bytes)?;
let prior_turn_count = selected_turns.len();
messages.splice(first_user_message_index..first_user_message_index, selected_turns);
let prompt_message_positions = (0..request.prompt.messages.len())
let prompt_message_positions = (0..prompt_messages.len())
.map(|index| {
if index < first_user_message_index {
index as u32
@@ -81,52 +106,63 @@ pub(super) fn render_session_prompt(
} else {
latest_turn.as_ref().map(message_params).unwrap_or_default()
};
let rendered = render_prompt_response(&request.prompt.messages, template_params, &final_params)?;
let rendered = render_prompt_response(prompt_messages, template_params, &final_params)?;
let trailing_turns = selected_turns
let latest_turns = latest_turn
.into_iter()
.chain(latest_turn)
.filter(prompt_message_should_survive)
.collect::<Vec<_>>();
let mut messages = rendered.messages;
messages.extend(trailing_turns);
let selected_turns = if action.is_some() {
ensure_messages_fit(&messages, &latest_turns, history_input_bytes)?;
Vec::new()
} else {
let mut fixed = messages.clone();
fixed.extend(latest_turns.clone());
select_history_turns(&fixed, prior_turns, history_input_bytes)?
};
messages.extend(selected_turns);
messages.extend(latest_turns);
Ok(PromptSessionResult {
messages,
warnings: rendered.warnings,
prompt_message_positions: (0..request.prompt.messages.len()).map(|index| index as u32).collect(),
prompt_message_positions: (0..prompt_messages.len()).map(|index| index as u32).collect(),
})
}
fn session_tokenizer(model: Option<&str>) -> Option<Tokenizer> {
let model = model?;
if model.starts_with("gpt") {
return from_model_name(model.to_string());
fn estimated_message_bytes(message: &PromptMessageContract) -> usize {
let mut size = MESSAGE_FRAMING_BYTES
.saturating_add(message.role.len())
.saturating_add(message.content.len());
for value in [
message.attachments.as_ref().map(serde_json::to_vec),
message.params.as_ref().map(serde_json::to_vec),
message.response_format.as_ref().map(serde_json::to_vec),
]
.into_iter()
.flatten()
{
size = size.saturating_add(value.map(|bytes| bytes.len()).unwrap_or(usize::MAX));
}
if model.starts_with("dall") {
return None;
}
from_model_name("gpt-4".to_string())
size
}
fn take_session_turns(
request: &PromptSessionContract,
tokenizer: Option<&Tokenizer>,
fn select_history_turns(
fixed_messages: &[PromptMessageContract],
history: &[PromptMessageContract],
history_input_bytes: usize,
) -> std::result::Result<Vec<PromptMessageContract>, String> {
if request.prompt.action.is_some() {
return Ok(request.turns.last().cloned().into_iter().collect());
let mut size = fixed_messages.iter().fold(0usize, |size, message| {
size.saturating_add(estimated_message_bytes(message))
});
if size > history_input_bytes {
return Err("session input exceeds history byte budget".to_string());
}
let mut picked = Vec::new();
let mut size = request.prompt.prompt_tokens;
for message in request.turns.iter().rev() {
let content = message.content.as_str();
size += tokenizer
.map(|tokenizer| tokenizer.count(content.to_string(), None))
.unwrap_or(0);
if size > request.max_token_size {
for message in history.iter().rev() {
size = size.saturating_add(estimated_message_bytes(message));
if size > history_input_bytes {
break;
}
picked.push(message.clone());
@@ -136,6 +172,16 @@ fn take_session_turns(
Ok(picked)
}
fn ensure_messages_fit(
prompt_messages: &[PromptMessageContract],
latest_turns: &[PromptMessageContract],
history_input_bytes: usize,
) -> std::result::Result<(), String> {
let mut messages = prompt_messages.to_vec();
messages.extend_from_slice(latest_turns);
select_history_turns(&messages, &[], history_input_bytes).map(|_| ())
}
fn prompt_uses_content(messages: &[PromptMessageContract]) -> std::result::Result<bool, String> {
for message in messages {
if template_uses_key(&parse_template(&message.content)?, "content") {
@@ -202,3 +248,102 @@ fn attachment_has_source(attachment: &Value) -> bool {
_ => false,
}
}
#[cfg(test)]
mod tests {
use serde_json::{Map, json};
use super::*;
fn message(role: &str, content: &str) -> PromptMessageContract {
serde_json::from_value(json!({ "role": role, "content": content })).unwrap()
}
#[test]
fn byte_estimate_is_model_independent_and_utf8_aware() {
let ascii = estimated_message_bytes(&message("user", "abc"));
let cjk = estimated_message_bytes(&message("user", "中文文"));
let emoji = estimated_message_bytes(&message("user", "😀😀😀"));
assert!(ascii < cjk);
assert!(cjk < emoji);
}
#[test]
fn keeps_latest_turn_and_only_a_contiguous_history_suffix() {
let prompt = vec![message("system", "system")];
let turns = vec![
message("user", "old"),
message("assistant", "recent"),
message("user", "latest"),
];
let fixed_bytes = estimated_message_bytes(&prompt[0]) + estimated_message_bytes(&turns[2]);
let budget = fixed_bytes + estimated_message_bytes(&turns[1]) + 8;
let result = render_session_prompt_with_budget(&prompt, None, &turns, &Map::new(), &Map::new(), budget).unwrap();
assert_eq!(
result
.messages
.iter()
.map(|message| message.content.as_str())
.collect::<Vec<_>>(),
["system", "recent", "latest"]
);
}
#[test]
fn rejects_oversized_rendered_prompt_or_latest_turn() {
let prompt = vec![message("system", "{{content}}")];
let turns = vec![message("user", "large input")];
let error = render_session_prompt_with_budget(&prompt, None, &turns, &Map::new(), &Map::new(), 1).unwrap_err();
assert_eq!(error, "session input exceeds history byte budget");
}
#[test]
fn merges_latest_user_content_params_and_file_handles_into_prompt() {
let prompt = vec![message("user", "{{content}} {{tone}}")];
let latest = serde_json::from_value(json!({
"role": "user",
"content": "Summarize this file",
"attachments": [{
"kind": "file_handle",
"fileHandle": "file-1",
"mimeType": "application/pdf"
}],
"params": { "tone": "brief" }
}))
.unwrap();
let result = render_session_prompt(&prompt, None, &[latest], &Map::new(), &Map::new()).unwrap();
assert_eq!(result.messages.len(), 1);
assert_eq!(result.messages[0].content, "Summarize this file brief");
assert_eq!(
result.messages[0].attachments,
Some(vec![json!({
"kind": "file_handle",
"fileHandle": "file-1",
"mimeType": "application/pdf"
})])
);
}
#[test]
fn action_prompt_drops_prior_history_but_keeps_latest_turn() {
let prompt = vec![message("system", "action")];
let turns = vec![
message("user", "old"),
message("assistant", "old answer"),
message("user", "latest"),
];
let result = render_session_prompt(&prompt, Some("edit"), &turns, &Map::new(), &Map::new()).unwrap();
assert_eq!(
result
.messages
.iter()
.map(|message| message.content.as_str())
.collect::<Vec<_>>(),
["action", "latest"]
);
}
}
packages/backend/native/src/llm/core/request_builder/mod.rs
+5 -37
View File
@@ -1,4 +1,4 @@
use llm_adapter::core::{self as adapter_core, EmbeddingRequest, ImageInput, ImageRequest, RerankRequest};
use llm_adapter::core::{self as adapter_core, EmbeddingRequest, ImageRequest, RerankRequest};
use napi::Result;
use napi_derive::napi;
use serde::Serialize;
@@ -8,7 +8,7 @@ use super::contracts::{
LlmImageRequestBuildContract, LlmImageRequestContract, LlmRequestContract, LlmRerankRequestContract,
LlmStructuredRequestContract, ModelConditionsContract, PromptMessageContract,
};
use crate::llm::{LlmDispatchPayload, LlmRerankDispatchPayload, LlmStructuredDispatchPayload, host::invalid_arg};
use crate::llm::{LlmDispatchPayload, LlmRerankDispatchPayload, LlmStructuredDispatchPayload, invalid_arg};
mod types;
@@ -62,38 +62,11 @@ pub(crate) fn build_image_request(request: ImageRequest) -> Result<ImageRequest>
}
pub(crate) fn build_image_request_from_messages(request: LlmImageRequestBuildContract) -> Result<ImageRequest> {
let protocol = request.protocol.clone();
let mut request =
let request =
adapter_core::build_image_request_from_prompt_messages(to_adapter(&request)?).map_err(map_builder_error)?;
if protocol == "fal_image" {
keep_fal_data_uri_inputs_as_urls(&mut request);
}
Ok(request)
}
fn keep_fal_data_uri_inputs_as_urls(request: &mut ImageRequest) {
let ImageRequest::Edit(edit) = request else {
return;
};
for image in &mut edit.images {
let replacement = match image {
ImageInput::Data {
data_base64,
media_type,
..
} => Some(ImageInput::Url {
url: format!("data:{media_type};base64,{data_base64}"),
media_type: Some(media_type.clone()),
}),
_ => None,
};
if let Some(replacement) = replacement {
*image = replacement;
}
}
}
pub(crate) fn infer_prompt_model_conditions(messages: Vec<PromptMessageInput>) -> Result<ModelConditionsContract> {
let messages = adapter_core::canonicalize_prompt_messages(to_adapter_prompt_messages(messages)?);
serde_json::to_value(adapter_core::infer_model_conditions_from_prompt_messages(messages))
@@ -471,11 +444,10 @@ mod tests {
}
#[test]
fn should_keep_fal_data_uri_image_inputs_as_urls() {
fn should_canonicalize_data_uri_image_inputs() {
let response = llm_build_image_request_from_messages(
serde_json::from_value(json!({
"model": "lora/image-to-image",
"protocol": "fal_image",
"messages": [{
"role": "user",
"content": "restyle",
@@ -493,11 +465,7 @@ mod tests {
let response = serde_json::to_value(response).unwrap();
assert_eq!(
response.pointer("/images/0"),
Some(&json!({
"kind": "url",
"url": "data:image/png;base64,aW1n",
"media_type": "image/png"
}))
Some(&json!({ "kind": "data", "data_base64": "aW1n", "media_type": "image/png" }))
);
}
packages/backend/native/src/llm/core/request_builder/types.rs
+1 -1
View File
@@ -15,7 +15,7 @@ use super::super::contracts::{
LlmRerankRequestContract, LlmStructuredRequestContract, RerankCandidate as ContractRerankCandidate, ToolContract,
};
use crate::llm::{
LlmDispatchPayload, LlmMiddlewarePayload, LlmRerankDispatchPayload, LlmStructuredDispatchPayload, host::invalid_arg,
LlmDispatchPayload, LlmMiddlewarePayload, LlmRerankDispatchPayload, LlmStructuredDispatchPayload, invalid_arg,
map_json_error,
};
packages/backend/native/src/llm/ffi/mod.rs
+1 -23
View File
@@ -1,27 +1,5 @@
mod dispatch;
mod middleware;
mod payload;
#[cfg(test)]
pub(crate) use dispatch::AsyncLlmDispatchPreparedTask;
pub(crate) use dispatch::{
dispatch_prepared_image_route_payloads, dispatch_prepared_structured_routes,
parse_prepared_chat_routes_with_middleware, parse_prepared_chat_routes_without_middleware,
};
pub use dispatch::{
llm_dispatch_prepared, llm_embedding_dispatch, llm_embedding_dispatch_prepared, llm_image_dispatch_prepared,
llm_plan_attachment_reference, llm_rerank_dispatch, llm_rerank_dispatch_prepared, llm_resolve_request_intent,
llm_structured_dispatch, llm_structured_dispatch_prepared,
};
pub(crate) use llm_adapter::middleware::StreamPipeline;
#[cfg(test)]
pub(crate) use middleware::resolve_request_chain;
pub(crate) use middleware::{
apply_request_middlewares, apply_structured_request_middlewares, backend_transport_error, map_backend_error,
map_json_error, parse_embedding_protocol, parse_protocol, parse_rerank_protocol, parse_structured_protocol,
resolve_stream_chain,
};
pub(crate) use payload::{
LlmDispatchPayload, LlmEmbeddingDispatchPayload, LlmMiddlewarePayload, LlmPreparedImageDispatchRoutePayload,
LlmRerankDispatchPayload, LlmRoutedBackendPayload, LlmStructuredDispatchPayload,
LlmDispatchPayload, LlmMiddlewarePayload, LlmRerankDispatchPayload, LlmStructuredDispatchPayload,
};
packages/backend/native/src/llm/ffi/payload.rs
+2 -155
View File
@@ -1,15 +1,10 @@
use llm_adapter::{
backend::BackendConfig,
core::{CoreRequest, EmbeddingRequest, RerankRequest, StructuredRequest},
core::{CoreRequest, RerankRequest, StructuredRequest},
middleware::MiddlewareConfig,
router::SerializablePreparedRoute,
};
use serde::{Deserialize, Serialize};
use crate::llm::core::contracts::{
LlmEmbeddingRequestContract, LlmImageRequestContract, LlmRequestContract, LlmRerankRequestContract,
LlmStructuredRequestContract,
};
use crate::llm::core::contracts::{LlmRequestContract, LlmRerankRequestContract, LlmStructuredRequestContract};
#[derive(Debug, Clone, Default, Deserialize, Serialize)]
#[serde(default)]
@@ -41,15 +36,6 @@ pub(crate) struct LlmDispatchPayload {
pub(crate) middleware: LlmMiddlewarePayload,
}
#[derive(Debug, Clone, Deserialize, Serialize)]
pub(crate) struct LlmRoutedBackendPayload {
pub(crate) provider_id: String,
pub(crate) protocol: String,
pub(crate) model: String,
#[serde(alias = "backendConfig")]
pub(crate) config: BackendConfig,
}
#[derive(Debug, Clone, Deserialize, Serialize)]
#[serde(try_from = "LlmStructuredRequestContract")]
pub(crate) struct LlmStructuredDispatchPayload {
@@ -59,20 +45,6 @@ pub(crate) struct LlmStructuredDispatchPayload {
pub(crate) middleware: LlmMiddlewarePayload,
}
#[derive(Debug, Clone, Deserialize, Serialize)]
#[serde(from = "LlmEmbeddingRequestContract")]
pub(crate) struct LlmEmbeddingDispatchPayload {
pub(crate) request: EmbeddingRequest,
}
impl From<LlmEmbeddingRequestContract> for LlmEmbeddingDispatchPayload {
fn from(request: LlmEmbeddingRequestContract) -> Self {
Self {
request: request.into(),
}
}
}
#[derive(Debug, Clone, Deserialize, Serialize)]
#[serde(from = "LlmRerankRequestContract")]
pub(crate) struct LlmRerankDispatchPayload {
@@ -87,128 +59,3 @@ impl From<LlmRerankRequestContract> for LlmRerankDispatchPayload {
}
}
}
pub(crate) type LlmPreparedImageDispatchRoutePayload = SerializablePreparedRoute<LlmImageRequestContract>;
#[cfg(test)]
mod tests {
use llm_adapter::router::SerializablePreparedRoute;
use super::{
LlmDispatchPayload, LlmPreparedImageDispatchRoutePayload, LlmRerankDispatchPayload, LlmStructuredDispatchPayload,
};
#[test]
fn prepared_chat_route_payload_deserializes_nested_request() {
let payload = serde_json::from_value::<Vec<SerializablePreparedRoute<LlmDispatchPayload>>>(serde_json::json!([
{
"provider_id": "openai-primary",
"protocol": "openai_chat",
"model": "gpt-5-mini",
"config": {
"base_url": "https://api.openai.com",
"auth_token": "test-key"
},
"request": {
"model": "gpt-5-mini",
"messages": [
{
"role": "user",
"content": [{ "type": "text", "text": "hello" }]
}
]
}
}
]))
.expect("prepared chat route payload should deserialize");
assert_eq!(payload[0].model, "gpt-5-mini");
assert_eq!(payload[0].request.request.model, "gpt-5-mini");
}
#[test]
fn prepared_structured_route_payload_deserializes_nested_request() {
let payload =
serde_json::from_value::<Vec<SerializablePreparedRoute<LlmStructuredDispatchPayload>>>(serde_json::json!([
{
"provider_id": "openai-primary",
"protocol": "openai_responses",
"model": "gpt-5-mini",
"config": {
"base_url": "https://api.openai.com",
"auth_token": "test-key"
},
"request": {
"model": "gpt-5-mini",
"messages": [
{
"role": "user",
"content": [{ "type": "text", "text": "hello" }]
}
],
"schema": {
"type": "object",
"properties": {
"summary": { "type": "string" }
},
"required": ["summary"]
}
}
}
]))
.expect("prepared structured route payload should deserialize");
assert_eq!(payload[0].model, "gpt-5-mini");
assert_eq!(payload[0].request.request.model, "gpt-5-mini");
}
#[test]
fn prepared_rerank_route_payload_deserializes_nested_request() {
let payload =
serde_json::from_value::<Vec<SerializablePreparedRoute<LlmRerankDispatchPayload>>>(serde_json::json!([
{
"provider_id": "openai-primary",
"protocol": "openai_chat",
"model": "gpt-5-mini",
"config": {
"base_url": "https://api.openai.com",
"auth_token": "test-key"
},
"request": {
"model": "gpt-5-mini",
"query": "hello",
"candidates": [{ "text": "world" }]
}
}
]))
.expect("prepared rerank route payload should deserialize");
assert_eq!(payload[0].model, "gpt-5-mini");
assert_eq!(payload[0].request.request.model, "gpt-5-mini");
}
#[test]
fn prepared_image_route_payload_deserializes_nested_request() {
let payload = serde_json::from_value::<Vec<LlmPreparedImageDispatchRoutePayload>>(serde_json::json!([
{
"provider_id": "openai-primary",
"protocol": "openai_images",
"model": "gpt-image-1",
"config": {
"base_url": "https://api.openai.com",
"auth_token": "test-key",
"request_layer": "openai_images"
},
"request": {
"model": "gpt-image-1",
"prompt": "draw",
"operation": "generate"
}
}
]))
.expect("prepared image route payload should deserialize");
assert_eq!(payload[0].model, "gpt-image-1");
assert_eq!(payload[0].request.prompt, "draw");
}
}
packages/backend/native/src/llm/mod.rs
+33 -30
View File
@@ -1,19 +1,17 @@
mod action;
pub(crate) mod byok;
mod contract_schema;
mod core;
mod ffi;
mod host;
mod prompt_catalog;
pub(crate) mod route;
#[cfg(test)]
mod tests;
pub use core::{
capability::{llm_match_model_capabilities, llm_resolve_requested_model_match},
pub use self::core::{
capability::llm_match_model_capabilities,
model_registry::{llm_match_model_registry, llm_resolve_model_registry_variant},
prompt::{
llm_collect_prompt_metadata, llm_count_prompt_tokens, llm_get_built_in_prompt_spec, llm_list_built_in_prompt_specs,
llm_render_built_in_prompt, llm_render_built_in_session_prompt, llm_render_prompt, llm_render_session_prompt,
llm_get_built_in_prompt_spec, llm_list_built_in_prompt_specs, llm_render_built_in_prompt,
llm_render_built_in_session_prompt,
},
request_builder::{
llm_build_canonical_request, llm_build_canonical_structured_request, llm_build_embedding_request,
@@ -22,29 +20,34 @@ pub use core::{
structured_output::{llm_canonical_json_schema_hash, llm_validate_json_schema},
};
pub use action::run_native_action_recipe_prepared_stream;
pub use contract_schema::{
llm_compile_execution_plan, llm_get_contract_schema, llm_normalize_prepared_routes, llm_validate_contract,
pub use action::copilot_action_recipe;
pub use byok::{
ByokCapabilityInput, ByokCatalogModelOutput, ByokCatalogOutput, ByokCatalogProviderOutput, ByokEndpointInput,
ByokLocalLeaseOutput, ByokModelDeclarationInput, ByokModelProbeCheckOutput, ByokModelProbeOutput,
ByokProbeCheckInput, ByokProbeResultOutput, ByokProbeStatusOutput, ByokProfileDefinitionInput, ByokProfileOutput,
ByokValidationOutput, CreateByokLocalLeaseInput, CreateByokLocalLeaseProviderInput, CreateByokProfileInput,
ProbeByokDraftInput, ProbeByokProfileInput, ReorderByokProfilesInput, ReplaceByokProfileInput,
RotateByokCredentialInput, byok_catalog,
};
#[cfg(test)]
pub(crate) use ffi::{AsyncLlmDispatchPreparedTask, resolve_request_chain};
#[napi_derive::napi(catch_unwind)]
pub fn llm_get_byok_catalog() -> ByokCatalogOutput {
byok_catalog()
}
pub(crate) use byok::{ByokProfileDefinition, validate_definition};
pub use contract_schema::{llm_get_contract_schema, llm_validate_contract};
pub(crate) use ffi::{
LlmDispatchPayload, LlmEmbeddingDispatchPayload, LlmMiddlewarePayload, LlmPreparedImageDispatchRoutePayload,
LlmRerankDispatchPayload, LlmRoutedBackendPayload, LlmStructuredDispatchPayload, StreamPipeline,
apply_request_middlewares, apply_structured_request_middlewares, backend_transport_error,
dispatch_prepared_image_route_payloads, dispatch_prepared_structured_routes, map_backend_error, map_json_error,
parse_embedding_protocol, parse_prepared_chat_routes_with_middleware, parse_prepared_chat_routes_without_middleware,
parse_protocol, parse_rerank_protocol, parse_structured_protocol, resolve_stream_chain,
LlmDispatchPayload, LlmMiddlewarePayload, LlmRerankDispatchPayload, LlmStructuredDispatchPayload,
};
pub use ffi::{
llm_dispatch_prepared, llm_embedding_dispatch, llm_embedding_dispatch_prepared, llm_image_dispatch_prepared,
llm_plan_attachment_reference, llm_rerank_dispatch, llm_rerank_dispatch_prepared, llm_resolve_request_intent,
llm_structured_dispatch, llm_structured_dispatch_prepared,
};
pub(crate) use host::{
LlmStreamHandle, STREAM_ABORTED_REASON, STREAM_CALLBACK_DISPATCH_FAILED_REASON, STREAM_END_MARKER, emit_error_event,
};
pub use host::{
llm_dispatch_prepared_stream, llm_dispatch_tool_loop_stream, llm_dispatch_tool_loop_stream_prepared,
llm_dispatch_tool_loop_stream_routed,
pub use prompt_catalog::llm_get_built_in_route_options;
pub use route::{
CopilotAccessProjection, CopilotExecuteInput, CopilotManagedTier, CopilotRouteCheckInput, CopilotTargetOverrideInput,
};
pub(crate) fn invalid_arg(message: impl Into<String>) -> napi::Error {
napi::Error::new(napi::Status::InvalidArg, message.into())
}
pub(crate) fn map_json_error(error: serde_json::Error) -> napi::Error {
invalid_arg(error.to_string())
}
packages/backend/native/src/llm/prompt_catalog.rs
+249 -30
View File
@@ -46,15 +46,12 @@ pub struct PromptSpecMessage {
}
#[napi(object)]
#[derive(Debug, Clone, Deserialize, Serialize)]
#[derive(Clone, Serialize)]
#[serde(rename_all = "camelCase")]
pub struct BuiltInPromptSpec {
pub name: String,
#[serde(skip_serializing_if = "Option::is_none")]
pub action: Option<String>,
pub model: String,
#[serde(skip_serializing_if = "Option::is_none")]
pub optional_models: Option<Vec<String>>,
#[serde(skip_serializing_if = "Option::is_none")]
pub config: Option<Value>,
#[serde(skip_serializing_if = "Option::is_none")]
@@ -64,6 +61,69 @@ pub struct BuiltInPromptSpec {
pub messages: Vec<PromptSpecMessage>,
}
#[derive(Clone, Deserialize)]
#[serde(rename_all = "camelCase", deny_unknown_fields)]
struct PromptCatalogSpec {
name: String,
#[serde(default)]
action: Option<String>,
#[serde(default)]
managed_route: Option<BuiltInManagedRouteSpec>,
#[serde(default)]
config: Option<Value>,
#[serde(default)]
params: Option<BTreeMap<String, PromptParamSpec>>,
#[serde(default)]
builtins: Option<Vec<PromptBuiltin>>,
messages: Vec<PromptSpecMessage>,
}
#[derive(Clone, Deserialize)]
#[serde(rename_all = "camelCase", deny_unknown_fields)]
struct BuiltInManagedRouteSpec {
targets: Vec<String>,
#[serde(default)]
premium_targets: Option<Vec<String>>,
#[serde(default)]
selectable_targets: Vec<BuiltInManagedTargetSpec>,
}
#[derive(Clone, Deserialize)]
#[serde(rename_all = "camelCase", deny_unknown_fields)]
struct BuiltInManagedTargetSpec {
id: String,
model_id: String,
display_name: String,
minimum_tier: BuiltInManagedTargetTier,
}
#[napi(string_enum)]
#[derive(Clone, Copy, PartialEq, Eq, Deserialize, Serialize)]
#[serde(rename_all = "lowercase")]
pub enum BuiltInManagedTargetTier {
Standard,
Premium,
}
#[napi(object)]
#[derive(Clone, Serialize)]
#[serde(rename_all = "camelCase")]
pub struct BuiltInManagedTarget {
pub id: String,
pub display_name: String,
pub minimum_tier: BuiltInManagedTargetTier,
}
#[napi(object)]
#[derive(Clone, Serialize)]
#[serde(rename_all = "camelCase")]
pub struct BuiltInRouteOptions {
pub route_id: String,
pub standard_default_target_id: Option<String>,
pub premium_default_target_id: Option<String>,
pub choices: Vec<BuiltInManagedTarget>,
}
#[derive(Debug, Clone, Serialize)]
#[serde(rename_all = "camelCase")]
pub(crate) struct BuiltInPromptMessage {
@@ -73,20 +133,29 @@ pub(crate) struct BuiltInPromptMessage {
pub(crate) params: Option<Map<String, Value>>,
}
#[derive(Debug, Clone, Serialize)]
#[derive(Clone, Serialize)]
#[serde(rename_all = "camelCase")]
pub(crate) struct BuiltInPrompt {
pub(crate) name: String,
#[serde(skip_serializing_if = "Option::is_none")]
pub(crate) action: Option<String>,
pub(crate) model: String,
#[serde(skip_serializing_if = "Option::is_none")]
pub(crate) optional_models: Option<Vec<String>>,
pub(crate) managed_targets: Vec<String>,
pub(crate) managed_premium_targets: Option<Vec<String>>,
#[serde(skip)]
pub(crate) managed_selectable_targets: Vec<BuiltInManagedTargetDefinition>,
#[serde(skip_serializing_if = "Option::is_none")]
pub(crate) config: Option<Value>,
pub(crate) messages: Vec<BuiltInPromptMessage>,
}
#[derive(Clone)]
pub(crate) struct BuiltInManagedTargetDefinition {
pub(crate) id: String,
pub(crate) model_id: String,
pub(crate) display_name: String,
pub(crate) minimum_tier: BuiltInManagedTargetTier,
}
struct PromptCatalog {
specs: Vec<BuiltInPromptSpec>,
prompts: Vec<BuiltInPrompt>,
@@ -112,16 +181,89 @@ pub(crate) fn built_in_prompt(name: &str) -> Option<&'static BuiltInPrompt> {
.and_then(|index| BUILTIN_PROMPT_CATALOG.prompts.get(*index))
}
pub(crate) fn built_in_managed_targets(name: &str, premium: bool) -> Option<&'static [String]> {
built_in_prompt(name).and_then(|prompt| {
let targets = if premium {
prompt
.managed_premium_targets
.as_deref()
.unwrap_or(prompt.managed_targets.as_slice())
} else {
prompt.managed_targets.as_slice()
};
(!targets.is_empty()).then_some(targets)
})
}
pub(crate) fn built_in_managed_target(
name: &str,
target_id: &str,
premium: bool,
) -> Option<&'static BuiltInManagedTargetDefinition> {
built_in_prompt(name)?
.managed_selectable_targets
.iter()
.find(|target| target.id == target_id && (premium || target.minimum_tier == BuiltInManagedTargetTier::Standard))
}
#[napi(catch_unwind)]
pub fn llm_get_built_in_route_options(name: String) -> Option<BuiltInRouteOptions> {
let prompt = built_in_prompt(&name)?;
if prompt.managed_selectable_targets.is_empty() {
return None;
}
let target_id_for_model = |model: Option<&String>| {
model.and_then(|model| {
prompt
.managed_selectable_targets
.iter()
.find(|target| &target.model_id == model)
.map(|target| target.id.clone())
})
};
Some(BuiltInRouteOptions {
route_id: prompt.name.clone(),
standard_default_target_id: target_id_for_model(prompt.managed_targets.first()),
premium_default_target_id: target_id_for_model(
prompt
.managed_premium_targets
.as_ref()
.and_then(|targets| targets.first())
.or_else(|| prompt.managed_targets.first()),
),
choices: prompt
.managed_selectable_targets
.iter()
.map(|target| BuiltInManagedTarget {
id: target.id.clone(),
display_name: target.display_name.clone(),
minimum_tier: target.minimum_tier,
})
.collect(),
})
}
impl PromptCatalog {
fn load() -> Result<Self, String> {
let partials: BTreeMap<String, String> =
serde_json::from_str(PROMPT_PARTIALS_SOURCE).map_err(|error| format!("invalid prompt partials JSON: {error}"))?;
let specs: Vec<BuiltInPromptSpec> =
let catalog_specs: Vec<PromptCatalogSpec> =
serde_json::from_str(PROMPT_SPECS_SOURCE).map_err(|error| format!("invalid prompt spec JSON: {error}"))?;
let prompts = specs
let prompts = catalog_specs
.iter()
.map(|spec| compile_prompt_spec(spec, &partials))
.collect::<Result<Vec<_>, _>>()?;
let specs = catalog_specs
.into_iter()
.map(|spec| BuiltInPromptSpec {
name: spec.name,
action: spec.action,
config: spec.config.filter(|value| !value.is_null()),
params: spec.params,
builtins: spec.builtins,
messages: spec.messages,
})
.collect::<Vec<_>>();
Ok(Self {
specs_by_name: specs
@@ -140,7 +282,17 @@ impl PromptCatalog {
}
}
fn compile_prompt_spec(spec: &BuiltInPromptSpec, partials: &BTreeMap<String, String>) -> Result<BuiltInPrompt, String> {
fn compile_prompt_spec(spec: &PromptCatalogSpec, partials: &BTreeMap<String, String>) -> Result<BuiltInPrompt, String> {
if spec
.managed_route
.as_ref()
.is_some_and(|route| !valid_managed_route(route))
{
return Err(format!("Prompt \"{}\" has an invalid managed route", spec.name));
}
if !spec.messages.is_empty() && spec.managed_route.is_none() {
return Err(format!("Executable prompt \"{}\" requires a managed route", spec.name));
}
let resolved_templates = spec
.messages
.iter()
@@ -186,13 +338,71 @@ fn compile_prompt_spec(spec: &BuiltInPromptSpec, partials: &BTreeMap<String, Str
Ok(BuiltInPrompt {
name: spec.name.clone(),
action: spec.action.clone(),
model: spec.model.clone(),
optional_models: spec.optional_models.clone(),
managed_targets: spec
.managed_route
.as_ref()
.map(|route| route.targets.clone())
.unwrap_or_default(),
managed_premium_targets: spec
.managed_route
.as_ref()
.and_then(|route| route.premium_targets.clone()),
managed_selectable_targets: spec
.managed_route
.as_ref()
.map(|route| {
route
.selectable_targets
.iter()
.map(|target| BuiltInManagedTargetDefinition {
id: target.id.clone(),
model_id: target.model_id.clone(),
display_name: target.display_name.clone(),
minimum_tier: target.minimum_tier,
})
.collect()
})
.unwrap_or_default(),
config: spec.config.clone().filter(|value| !value.is_null()),
messages,
})
}
fn valid_managed_route(route: &BuiltInManagedRouteSpec) -> bool {
if route.targets.is_empty()
|| route.targets.iter().any(|target| target.trim().is_empty())
|| route
.premium_targets
.as_ref()
.is_some_and(|targets| targets.is_empty() || targets.iter().any(|target| target.trim().is_empty()))
{
return false;
}
let ids = route
.selectable_targets
.iter()
.map(|target| target.id.as_str())
.collect::<BTreeSet<_>>();
let models = route
.selectable_targets
.iter()
.map(|target| target.model_id.as_str())
.collect::<BTreeSet<_>>();
if route.selectable_targets.iter().any(|target| {
target.id.trim().is_empty() || target.model_id.trim().is_empty() || target.display_name.trim().is_empty()
}) || ids.len() != route.selectable_targets.len()
|| models.len() != route.selectable_targets.len()
{
return false;
}
route.selectable_targets.is_empty()
|| route
.targets
.iter()
.chain(route.premium_targets.iter().flatten())
.all(|model| models.contains(model.as_str()))
}
fn normalize_prompt_param(spec: &PromptParamSpec) -> Value {
match spec.enum_values.as_ref() {
Some(values) if !values.is_empty() => {
@@ -250,7 +460,7 @@ fn resolve_prompt_template(template: &str, partials: &BTreeMap<String, String>)
Err("Prompt partial expansion exceeded maximum depth".to_string())
}
fn validate_builtins(spec: &BuiltInPromptSpec, templates: &[String]) -> Result<(), String> {
fn validate_builtins(spec: &PromptCatalogSpec, templates: &[String]) -> Result<(), String> {
let declared = spec
.builtins
.clone()
@@ -355,26 +565,35 @@ mod tests {
);
let chat = built_in_prompt("Chat With AFFiNE AI").expect("chat prompt");
assert_eq!(chat.model, "gpt-5.6-luna");
assert_eq!(chat.managed_targets, ["gpt-5.6-luna"]);
assert_eq!(
chat
.optional_models
.as_ref()
.map(|models| models.iter().map(String::as_str).collect::<Vec<_>>()),
Some(vec![
"gpt-5.6-luna",
"gpt-5.6-terra",
"gemini-3.6-flash",
"claude-sonnet-4-6"
])
.managed_premium_targets
.as_deref()
.map(|targets| targets.iter().map(String::as_str).collect::<Vec<_>>()),
Some(vec!["gpt-5.6-luna"])
);
let options = llm_get_built_in_route_options(chat.name.clone()).expect("chat route options");
assert_eq!(options.standard_default_target_id.as_deref(), Some("luna"));
assert_eq!(options.premium_default_target_id.as_deref(), Some("luna"));
assert_eq!(options.choices.len(), 4);
assert_eq!(
built_in_managed_target(&chat.name, "terra", false).map(|target| target.model_id.as_str()),
None
);
assert_eq!(
chat.config.as_ref().and_then(|config| config.get("proModels")),
Some(&serde_json::json!([
"gpt-5.6-terra",
"gemini-3.6-flash",
"claude-sonnet-4-6"
]))
built_in_managed_target(&chat.name, "terra", true).map(|target| target.model_id.as_str()),
Some("gpt-5.6-terra")
);
let transcript = built_in_prompt("Transcript audio structured").expect("transcript prompt");
assert_eq!(transcript.managed_targets, ["gemini-3.5-flash-lite"]);
assert_eq!(
transcript
.managed_premium_targets
.as_deref()
.map(|targets| targets.iter().map(String::as_str).collect::<Vec<_>>()),
Some(vec!["gemini-3.6-flash"])
);
}
}
packages/backend/native/src/llm/route/catalog.rs
+288
View File
@@ -0,0 +1,288 @@
use llm_adapter::capability::{
AttachmentKind, AttachmentSource, ModelFeature, ModelInput, ModelOutput, ModelRequirements,
};
use crate::llm::{
prompt_catalog::{built_in_managed_target, built_in_managed_targets},
route::CopilotManagedTier,
};
#[derive(Clone, Copy, PartialEq, Eq)]
pub(crate) enum RouteOperation {
Chat,
Structured,
Embedding,
Rerank,
Image,
Transcription,
}
#[derive(Clone, PartialEq, Eq)]
pub(crate) struct CatalogSlot {
pub(crate) id: &'static str,
pub(crate) operation: RouteOperation,
pub(crate) requirements: ModelRequirements,
}
#[derive(Clone, Copy, PartialEq, Eq)]
pub(crate) enum QuotaPolicy {
Metered,
Internal,
System,
}
pub(crate) fn quota_policy(slot: &CatalogSlot, built_in_route_id: Option<&str>) -> QuotaPolicy {
if matches!(slot.id, "index.embedding" | "search.rerank") {
QuotaPolicy::System
} else if built_in_route_id == Some("Summary as title") {
QuotaPolicy::Internal
} else {
QuotaPolicy::Metered
}
}
pub(crate) fn managed_targets(
slot: &CatalogSlot,
built_in_route_id: Option<&str>,
managed_tier: CopilotManagedTier,
) -> Option<Vec<String>> {
if let Some(route_id) = built_in_route_id {
return built_in_managed_targets(route_id, managed_tier == CopilotManagedTier::Premium).map(<[String]>::to_vec);
}
match slot.id {
"index.embedding" => Some(vec!["gemini-embedding-001".to_string()]),
"search.rerank" => Some(vec!["gpt-4o-mini".to_string()]),
_ => None,
}
}
pub(crate) fn managed_selected_target(
built_in_route_id: Option<&str>,
target_id: &str,
managed_tier: CopilotManagedTier,
) -> Option<String> {
built_in_managed_target(
built_in_route_id?,
target_id,
managed_tier == CopilotManagedTier::Premium,
)
.map(|target| target.model_id.clone())
}
pub(crate) fn slot(id: &str) -> Option<CatalogSlot> {
let (canonical, operation, input, output, features, attachment_kinds, attachment_sources) = match id {
"chat.default" | "prompt.text" => (
if id == "chat.default" {
"chat.default"
} else {
"prompt.text"
},
RouteOperation::Chat,
vec![ModelInput::Text],
vec![ModelOutput::Text],
vec![],
vec![],
vec![],
),
"chat.structured" | "prompt.structured" => (
if id == "chat.structured" {
"chat.structured"
} else {
"prompt.structured"
},
RouteOperation::Structured,
vec![ModelInput::Text],
vec![ModelOutput::Structured],
vec![],
vec![],
vec![],
),
"action.mindmap.generate" => (
"action.mindmap.generate",
RouteOperation::Structured,
vec![ModelInput::Text],
vec![ModelOutput::Structured],
vec![],
vec![],
vec![],
),
"action.slides.outline" => (
"action.slides.outline",
RouteOperation::Structured,
vec![ModelInput::Text],
vec![ModelOutput::Structured],
vec![],
vec![],
vec![],
),
"image.generate"
| "action.image.filter.sketch"
| "action.image.filter.clay"
| "action.image.filter.anime"
| "action.image.filter.pixel" => (
match id {
"image.generate" => "image.generate",
"action.image.filter.sketch" => "action.image.filter.sketch",
"action.image.filter.clay" => "action.image.filter.clay",
"action.image.filter.anime" => "action.image.filter.anime",
_ => "action.image.filter.pixel",
},
RouteOperation::Image,
vec![ModelInput::Text],
vec![ModelOutput::Image],
vec![],
vec![],
vec![],
),
"index.embedding" => (
"index.embedding",
RouteOperation::Embedding,
vec![ModelInput::Text],
vec![ModelOutput::Embedding],
vec![],
vec![],
vec![],
),
"search.rerank" => (
"search.rerank",
RouteOperation::Rerank,
vec![ModelInput::Text],
vec![ModelOutput::Rerank],
vec![],
vec![],
vec![],
),
"transcript.audio" => (
"transcript.audio",
RouteOperation::Transcription,
vec![ModelInput::Audio],
vec![ModelOutput::Structured],
vec![],
vec![AttachmentKind::Audio],
vec![
AttachmentSource::Url,
AttachmentSource::Data,
AttachmentSource::Bytes,
AttachmentSource::FileHandle,
],
),
_ => return None,
};
Some(CatalogSlot {
id: canonical,
operation,
requirements: ModelRequirements {
input,
output,
features,
attachment_kinds,
attachment_sources,
},
})
}
pub(crate) fn with_request_requirements(
mut slot: CatalogSlot,
needs_tools: bool,
attachment_kinds: Vec<AttachmentKind>,
attachment_sources: Vec<AttachmentSource>,
) -> CatalogSlot {
if needs_tools {
slot.requirements.features.push(ModelFeature::ToolCalling);
}
for kind in attachment_kinds {
let input = match kind {
AttachmentKind::Image => ModelInput::Image,
AttachmentKind::Audio => ModelInput::Audio,
AttachmentKind::File => ModelInput::File,
};
if !slot.requirements.input.contains(&input) {
slot.requirements.input.push(input);
}
if !slot.requirements.attachment_kinds.contains(&kind) {
slot.requirements.attachment_kinds.push(kind);
}
}
for source in attachment_sources {
if !slot.requirements.attachment_sources.contains(&source) {
slot.requirements.attachment_sources.push(source);
}
}
slot
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn inventory_slots_have_one_operation_and_explicit_requirements() {
let slots = [
"chat.default",
"chat.structured",
"prompt.text",
"prompt.structured",
"action.mindmap.generate",
"action.slides.outline",
"action.image.filter.sketch",
"action.image.filter.clay",
"action.image.filter.anime",
"action.image.filter.pixel",
"image.generate",
"index.embedding",
"search.rerank",
"transcript.audio",
];
for id in slots {
let slot = slot(id).expect("inventory slot must exist");
assert!(!slot.requirements.input.is_empty());
assert!(!slot.requirements.output.is_empty());
}
}
#[test]
fn managed_routes_are_built_in_by_prompt_or_system_slot() {
let prompt_slot = slot("prompt.text").unwrap();
assert_eq!(
managed_targets(&prompt_slot, Some("Summary as title"), CopilotManagedTier::Standard).unwrap(),
["gpt-5.6-luna"]
);
assert!(
managed_targets(
&prompt_slot,
Some("workflow:presentation"),
CopilotManagedTier::Standard
)
.is_none()
);
assert_eq!(
managed_targets(
&prompt_slot,
Some("Transcript audio structured"),
CopilotManagedTier::Premium
)
.unwrap(),
["gemini-3.6-flash"]
);
assert_eq!(
managed_targets(&slot("index.embedding").unwrap(), None, CopilotManagedTier::Standard).unwrap(),
["gemini-embedding-001"]
);
assert_eq!(
managed_targets(&slot("search.rerank").unwrap(), None, CopilotManagedTier::Standard).unwrap(),
["gpt-4o-mini"]
);
assert!(matches!(
quota_policy(&prompt_slot, Some("Summary as title")),
QuotaPolicy::Internal
));
assert!(matches!(
quota_policy(&prompt_slot, Some("Chat With AFFiNE AI")),
QuotaPolicy::Metered
));
assert!(matches!(
quota_policy(&slot("index.embedding").unwrap(), None),
QuotaPolicy::System
));
}
}
packages/backend/native/src/llm/route/contract.rs
+50
View File
@@ -0,0 +1,50 @@
#[napi_derive::napi(string_enum)]
#[derive(Clone, Copy, PartialEq, Eq)]
pub enum CopilotManagedTier {
Standard,
Premium,
}
#[derive(Clone)]
#[napi_derive::napi(object)]
pub struct CopilotAccessProjection {
pub route_allowed: bool,
pub managed_tier: CopilotManagedTier,
pub server_byok: bool,
pub local_byok: bool,
}
#[derive(Clone)]
#[napi_derive::napi(object)]
pub struct CopilotTargetOverrideInput {
pub profile_id: String,
pub model_id: String,
}
#[derive(Clone)]
#[napi_derive::napi(object)]
pub struct CopilotRouteCheckInput {
pub slot: String,
pub built_in_route_id: Option<String>,
pub workspace_id: Option<String>,
pub user_id: Option<String>,
pub local_lease_id: Option<String>,
pub access: CopilotAccessProjection,
pub managed_target_id: Option<String>,
pub target_override: Option<CopilotTargetOverrideInput>,
}
#[derive(Clone)]
#[napi_derive::napi(object)]
pub struct CopilotExecuteInput {
pub slot: String,
pub built_in_route_id: Option<String>,
pub workspace_id: Option<String>,
pub user_id: Option<String>,
pub local_lease_id: Option<String>,
pub access: CopilotAccessProjection,
pub managed_target_id: Option<String>,
pub target_override: Option<CopilotTargetOverrideInput>,
#[napi(ts_type = "unknown")]
pub request: serde_json::Value,
}
packages/backend/native/src/llm/route/mod.rs
+15
View File
@@ -0,0 +1,15 @@
mod catalog;
mod contract;
mod policy;
pub(crate) use catalog::{
CatalogSlot, QuotaPolicy, RouteOperation, managed_selected_target, managed_targets, quota_policy, slot,
with_request_requirements,
};
pub use contract::{
CopilotAccessProjection, CopilotExecuteInput, CopilotManagedTier, CopilotRouteCheckInput, CopilotTargetOverrideInput,
};
pub(crate) use policy::{
AuthorizedProfileRef, AuthorizedTargetRef, CredentialRef, Deployment, ProfileSource, RouteDecision,
RouteDecisionReason, RoutePolicyInput, TargetOverride, decide,
};
packages/backend/native/src/llm/route/policy.rs
+311
View File
@@ -0,0 +1,311 @@
use llm_adapter::capability::declared_model_matches;
use super::CatalogSlot;
use crate::llm::byok::ByokProfileDefinition;
#[derive(Clone, Copy, PartialEq, Eq)]
pub(crate) enum Deployment {
Cloud,
SelfHosted,
}
#[derive(Clone, Copy, PartialEq, Eq)]
pub(crate) enum ProfileSource {
Server,
Local,
Managed,
}
pub(crate) struct AuthorizedProfileRef {
pub(crate) profile_id: String,
pub(crate) source: ProfileSource,
pub(crate) provider: String,
pub(crate) definition: ByokProfileDefinition,
pub(crate) sort_order: i32,
pub(crate) credential_ref: CredentialRef,
}
pub(crate) enum CredentialRef {
Envelope { encrypted: String, aad: Vec<u8> },
Managed { profile_id: String },
}
#[derive(Clone, PartialEq, Eq)]
pub(crate) struct TargetOverride {
pub(crate) profile_id: String,
pub(crate) model_id: String,
}
#[derive(Clone, Copy, PartialEq, Eq)]
pub(crate) enum RouteDecisionReason {
ByokDisabled,
AccessUnavailable,
ExplicitTargetUnavailable,
NoCompatibleTarget,
ManagedPresetUnavailable,
}
pub(crate) enum RouteDecision {
Ready(Vec<AuthorizedTargetRef>),
Denied(RouteDecisionReason),
NoRoute(RouteDecisionReason),
}
pub(crate) struct AuthorizedTargetRef {
pub(crate) profile_index: usize,
pub(crate) model_index: usize,
}
pub(crate) struct RoutePolicyInput<'a> {
pub(crate) slot: &'a CatalogSlot,
pub(crate) deployment: Deployment,
pub(crate) byok_enabled: bool,
pub(crate) access_available: bool,
pub(crate) profiles: &'a [AuthorizedProfileRef],
pub(crate) target_override: Option<&'a TargetOverride>,
pub(crate) target_override_managed: bool,
}
pub(crate) fn decide(input: RoutePolicyInput<'_>) -> RouteDecision {
if input.deployment == Deployment::SelfHosted && !input.byok_enabled {
return RouteDecision::NoRoute(RouteDecisionReason::ByokDisabled);
}
if input.target_override_managed {
if input.deployment != Deployment::Cloud {
return RouteDecision::Denied(RouteDecisionReason::ExplicitTargetUnavailable);
}
if !input.access_available {
return RouteDecision::Denied(RouteDecisionReason::AccessUnavailable);
}
}
if let Some(target) = input.target_override {
let mut selected = compatible_targets(&input, input.target_override_managed);
selected.retain(|candidate| {
let profile = &input.profiles[candidate.profile_index];
let model = &profile.definition.models[candidate.model_index];
profile.profile_id == target.profile_id && model.model_id == target.model_id
});
return if selected.is_empty() {
RouteDecision::Denied(RouteDecisionReason::ExplicitTargetUnavailable)
} else {
RouteDecision::Ready(selected)
};
}
let byok = compatible_targets(&input, false);
if !byok.is_empty() {
return RouteDecision::Ready(byok);
}
if !input.access_available {
return RouteDecision::Denied(RouteDecisionReason::AccessUnavailable);
}
if input.deployment == Deployment::Cloud {
let managed = compatible_targets(&input, true);
if managed.is_empty() {
RouteDecision::NoRoute(RouteDecisionReason::ManagedPresetUnavailable)
} else {
RouteDecision::Ready(managed)
}
} else {
RouteDecision::NoRoute(RouteDecisionReason::NoCompatibleTarget)
}
}
fn compatible_targets(input: &RoutePolicyInput<'_>, managed: bool) -> Vec<AuthorizedTargetRef> {
let mut profiles = input
.profiles
.iter()
.enumerate()
.filter(|(_, profile)| (profile.source == ProfileSource::Managed) == managed)
.collect::<Vec<_>>();
profiles.sort_by_key(|(_, profile)| profile.sort_order);
profiles
.into_iter()
.flat_map(|(profile_index, profile)| {
profile
.definition
.models
.iter()
.enumerate()
.filter(|(_, model)| model.enabled && declared_model_matches(&model.capabilities, &input.slot.requirements))
.map(move |(model_index, _)| AuthorizedTargetRef {
profile_index,
model_index,
})
})
.collect()
}
#[cfg(test)]
mod tests {
use llm_adapter::capability::{DeclaredModelCapability, ModelInput, ModelOutput};
use super::*;
use crate::llm::{
byok::{ByokEndpoint, ByokModelDeclaration},
route::catalog,
};
fn profile(id: &str, source: ProfileSource, model: &str, output: ModelOutput) -> AuthorizedProfileRef {
AuthorizedProfileRef {
profile_id: id.to_string(),
source,
provider: "openai".to_string(),
definition: ByokProfileDefinition {
version: 1,
endpoint: ByokEndpoint::Custom {
url: "https://example.test/v1".to_string(),
},
models: vec![ByokModelDeclaration {
model_id: model.to_string(),
enabled: true,
capabilities: vec![DeclaredModelCapability {
input: vec![ModelInput::Text],
output: vec![output],
features: vec![],
attachment_kinds: vec![],
attachment_sources: vec![],
}],
}],
},
sort_order: 0,
credential_ref: CredentialRef::Managed {
profile_id: id.to_string(),
},
}
}
#[test]
fn deployment_matrix_and_unsupported_only_fallback() {
let slot = catalog::slot("chat.default").unwrap();
let cases = [
(
Deployment::Cloud,
true,
vec![profile("managed", ProfileSource::Managed, "A", ModelOutput::Text)],
true,
),
(
Deployment::SelfHosted,
false,
vec![profile("managed", ProfileSource::Managed, "A", ModelOutput::Text)],
false,
),
(Deployment::SelfHosted, true, vec![], false),
(
Deployment::Cloud,
true,
vec![
profile("byok", ProfileSource::Server, "B", ModelOutput::Image),
profile("managed", ProfileSource::Managed, "A", ModelOutput::Text),
],
true,
),
];
for (deployment, byok_enabled, profiles, ready) in cases {
let decision = decide(RoutePolicyInput {
slot: &slot,
deployment,
byok_enabled,
access_available: true,
profiles: &profiles,
target_override: None,
target_override_managed: false,
});
assert_eq!(matches!(decision, RouteDecision::Ready(_)), ready);
}
let profiles = vec![
profile("byok", ProfileSource::Server, "B", ModelOutput::Text),
profile("managed", ProfileSource::Managed, "A", ModelOutput::Text),
];
assert!(matches!(
decide(RoutePolicyInput {
slot: &slot,
deployment: Deployment::Cloud,
byok_enabled: true,
access_available: false,
profiles: &profiles,
target_override: None,
target_override_managed: false,
}),
RouteDecision::Ready(_)
));
assert!(matches!(
decide(RoutePolicyInput {
slot: &slot,
deployment: Deployment::Cloud,
byok_enabled: true,
access_available: false,
profiles: &profiles[1..],
target_override: None,
target_override_managed: false,
}),
RouteDecision::Denied(RouteDecisionReason::AccessUnavailable)
));
}
#[test]
fn override_is_complete_and_custom_model_is_not_replaced_by_managed_model() {
let slot = catalog::slot("chat.default").unwrap();
let profiles = vec![
profile("byok", ProfileSource::Server, "vendor/model:B", ModelOutput::Text),
profile("managed", ProfileSource::Managed, "model:A", ModelOutput::Text),
];
let target = TargetOverride {
profile_id: "byok".to_string(),
model_id: "vendor/model:B".to_string(),
};
let RouteDecision::Ready(candidates) = decide(RoutePolicyInput {
slot: &slot,
deployment: Deployment::Cloud,
byok_enabled: true,
access_available: true,
profiles: &profiles,
target_override: Some(&target),
target_override_managed: false,
}) else {
panic!("override should resolve");
};
assert_eq!(
profiles[candidates[0].profile_index].definition.models[candidates[0].model_index].model_id,
"vendor/model:B"
);
let managed_target = TargetOverride {
profile_id: "managed".to_string(),
model_id: "model:A".to_string(),
};
let RouteDecision::Ready(candidates) = decide(RoutePolicyInput {
slot: &slot,
deployment: Deployment::Cloud,
byok_enabled: true,
access_available: true,
profiles: &profiles,
target_override: Some(&managed_target),
target_override_managed: true,
}) else {
panic!("managed selection should resolve");
};
assert!(matches!(
profiles[candidates[0].profile_index].source,
ProfileSource::Managed
));
let mut disabled = profile("disabled", ProfileSource::Server, "model:C", ModelOutput::Text);
disabled.definition.models[0].enabled = false;
assert!(matches!(
decide(RoutePolicyInput {
slot: &slot,
deployment: Deployment::SelfHosted,
byok_enabled: true,
access_available: true,
profiles: &[disabled],
target_override: None,
target_override_managed: false,
}),
RouteDecision::NoRoute(RouteDecisionReason::NoCompatibleTarget)
));
}
}
packages/backend/native/src/runtime/backend_runtime/byok/admission.rs
+74
View File
@@ -0,0 +1,74 @@
use std::{
net::{IpAddr, Ipv4Addr},
time::Duration,
};
use super::{RuntimeError, RuntimeResult};
use crate::{
llm::byok::{ByokEndpoint, ByokProfileDefinition},
runtime::config::CopilotByokRuntimeConfig,
};
const DNS_RESOLUTION_TIMEOUT: Duration = Duration::from_secs(5);
pub(super) async fn admit_endpoint(
definition: &ByokProfileDefinition,
policy: &CopilotByokRuntimeConfig,
) -> RuntimeResult<()> {
let ByokEndpoint::Custom { url } = &definition.endpoint else {
return Ok(());
};
if !policy.allow_custom_endpoint {
return Err(RuntimeError::invalid_input("custom BYOK endpoints are disabled"));
}
if policy.allow_private_endpoint {
return Ok(());
}
let parsed = url::Url::parse(url).map_err(|_| RuntimeError::invalid_input("invalid BYOK endpoint"))?;
let host = parsed
.host_str()
.ok_or_else(|| RuntimeError::invalid_input("invalid BYOK endpoint"))?;
if host.eq_ignore_ascii_case("localhost") {
return Err(RuntimeError::invalid_input("private BYOK endpoints are disabled"));
}
let port = parsed.port_or_known_default().unwrap_or(443);
let addresses = tokio::time::timeout(DNS_RESOLUTION_TIMEOUT, tokio::net::lookup_host((host, port)))
.await
.map_err(|_| RuntimeError::invalid_input("BYOK endpoint DNS resolution timed out"))?
.map_err(|_| RuntimeError::invalid_input("BYOK endpoint DNS resolution failed"))?;
let mut resolved = false;
for address in addresses {
resolved = true;
if is_private_address(address.ip()) {
return Err(RuntimeError::invalid_input("private BYOK endpoints are disabled"));
}
}
if !resolved {
return Err(RuntimeError::invalid_input("BYOK endpoint DNS resolution failed"));
}
Ok(())
}
fn is_private_address(address: IpAddr) -> bool {
match address {
IpAddr::V4(address) => {
address.is_private()
|| address.is_loopback()
|| address.is_link_local()
|| address.is_broadcast()
|| address.is_documentation()
|| address.is_unspecified()
|| address.octets()[0] == 0
|| Ipv4Addr::new(100, 64, 0, 0) <= address && address <= Ipv4Addr::new(100, 127, 255, 255)
}
IpAddr::V6(address) => {
address.is_loopback()
|| address.is_unspecified()
|| address.is_unique_local()
|| address.is_unicast_link_local()
|| address
.to_ipv4_mapped()
.is_some_and(|address| is_private_address(IpAddr::V4(address)))
}
}
}
packages/backend/native/src/runtime/backend_runtime/byok/local.rs
+198
View File
@@ -0,0 +1,198 @@
use hmac::{Hmac, KeyInit, Mac};
use serde::{Deserialize, Serialize};
use sha2::Sha256;
use sqlx::{PgPool, Row};
use uuid::Uuid;
use super::{RuntimeError, RuntimeResult, admit_endpoint, envelope_key, require_text, token_hash};
use crate::{
llm::{
ByokLocalLeaseOutput, ByokProfileDefinition, CreateByokLocalLeaseInput,
byok::{SensitiveCredential, local_aad},
validate_definition,
},
runtime::config::CopilotByokRuntimeConfig,
};
const LOCAL_LEASE_PURPOSE: &str = "copilot_byok_local_lease";
const LOCAL_LEASE_ACTIVE_PURPOSE: &str = "copilot_byok_local_lease:active";
const LOCAL_LEASE_TTL_MS: i64 = 10 * 60 * 1000;
#[derive(Serialize, Deserialize)]
#[serde(rename_all = "camelCase", deny_unknown_fields)]
pub(crate) struct LocalLeasePayload {
pub(crate) version: u32,
pub(crate) workspace_id: String,
pub(crate) user_id: String,
pub(crate) providers: Vec<LocalLeaseProvider>,
}
#[derive(Serialize, Deserialize)]
#[serde(rename_all = "camelCase", deny_unknown_fields)]
pub(crate) struct LocalLeaseProvider {
pub(crate) provider: String,
pub(crate) name: String,
pub(crate) description: Option<String>,
pub(crate) encrypted_credential: String,
pub(crate) definition: ByokProfileDefinition,
pub(crate) enabled: bool,
}
pub(in super::super) async fn create(
pool: &PgPool,
root_secret: &[u8],
policy: &CopilotByokRuntimeConfig,
input: CreateByokLocalLeaseInput,
) -> RuntimeResult<ByokLocalLeaseOutput> {
require_text(&input.workspace_id, "workspaceId")?;
require_text(&input.user_id, "userId")?;
if input.providers.is_empty() {
return Err(RuntimeError::invalid_input("providers is required"));
}
let lease_id = Uuid::new_v4().to_string();
let key = envelope_key(root_secret)?;
let mut fingerprint =
Hmac::<Sha256>::new_from_slice(root_secret).map_err(|_| RuntimeError::invalid_state("credential_unavailable"))?;
fingerprint.update(input.workspace_id.as_bytes());
fingerprint.update(&[0]);
fingerprint.update(input.user_id.as_bytes());
let mut providers = Vec::with_capacity(input.providers.len());
for (index, provider) in input.providers.into_iter().enumerate() {
require_text(&provider.name, "name")?;
require_text(&provider.credential, "credential")?;
let definition = validate_definition(&provider.provider, provider.definition)
.map_err(|error| RuntimeError::invalid_input(error.to_string()))?;
admit_endpoint(&definition, policy).await?;
fingerprint.update(&[0]);
fingerprint.update(provider.provider.as_bytes());
fingerprint.update(&[0]);
fingerprint.update(provider.credential.as_bytes());
fingerprint.update(&[0]);
fingerprint.update(
&serde_json::to_vec(&definition)
.map_err(|error| RuntimeError::json("serialize BYOK local definition failed", error))?,
);
let encrypted_credential = key
.encrypt(
&SensitiveCredential::new(provider.credential.into_bytes()),
&local_aad(
&input.workspace_id,
&input.user_id,
&lease_id,
index,
&provider.provider,
definition.endpoint_identity(),
),
)
.map_err(|_| RuntimeError::invalid_state("credential_unavailable"))?;
providers.push(LocalLeaseProvider {
provider: provider.provider,
name: provider.name,
description: provider.description,
encrypted_credential,
definition,
enabled: provider.enabled,
});
}
let active_key = hex::encode(fingerprint.finalize().into_bytes());
let payload = serde_json::to_value(LocalLeasePayload {
version: 1,
workspace_id: input.workspace_id,
user_id: input.user_id,
providers,
})
.map_err(|error| RuntimeError::json("serialize BYOK local lease failed", error))?;
let mut tx = pool
.begin()
.await
.map_err(|error| RuntimeError::database("create BYOK local lease transaction failed", error))?;
sqlx::query("SELECT pg_advisory_xact_lock(hashtextextended($1, 0))")
.bind(&active_key)
.execute(&mut *tx)
.await
.map_err(|error| RuntimeError::database("lock BYOK local lease failed", error))?;
if let Some(row) = sqlx::query(
r#"
SELECT payload->>'leaseId' AS lease_id
FROM runtime_states
WHERE purpose = $1 AND token_hash = $2
AND consumed_at IS NULL AND expires_at > clock_timestamp()
FOR UPDATE
"#,
)
.bind(LOCAL_LEASE_ACTIVE_PURPOSE)
.bind(token_hash(&active_key))
.fetch_optional(&mut *tx)
.await
.map_err(|error| RuntimeError::database("read active BYOK local lease failed", error))?
{
let existing_lease_id: String = row.get("lease_id");
if let Some(expires_at_ms) = sqlx::query_scalar::<_, i64>(
r#"
SELECT (EXTRACT(EPOCH FROM expires_at) * 1000)::BIGINT
FROM runtime_states
WHERE purpose = $1 AND token_hash = $2
AND consumed_at IS NULL AND expires_at > clock_timestamp()
FOR UPDATE
"#,
)
.bind(LOCAL_LEASE_PURPOSE)
.bind(token_hash(&existing_lease_id))
.fetch_optional(&mut *tx)
.await
.map_err(|error| RuntimeError::database("read active BYOK local lease payload failed", error))?
{
tx.commit()
.await
.map_err(|error| RuntimeError::database("reuse BYOK local lease commit failed", error))?;
return Ok(ByokLocalLeaseOutput {
lease_id: existing_lease_id,
expires_at_ms,
});
}
}
sqlx::query("DELETE FROM runtime_states WHERE purpose = $1 AND token_hash = $2")
.bind(LOCAL_LEASE_ACTIVE_PURPOSE)
.bind(token_hash(&active_key))
.execute(&mut *tx)
.await
.map_err(|error| RuntimeError::database("delete stale BYOK local lease active record failed", error))?;
let expires_at_ms = sqlx::query_scalar::<_, i64>(
r#"
INSERT INTO runtime_states (purpose, token_hash, lookup_key, payload, expires_at)
VALUES ($1, $2, $3, $4, clock_timestamp() + ($5 * INTERVAL '1 millisecond'))
RETURNING (EXTRACT(EPOCH FROM expires_at) * 1000)::BIGINT
"#,
)
.bind(LOCAL_LEASE_PURPOSE)
.bind(token_hash(&lease_id))
.bind(&active_key)
.bind(payload)
.bind(LOCAL_LEASE_TTL_MS as f64)
.fetch_one(&mut *tx)
.await
.map_err(|error| RuntimeError::database("create BYOK local lease failed", error))?;
sqlx::query(
r#"
INSERT INTO runtime_states (purpose, token_hash, lookup_key, payload, expires_at)
VALUES ($1, $2, $3, jsonb_build_object('leaseId', $4::text), clock_timestamp() + ($5 * INTERVAL '1 millisecond'))
"#,
)
.bind(LOCAL_LEASE_ACTIVE_PURPOSE)
.bind(token_hash(&active_key))
.bind(&active_key)
.bind(&lease_id)
.bind(LOCAL_LEASE_TTL_MS as f64)
.execute(&mut *tx)
.await
.map_err(|error| RuntimeError::database("create BYOK local lease active record failed", error))?;
tx.commit()
.await
.map_err(|error| RuntimeError::database("create BYOK local lease commit failed", error))?;
Ok(ByokLocalLeaseOutput {
lease_id,
expires_at_ms,
})
}
packages/backend/native/src/runtime/backend_runtime/byok/mod.rs
+11
View File
@@ -0,0 +1,11 @@
mod admission;
mod local;
mod probe;
mod profile;
use admission::admit_endpoint;
pub(super) use local::{LocalLeasePayload, create as create_local_lease};
pub(super) use profile::{create, delete, list, probe_draft, probe_profile, reorder, replace, rotate};
use profile::{envelope_key, require_text};
use super::{RuntimeError, RuntimeResult, backend_provider, byok_endpoint, executable_protocol, token_hash};
packages/backend/native/src/runtime/backend_runtime/byok/probe.rs
+465
View File
@@ -0,0 +1,465 @@
use std::collections::{HashMap, HashSet};
use llm_adapter::{
backend::{BackendError, DefaultHttpClient},
capability::{
AttachmentKind, AttachmentSource, ModelFeature, ModelInput, ModelOutput, ModelRequirements, declared_model_matches,
},
core::{
CoreContent, CoreMessage, CoreRequest, CoreRole, CoreToolDefinition, EmbeddingRequest, ImageOptions,
ImageProviderOptions, ImageRequest, RerankCandidate, RerankRequest, StructuredRequest,
},
router::{ExecutablePreparedRoute, ExecutableRequest, dispatch_prepared_route},
target::{BackendCredential, BackendOperation, BackendTargetInput, EgressPolicy, compile_backend_target},
};
use serde_json::json;
use super::{RuntimeError, RuntimeResult, backend_provider, byok_endpoint, executable_protocol};
use crate::{
llm::{
ByokModelProbeCheckOutput, ByokModelProbeOutput, ByokProbeCheckInput, ByokProbeResultOutput, ByokProbeStatusOutput,
byok::{ByokEndpoint, ByokProfileDefinition, SensitiveCredential, definition_fingerprint},
},
runtime::config::CopilotByokRuntimeConfig,
};
pub(super) async fn execute_probe(
provider: &str,
definition: &ByokProfileDefinition,
credential: SensitiveCredential,
policy: &CopilotByokRuntimeConfig,
checks: Vec<ByokProbeCheckInput>,
) -> RuntimeResult<ByokProbeResultOutput> {
let tested_at_ms = chrono::Utc::now().timestamp_millis();
let connection_error = connection_probe(provider, &definition.endpoint, &credential, policy).await;
let connection = status(tested_at_ms, connection_error.as_deref());
let mut requested = HashSet::new();
for check in checks {
if !requested.insert((check.model_id.clone(), check.operation.clone())) {
return Err(RuntimeError::invalid_input("duplicate BYOK probe check"));
}
if !matches!(
check.operation.as_str(),
"chat" | "structured" | "tools" | "vision" | "embedding" | "rerank" | "image" | "transcript"
) {
return Err(RuntimeError::invalid_input("unknown BYOK probe operation"));
}
}
let mut models = Vec::new();
for model in &definition.models {
let model_checks = requested
.iter()
.filter(|(model_id, _)| model_id == &model.model_id)
.map(|(_, operation)| operation.clone())
.collect::<Vec<_>>();
if model_checks.is_empty() {
continue;
}
let mut outputs = Vec::with_capacity(model_checks.len());
for operation in model_checks {
let probe_status = if connection_error.is_some() {
not_tested()
} else if !model.enabled {
failed(tested_at_ms, "model_disabled")
} else if !declared_model_matches(&model.capabilities, &requirements(&operation)) {
failed(tested_at_ms, "capability_not_declared")
} else if matches!(operation.as_str(), "vision" | "transcript") {
not_tested()
} else {
let provider = provider.to_string();
let endpoint = definition.endpoint.clone();
let model_id = model.model_id.clone();
let credential = String::from_utf8(credential.expose().to_vec())
.map_err(|_| RuntimeError::invalid_state("credential_unavailable"))?;
let operation_for_task = operation.clone();
let allow_private = policy.allow_custom_endpoint && policy.allow_private_endpoint;
tokio::task::spawn_blocking(move || {
dispatch_check(
&provider,
&endpoint,
&model_id,
credential,
&operation_for_task,
allow_private,
)
})
.await
.map_err(|error| RuntimeError::invalid_state(format!("BYOK model probe task failed: {error}")))?
};
outputs.push(ByokModelProbeCheckOutput {
operation,
status: probe_status,
});
}
models.push(ByokModelProbeOutput {
model_id: model.model_id.clone(),
checks: outputs,
});
}
if requested
.iter()
.any(|(model_id, _)| !definition.models.iter().any(|model| &model.model_id == model_id))
{
return Err(RuntimeError::invalid_input("BYOK probe model not found"));
}
Ok(ByokProbeResultOutput {
definition_fingerprint: definition_fingerprint(definition),
stale: false,
connection,
models,
})
}
async fn connection_probe(
provider: &str,
endpoint: &ByokEndpoint,
credential: &SensitiveCredential,
policy: &CopilotByokRuntimeConfig,
) -> Option<String> {
let credential = match std::str::from_utf8(credential.expose()) {
Ok(value) => value.to_string(),
Err(_) => return Some("credential_unavailable".to_string()),
};
let (url, headers) = probe_request(provider, endpoint, credential);
let allow_private = policy.allow_custom_endpoint && policy.allow_private_endpoint;
let result = tokio::task::spawn_blocking(move || {
safefetch::safe_fetch(&safefetch::SafeFetchRequest {
url,
method: Some(safefetch::SafeFetchMethod::Get),
headers: Some(headers.clone()),
body: None,
timeout_ms: Some(10_000),
max_redirects: Some(0),
max_bytes: Some(1024 * 1024),
allowed_headers: Some(headers.keys().cloned().collect()),
allowed_hosts: None,
allow_http: Some(allow_private),
allow_private_target_origin: Some(allow_private),
ech_config_list: None,
})
})
.await;
match result {
Ok(Ok(response))
if (200..300).contains(&response.status) && valid_connection_response(provider, &response.body) =>
{
None
}
Ok(Ok(response)) => Some(http_error_kind(response.status).to_string()),
_ => Some("transport".to_string()),
}
}
fn dispatch_check(
provider: &str,
endpoint: &ByokEndpoint,
model_id: &str,
credential: String,
operation: &str,
allow_private: bool,
) -> ByokProbeStatusOutput {
let checked_at = chrono::Utc::now().timestamp_millis();
let operation_kind = match operation {
"chat" | "tools" => BackendOperation::Chat,
"structured" => BackendOperation::Structured,
"embedding" => BackendOperation::Embedding,
"rerank" => BackendOperation::Rerank,
"image" => BackendOperation::Image,
_ => return not_tested(),
};
let target = compile_backend_target(BackendTargetInput {
provider: match backend_provider(provider) {
Ok(provider) => provider,
Err(_) => return failed(checked_at, "unsupported_provider"),
},
operation: operation_kind,
endpoint: byok_endpoint(provider, endpoint),
model: model_id.to_string(),
credential: BackendCredential::new(credential),
timeout_ms: Some(15_000),
egress_policy: if allow_private {
EgressPolicy::AllowPrivate
} else {
EgressPolicy::PublicOnly
},
});
let target = match target {
Ok(target) => target,
Err(_) => return failed(checked_at, "unsupported_operation"),
};
let route = ExecutablePreparedRoute::new(
executable_protocol(target.protocol),
target.model,
target.config,
probe_request_for_operation(operation),
);
let route = match route {
Ok(route) => route,
Err(_) => return failed(checked_at, "invalid_probe_request"),
};
match dispatch_prepared_route(&DefaultHttpClient::default(), &route) {
Ok(_) => verified(checked_at),
Err(error) => failed(checked_at, backend_error_kind(&error)),
}
}
fn probe_request_for_operation(operation: &str) -> ExecutableRequest {
let message = CoreMessage {
role: CoreRole::User,
content: vec![CoreContent::Text {
text: "Reply with OK.".to_string(),
}],
};
match operation {
"chat" | "tools" => ExecutableRequest::Chat(CoreRequest {
model: String::new(),
messages: vec![message],
stream: false,
max_tokens: Some(8),
temperature: Some(0.0),
tools: if operation == "tools" {
vec![CoreToolDefinition {
name: "byok_probe".to_string(),
description: Some("Probe tool compatibility".to_string()),
parameters: json!({ "type": "object", "properties": {} }),
}]
} else {
vec![]
},
tool_choice: None,
include: None,
reasoning: None,
response_schema: None,
}),
"structured" => ExecutableRequest::Structured(StructuredRequest {
model: String::new(),
messages: vec![message],
schema: json!({
"type": "object",
"properties": { "ok": { "type": "boolean" } },
"required": ["ok"],
"additionalProperties": false
}),
max_tokens: Some(16),
temperature: Some(0.0),
reasoning: None,
strict: Some(true),
response_mime_type: Some("application/json".to_string()),
}),
"embedding" => ExecutableRequest::Embedding(EmbeddingRequest {
model: String::new(),
inputs: vec!["BYOK probe".to_string()],
dimensions: None,
task_type: None,
}),
"rerank" => ExecutableRequest::Rerank(RerankRequest {
model: String::new(),
query: "probe".to_string(),
candidates: vec![RerankCandidate {
id: None,
text: "probe".to_string(),
}],
top_n: Some(1),
}),
"image" => ExecutableRequest::Image(Box::new(ImageRequest::generate(
String::new(),
"A single black pixel".to_string(),
ImageOptions::default(),
ImageProviderOptions::default(),
))),
_ => unreachable!("validated probe operation"),
}
}
fn requirements(operation: &str) -> ModelRequirements {
let (input, output, features, attachment_kinds, attachment_sources) = match operation {
"chat" => (vec![ModelInput::Text], vec![ModelOutput::Text], vec![], vec![], vec![]),
"structured" => (
vec![ModelInput::Text],
vec![ModelOutput::Structured],
vec![],
vec![],
vec![],
),
"tools" => (
vec![ModelInput::Text],
vec![ModelOutput::Text],
vec![ModelFeature::ToolCalling],
vec![],
vec![],
),
"vision" => (
vec![ModelInput::Text, ModelInput::Image],
vec![ModelOutput::Text],
vec![],
vec![AttachmentKind::Image],
vec![AttachmentSource::Data],
),
"embedding" => (
vec![ModelInput::Text],
vec![ModelOutput::Embedding],
vec![],
vec![],
vec![],
),
"rerank" => (
vec![ModelInput::Text],
vec![ModelOutput::Rerank],
vec![],
vec![],
vec![],
),
"image" => (vec![ModelInput::Text], vec![ModelOutput::Image], vec![], vec![], vec![]),
"transcript" => (
vec![ModelInput::Audio],
vec![ModelOutput::Structured],
vec![],
vec![AttachmentKind::Audio],
vec![AttachmentSource::Data],
),
_ => unreachable!("validated probe operation"),
};
ModelRequirements {
input,
output,
features,
attachment_kinds,
attachment_sources,
}
}
fn probe_request(provider: &str, endpoint: &ByokEndpoint, credential: String) -> (String, HashMap<String, String>) {
let base = match endpoint {
ByokEndpoint::Custom { url } => url.as_str(),
ByokEndpoint::ProviderDefault => match provider {
"openai" => "https://api.openai.com/v1",
"anthropic" => "https://api.anthropic.com/v1",
"gemini" => "https://generativelanguage.googleapis.com/v1beta",
"fal" => "https://api.fal.ai/v1",
_ => unreachable!("validated provider"),
},
};
let mut headers = HashMap::new();
match provider {
"openai" => {
headers.insert("authorization".to_string(), format!("Bearer {credential}"));
}
"anthropic" => {
headers.insert("x-api-key".to_string(), credential);
headers.insert("anthropic-version".to_string(), "2023-06-01".to_string());
}
"gemini" => {
headers.insert("x-goog-api-key".to_string(), credential);
}
"fal" => {
headers.insert("authorization".to_string(), format!("Key {credential}"));
}
_ => unreachable!("validated provider"),
}
let suffix = if provider == "fal" { "models?limit=10" } else { "models" };
(format!("{}/{suffix}", base.trim_end_matches('/')), headers)
}
fn valid_connection_response(provider: &str, body: &[u8]) -> bool {
let Ok(body) = serde_json::from_slice::<serde_json::Value>(body) else {
return false;
};
match provider {
"openai" | "anthropic" => body.get("data").is_some_and(serde_json::Value::is_array),
"gemini" => body.get("models").is_some_and(serde_json::Value::is_array),
"fal" => body.get("error").is_none(),
_ => false,
}
}
fn status(tested_at_ms: i64, error: Option<&str>) -> ByokProbeStatusOutput {
match error {
Some(error) => failed(tested_at_ms, error),
None => verified(tested_at_ms),
}
}
fn verified(tested_at_ms: i64) -> ByokProbeStatusOutput {
ByokProbeStatusOutput {
kind: "verified".to_string(),
tested_at_ms: Some(tested_at_ms),
error_kind: None,
}
}
fn failed(tested_at_ms: i64, error: &str) -> ByokProbeStatusOutput {
ByokProbeStatusOutput {
kind: "failed".to_string(),
tested_at_ms: Some(tested_at_ms),
error_kind: Some(error.to_string()),
}
}
fn not_tested() -> ByokProbeStatusOutput {
ByokProbeStatusOutput {
kind: "not_tested".to_string(),
tested_at_ms: None,
error_kind: None,
}
}
fn http_error_kind(status: u16) -> &'static str {
match status {
401 => "authentication",
403 => "permission",
404 => "not_found",
429 => "rate_limited",
500..=599 => "unavailable",
_ => "rejected",
}
}
fn backend_error_kind(error: &BackendError) -> &'static str {
match error {
BackendError::UpstreamStatus { status, .. } => http_error_kind(*status),
BackendError::Transport { .. } => "transport",
BackendError::Timeout { .. } => "timeout",
BackendError::InvalidConfig { .. } | BackendError::InvalidRequest { .. } => "unsupported_operation",
BackendError::InvalidResponse { .. }
| BackendError::InvalidStructuredOutput { .. }
| BackendError::Json(_)
| BackendError::Stream(_) => "invalid_response",
BackendError::NoBackendAvailable => "unavailable",
}
}
#[cfg(test)]
mod tests {
use super::*;
use llm_adapter::target::BackendEndpoint;
#[test]
fn connection_probe_errors_are_low_information() {
let (url, headers) = probe_request("openai", &ByokEndpoint::ProviderDefault, "secret".to_string());
assert_eq!(url, "https://api.openai.com/v1/models");
assert_eq!(headers.get("authorization").map(String::as_str), Some("Bearer secret"));
assert_eq!(http_error_kind(401), "authentication");
assert_eq!(http_error_kind(403), "permission");
assert_eq!(http_error_kind(429), "rate_limited");
assert_eq!(http_error_kind(503), "unavailable");
let custom = ByokEndpoint::Custom {
url: "http://127.0.0.1:1234/v1".to_string(),
};
assert_eq!(
probe_request("openai", &custom, "secret".to_string()).0,
"http://127.0.0.1:1234/v1/models"
);
assert_eq!(
byok_endpoint("openai", &custom),
BackendEndpoint::Custom("http://127.0.0.1:1234".to_string())
);
assert!(valid_connection_response("openai", br#"{"data":[]}"#));
assert!(!valid_connection_response(
"openai",
br#"{"error":"Unexpected endpoint"}"#
));
}
}
packages/backend/native/src/runtime/backend_runtime/byok/profile.rs
+577
View File
@@ -0,0 +1,577 @@
use std::collections::{HashMap, HashSet};
use sqlx::{FromRow, PgPool};
use uuid::Uuid;
use super::{RuntimeError, RuntimeResult, admit_endpoint};
use crate::{
llm::{
ByokProfileDefinition, ByokProfileOutput, ByokValidationOutput, CreateByokProfileInput, ProbeByokDraftInput,
ProbeByokProfileInput, ReorderByokProfilesInput, ReplaceByokProfileInput, RotateByokCredentialInput,
byok::{CredentialEnvelopeKey, SensitiveCredential, reconcile_validation, server_aad},
validate_definition,
},
runtime::config::CopilotByokRuntimeConfig,
};
#[derive(FromRow)]
struct ProfileRow {
id: String,
workspace_id: String,
provider: String,
name: String,
description: Option<String>,
encrypted_api_key: String,
definition: serde_json::Value,
sort_order: i32,
enabled: bool,
revision: i32,
credential_generation: i32,
validation: Option<serde_json::Value>,
}
#[derive(FromRow)]
struct ProfileAdmissionRow {
provider: String,
revision: i32,
}
pub(in super::super) async fn list(pool: &PgPool, workspace_id: &str) -> RuntimeResult<Vec<ByokProfileOutput>> {
let rows = sqlx::query_as::<_, ProfileRow>(
r#"
SELECT id, workspace_id, provider, name, description, encrypted_api_key,
definition, sort_order, enabled, revision, credential_generation, validation
FROM ai_workspace_byok_configs
WHERE workspace_id = $1
ORDER BY sort_order ASC, created_at ASC
"#,
)
.bind(workspace_id)
.fetch_all(pool)
.await
.map_err(|error| RuntimeError::database("list BYOK profiles failed", error))?;
rows.into_iter().map(profile_output).collect()
}
pub(in super::super) async fn create(
pool: &PgPool,
root_secret: &[u8],
policy: &CopilotByokRuntimeConfig,
input: CreateByokProfileInput,
) -> RuntimeResult<ByokProfileOutput> {
require_text(&input.workspace_id, "workspaceId")?;
require_text(&input.name, "name")?;
require_text(&input.credential, "credential")?;
require_text(&input.actor_user_id, "actorUserId")?;
let definition = validate_definition(&input.provider, input.definition)
.map_err(|error| RuntimeError::invalid_input(error.to_string()))?;
admit_endpoint(&definition, policy).await?;
let key = envelope_key(root_secret)?;
let profile_id = Uuid::new_v4().to_string();
let aad = server_aad(
&input.workspace_id,
&profile_id,
&input.provider,
definition.endpoint_identity(),
);
let encrypted = key
.encrypt(&SensitiveCredential::new(input.credential.into_bytes()), &aad)
.map_err(|_| RuntimeError::invalid_state("credential_unavailable"))?;
let definition_json =
serde_json::to_value(&definition).map_err(|error| RuntimeError::json("serialize BYOK definition failed", error))?;
let mut tx = pool
.begin()
.await
.map_err(|error| RuntimeError::database("create BYOK profile transaction failed", error))?;
sqlx::query("SELECT pg_advisory_xact_lock(hashtextextended($1, 0))")
.bind(&input.workspace_id)
.execute(&mut *tx)
.await
.map_err(|error| RuntimeError::database("lock BYOK profile order failed", error))?;
let sort_order = sqlx::query_scalar::<_, i32>(
"SELECT COALESCE(MAX(sort_order) + 1, 0)::int FROM ai_workspace_byok_configs WHERE workspace_id = $1",
)
.bind(&input.workspace_id)
.fetch_one(&mut *tx)
.await
.map_err(|error| RuntimeError::database("resolve BYOK profile order failed", error))?;
let row = sqlx::query_as::<_, ProfileRow>(
r#"
INSERT INTO ai_workspace_byok_configs (
id, workspace_id, provider, name, description, encrypted_api_key,
definition, sort_order, enabled, created_by, updated_by, created_at, updated_at
)
VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $10, CURRENT_TIMESTAMP, CURRENT_TIMESTAMP)
RETURNING id, workspace_id, provider, name, description, encrypted_api_key,
definition, sort_order, enabled, revision, credential_generation, validation
"#,
)
.bind(&profile_id)
.bind(&input.workspace_id)
.bind(&input.provider)
.bind(&input.name)
.bind(&input.description)
.bind(encrypted)
.bind(definition_json)
.bind(sort_order)
.bind(input.enabled)
.bind(&input.actor_user_id)
.fetch_one(&mut *tx)
.await
.map_err(|error| RuntimeError::database("create BYOK profile failed", error))?;
tx.commit()
.await
.map_err(|error| RuntimeError::database("create BYOK profile commit failed", error))?;
profile_output(row)
}
pub(in super::super) async fn replace(
pool: &PgPool,
root_secret: &[u8],
policy: &CopilotByokRuntimeConfig,
input: ReplaceByokProfileInput,
) -> RuntimeResult<ByokProfileOutput> {
require_text(&input.workspace_id, "workspaceId")?;
require_text(&input.name, "name")?;
require_text(&input.actor_user_id, "actorUserId")?;
if let Some(credential) = input.credential.as_deref() {
require_text(credential, "credential")?;
}
if input.expected_revision < 1 {
return Err(RuntimeError::invalid_input("expectedRevision is required"));
}
let admission = select_profile_for_admission(pool, &input.workspace_id, &input.profile_id).await?;
if admission.revision != input.expected_revision {
return Err(RuntimeError::invalid_input("byok_revision_conflict"));
}
let definition = validate_definition(&admission.provider, input.definition)
.map_err(|error| RuntimeError::invalid_input(error.to_string()))?;
admit_endpoint(&definition, policy).await?;
let mut tx = pool
.begin()
.await
.map_err(|error| RuntimeError::database("replace BYOK profile transaction failed", error))?;
let old = select_profile_for_update(&mut tx, &input.workspace_id, &input.profile_id).await?;
if old.revision != input.expected_revision {
return Err(RuntimeError::invalid_input("byok_revision_conflict"));
}
let old_definition = parse_definition(old.definition.clone())?;
let key = envelope_key(root_secret)?;
let credential_changed = input.credential.is_some();
let credential = if let Some(credential) = input.credential {
SensitiveCredential::new(credential.into_bytes())
} else {
key
.decrypt(
&old.encrypted_api_key,
&server_aad(
&old.workspace_id,
&old.id,
&old.provider,
old_definition.endpoint_identity(),
),
)
.map_err(|_| RuntimeError::invalid_state("credential_unavailable"))?
};
let encrypted = key
.encrypt(
&credential,
&server_aad(
&old.workspace_id,
&old.id,
&old.provider,
definition.endpoint_identity(),
),
)
.map_err(|_| RuntimeError::invalid_state("credential_unavailable"))?;
let definition_json =
serde_json::to_value(&definition).map_err(|error| RuntimeError::json("serialize BYOK definition failed", error))?;
let credential_generation = old.credential_generation + i32::from(credential_changed);
let validation = reconcile_validation(
parse_validation(old.validation)?,
&old_definition,
&definition,
credential_generation,
credential_changed,
)
.map(serde_json::to_value)
.transpose()
.map_err(|error| RuntimeError::json("serialize BYOK validation failed", error))?;
let row = sqlx::query_as::<_, ProfileRow>(
r#"
UPDATE ai_workspace_byok_configs
SET name = $3, description = $4, definition = $5, encrypted_api_key = $6,
enabled = $7, credential_generation = $8, validation = $9,
revision = revision + 1, updated_by = $10, updated_at = CURRENT_TIMESTAMP
WHERE workspace_id = $1 AND id = $2 AND revision = $11
RETURNING id, workspace_id, provider, name, description, encrypted_api_key,
definition, sort_order, enabled, revision, credential_generation, validation
"#,
)
.bind(&input.workspace_id)
.bind(&input.profile_id)
.bind(&input.name)
.bind(&input.description)
.bind(definition_json)
.bind(encrypted)
.bind(input.enabled)
.bind(credential_generation)
.bind(validation)
.bind(&input.actor_user_id)
.bind(input.expected_revision)
.fetch_optional(&mut *tx)
.await
.map_err(|error| RuntimeError::database("replace BYOK profile failed", error))?
.ok_or_else(|| RuntimeError::invalid_input("byok_revision_conflict"))?;
tx.commit()
.await
.map_err(|error| RuntimeError::database("replace BYOK profile commit failed", error))?;
profile_output(row)
}
async fn select_profile_for_admission(
pool: &PgPool,
workspace_id: &str,
profile_id: &str,
) -> RuntimeResult<ProfileAdmissionRow> {
sqlx::query_as::<_, ProfileAdmissionRow>(
"SELECT provider, revision FROM ai_workspace_byok_configs WHERE workspace_id = $1 AND id = $2",
)
.bind(workspace_id)
.bind(profile_id)
.fetch_optional(pool)
.await
.map_err(|error| RuntimeError::database("load BYOK profile admission data failed", error))?
.ok_or_else(|| RuntimeError::invalid_input("BYOK profile not found"))
}
pub(in super::super) async fn rotate(
pool: &PgPool,
root_secret: &[u8],
input: RotateByokCredentialInput,
) -> RuntimeResult<ByokProfileOutput> {
require_text(&input.credential, "credential")?;
require_text(&input.actor_user_id, "actorUserId")?;
if input.expected_revision < 1 {
return Err(RuntimeError::invalid_input("expectedRevision is required"));
}
let mut tx = pool
.begin()
.await
.map_err(|error| RuntimeError::database("rotate BYOK credential transaction failed", error))?;
let profile = select_profile_for_update(&mut tx, &input.workspace_id, &input.profile_id).await?;
if profile.revision != input.expected_revision {
return Err(RuntimeError::invalid_input("byok_revision_conflict"));
}
let definition = parse_definition(profile.definition.clone())?;
let encrypted = envelope_key(root_secret)?
.encrypt(
&SensitiveCredential::new(input.credential.into_bytes()),
&server_aad(
&profile.workspace_id,
&profile.id,
&profile.provider,
definition.endpoint_identity(),
),
)
.map_err(|_| RuntimeError::invalid_state("credential_unavailable"))?;
let row = sqlx::query_as::<_, ProfileRow>(
r#"
UPDATE ai_workspace_byok_configs
SET encrypted_api_key = $3, credential_generation = credential_generation + 1,
validation = NULL, revision = revision + 1, updated_by = $4,
updated_at = CURRENT_TIMESTAMP
WHERE workspace_id = $1 AND id = $2 AND revision = $5
RETURNING id, workspace_id, provider, name, description, encrypted_api_key,
definition, sort_order, enabled, revision, credential_generation, validation
"#,
)
.bind(&input.workspace_id)
.bind(&input.profile_id)
.bind(encrypted)
.bind(&input.actor_user_id)
.bind(input.expected_revision)
.fetch_optional(&mut *tx)
.await
.map_err(|error| RuntimeError::database("rotate BYOK credential failed", error))?
.ok_or_else(|| RuntimeError::invalid_input("byok_revision_conflict"))?;
tx.commit()
.await
.map_err(|error| RuntimeError::database("rotate BYOK credential commit failed", error))?;
profile_output(row)
}
pub(in super::super) async fn delete(pool: &PgPool, workspace_id: &str, profile_id: &str) -> RuntimeResult<bool> {
let affected = sqlx::query("DELETE FROM ai_workspace_byok_configs WHERE workspace_id = $1 AND id = $2")
.bind(workspace_id)
.bind(profile_id)
.execute(pool)
.await
.map_err(|error| RuntimeError::database("delete BYOK profile failed", error))?
.rows_affected();
Ok(affected == 1)
}
pub(in super::super) async fn reorder(
pool: &PgPool,
input: ReorderByokProfilesInput,
) -> RuntimeResult<Vec<ByokProfileOutput>> {
require_text(&input.workspace_id, "workspaceId")?;
require_text(&input.actor_user_id, "actorUserId")?;
let unique = input
.profiles
.iter()
.map(|profile| profile.profile_id.as_str())
.collect::<HashSet<_>>();
if unique.len() != input.profiles.len() {
return Err(RuntimeError::invalid_input("duplicate BYOK profile id"));
}
let mut tx = pool
.begin()
.await
.map_err(|error| RuntimeError::database("reorder BYOK profiles transaction failed", error))?;
sqlx::query("SELECT pg_advisory_xact_lock(hashtextextended($1, 0))")
.bind(&input.workspace_id)
.execute(&mut *tx)
.await
.map_err(|error| RuntimeError::database("lock BYOK profile order failed", error))?;
let current = sqlx::query_as::<_, (String, i32)>(
"SELECT id, revision FROM ai_workspace_byok_configs WHERE workspace_id = $1 ORDER BY sort_order, created_at",
)
.bind(&input.workspace_id)
.fetch_all(&mut *tx)
.await
.map_err(|error| RuntimeError::database("read BYOK profile order failed", error))?;
let current_set = current
.iter()
.map(|(profile_id, _)| profile_id.as_str())
.collect::<HashSet<_>>();
if current_set != unique {
return Err(RuntimeError::invalid_input(
"BYOK profile order must contain every server profile",
));
}
let revisions = current.into_iter().collect::<HashMap<_, _>>();
if input
.profiles
.iter()
.any(|profile| revisions.get(&profile.profile_id).copied() != Some(profile.expected_revision))
{
return Err(RuntimeError::invalid_input("byok_revision_conflict"));
}
for (sort_order, profile) in input.profiles.iter().enumerate() {
let updated = sqlx::query(
"UPDATE ai_workspace_byok_configs SET sort_order = $3, revision = revision + 1, updated_by = $4, updated_at = \
CURRENT_TIMESTAMP WHERE workspace_id = $1 AND id = $2 AND revision = $5",
)
.bind(&input.workspace_id)
.bind(&profile.profile_id)
.bind(sort_order as i32)
.bind(&input.actor_user_id)
.bind(profile.expected_revision)
.execute(&mut *tx)
.await
.map_err(|error| RuntimeError::database("update BYOK profile order failed", error))?
.rows_affected();
if updated != 1 {
return Err(RuntimeError::invalid_input("byok_revision_conflict"));
}
}
let rows = sqlx::query_as::<_, ProfileRow>(
r#"
SELECT id, workspace_id, provider, name, description, encrypted_api_key,
definition, sort_order, enabled, revision, credential_generation, validation
FROM ai_workspace_byok_configs
WHERE workspace_id = $1
ORDER BY sort_order, created_at
"#,
)
.bind(&input.workspace_id)
.fetch_all(&mut *tx)
.await
.map_err(|error| RuntimeError::database("read reordered BYOK profiles failed", error))?;
tx.commit()
.await
.map_err(|error| RuntimeError::database("reorder BYOK profiles commit failed", error))?;
rows.into_iter().map(profile_output).collect()
}
pub(in super::super) async fn probe_profile(
pool: &PgPool,
root_secret: &[u8],
policy: &CopilotByokRuntimeConfig,
input: ProbeByokProfileInput,
) -> RuntimeResult<crate::llm::ByokProbeResultOutput> {
let profile = sqlx::query_as::<_, ProfileRow>(
r#"
SELECT id, workspace_id, provider, name, description, encrypted_api_key,
definition, sort_order, enabled, revision, credential_generation, validation
FROM ai_workspace_byok_configs
WHERE workspace_id = $1 AND id = $2
"#,
)
.bind(&input.workspace_id)
.bind(&input.profile_id)
.fetch_optional(pool)
.await
.map_err(|error| RuntimeError::database("read BYOK profile for probe failed", error))?
.ok_or_else(|| RuntimeError::invalid_input("BYOK profile not found"))?;
let definition = parse_definition(profile.definition.clone())?;
admit_endpoint(&definition, policy).await?;
let credential = envelope_key(root_secret)?
.decrypt(
&profile.encrypted_api_key,
&server_aad(
&profile.workspace_id,
&profile.id,
&profile.provider,
definition.endpoint_identity(),
),
)
.map_err(|_| RuntimeError::invalid_state("credential_unavailable"))?;
let mut result =
super::probe::execute_probe(&profile.provider, &definition, credential, policy, input.checks).await?;
let validation = ByokValidationOutput {
definition_fingerprint: result.definition_fingerprint.clone(),
credential_generation: profile.credential_generation,
connection: result.connection.clone(),
models: result.models.clone(),
};
let validation =
serde_json::to_value(validation).map_err(|error| RuntimeError::json("serialize BYOK validation failed", error))?;
let updated = sqlx::query(
"UPDATE ai_workspace_byok_configs SET validation = $3 WHERE workspace_id = $1 AND id = $2 AND revision = $4 AND \
credential_generation = $5",
)
.bind(&input.workspace_id)
.bind(&input.profile_id)
.bind(validation)
.bind(profile.revision)
.bind(profile.credential_generation)
.execute(pool)
.await
.map_err(|error| RuntimeError::database("record BYOK profile probe failed", error))?
.rows_affected();
result.stale = updated != 1;
Ok(result)
}
pub(in super::super) async fn probe_draft(
pool: &PgPool,
root_secret: &[u8],
policy: &CopilotByokRuntimeConfig,
input: ProbeByokDraftInput,
) -> RuntimeResult<crate::llm::ByokProbeResultOutput> {
let definition = validate_definition(&input.provider, input.definition)
.map_err(|error| RuntimeError::invalid_input(error.to_string()))?;
admit_endpoint(&definition, policy).await?;
let credential = match (input.credential, input.profile_id, input.expected_revision) {
(Some(credential), None, None) => {
require_text(&credential, "credential")?;
SensitiveCredential::new(credential.into_bytes())
}
(None, Some(profile_id), Some(expected_revision)) => {
let profile = sqlx::query_as::<_, ProfileRow>(
r#"
SELECT id, workspace_id, provider, name, description, encrypted_api_key,
definition, sort_order, enabled, revision, credential_generation, validation
FROM ai_workspace_byok_configs
WHERE workspace_id = $1 AND id = $2
"#,
)
.bind(&input.workspace_id)
.bind(&profile_id)
.fetch_optional(pool)
.await
.map_err(|error| RuntimeError::database("read BYOK profile for draft probe failed", error))?
.ok_or_else(|| RuntimeError::invalid_input("BYOK profile not found"))?;
if profile.provider != input.provider || profile.revision != expected_revision {
return Err(RuntimeError::invalid_input("byok_revision_conflict"));
}
let stored_definition = parse_definition(profile.definition)?;
envelope_key(root_secret)?
.decrypt(
&profile.encrypted_api_key,
&server_aad(
&profile.workspace_id,
&profile.id,
&profile.provider,
stored_definition.endpoint_identity(),
),
)
.map_err(|_| RuntimeError::invalid_state("credential_unavailable"))?
}
_ => {
return Err(RuntimeError::invalid_input(
"draft probe requires either credential or stored profile revision",
));
}
};
super::probe::execute_probe(&input.provider, &definition, credential, policy, input.checks).await
}
async fn select_profile_for_update(
tx: &mut sqlx::Transaction<'_, sqlx::Postgres>,
workspace_id: &str,
profile_id: &str,
) -> RuntimeResult<ProfileRow> {
sqlx::query_as::<_, ProfileRow>(
r#"
SELECT id, workspace_id, provider, name, description, encrypted_api_key,
definition, sort_order, enabled, revision, credential_generation, validation
FROM ai_workspace_byok_configs
WHERE workspace_id = $1 AND id = $2
FOR UPDATE
"#,
)
.bind(workspace_id)
.bind(profile_id)
.fetch_optional(&mut **tx)
.await
.map_err(|error| RuntimeError::database("read BYOK profile failed", error))?
.ok_or_else(|| RuntimeError::invalid_input("BYOK profile not found"))
}
pub(super) fn envelope_key(root_secret: &[u8]) -> RuntimeResult<CredentialEnvelopeKey> {
CredentialEnvelopeKey::derive(root_secret)
.map_err(|_| RuntimeError::invalid_state("stable crypto.privateKey is required for persistent BYOK"))
}
fn parse_definition(value: serde_json::Value) -> RuntimeResult<ByokProfileDefinition> {
serde_json::from_value(value).map_err(|error| RuntimeError::json("invalid stored BYOK definition", error))
}
fn profile_output(row: ProfileRow) -> RuntimeResult<ByokProfileOutput> {
let definition = parse_definition(row.definition)?;
let validation = parse_validation(row.validation)?;
Ok(ByokProfileOutput {
profile_id: row.id,
workspace_id: row.workspace_id,
provider: row.provider,
name: row.name,
description: row.description,
definition: definition.into(),
enabled: row.enabled,
sort_order: row.sort_order,
revision: row.revision,
validation,
})
}
fn parse_validation(value: Option<serde_json::Value>) -> RuntimeResult<Option<ByokValidationOutput>> {
value
.map(|value| {
serde_json::from_value(value).map_err(|error| RuntimeError::json("invalid stored BYOK validation", error))
})
.transpose()
}
pub(super) fn require_text(value: &str, field: &'static str) -> RuntimeResult<()> {
if value.trim().is_empty() {
Err(RuntimeError::invalid_input(format!("{field} is required")))
} else {
Ok(())
}
}
packages/backend/native/src/runtime/backend_runtime/constants.rs
-2
View File
@@ -1,5 +1,3 @@
pub(super) const BYOK_LOCAL_LEASE_ACTIVE_PURPOSE: &str = "copilot_byok_local_lease:active";
pub(super) const BYOK_LOCAL_LEASE_PURPOSE: &str = "copilot_byok_local_lease";
pub(super) const MAGIC_LINK_OTP_PURPOSE: &str = "magic_link_otp";
pub(super) const MAX_MAGIC_LINK_OTP_ATTEMPTS: i32 = 10;
pub(super) const WORKSPACE_INVITE_LINK_ID_PURPOSE: &str = "workspace_invite_link:id";
packages/backend/native/src/runtime/backend_runtime/copilot/context.rs
+266
View File
@@ -0,0 +1,266 @@
use llm_adapter::capability::provider_default_capability_upper_bound;
use sqlx::{FromRow, PgPool, Row};
use super::super::{LocalLeasePayload, RuntimeError, RuntimeResult, token_hash};
use crate::{
llm::{
CopilotAccessProjection,
byok::{ByokEndpoint, ByokModelDeclaration, ByokProfileDefinition, local_aad, server_aad},
route::{self, AuthorizedProfileRef, CatalogSlot, CredentialRef, ProfileSource},
},
runtime::{CopilotManagedProfileConfig, CopilotRuntimeConfig},
};
#[derive(FromRow)]
struct ServerProfileRow {
id: String,
workspace_id: String,
provider: String,
encrypted_api_key: String,
definition: serde_json::Value,
sort_order: i32,
}
pub(super) struct ProfileLoadInput<'a> {
pub(super) slot: &'a CatalogSlot,
pub(super) built_in_route_id: Option<&'a str>,
pub(super) workspace_id: Option<&'a str>,
pub(super) user_id: Option<&'a str>,
pub(super) local_lease_id: Option<&'a str>,
pub(super) access: &'a CopilotAccessProjection,
pub(super) managed_target_id: Option<&'a str>,
}
pub(super) async fn load_profiles(
pool: &PgPool,
config: &CopilotRuntimeConfig,
input: ProfileLoadInput<'_>,
) -> RuntimeResult<Vec<AuthorizedProfileRef>> {
let mut profiles = Vec::new();
if let Some(workspace_id) = input.workspace_id
&& input.access.server_byok
{
profiles.extend(load_server_profiles(pool, workspace_id).await?);
}
if let (Some(workspace_id), Some(user_id), Some(lease_id)) = (input.workspace_id, input.user_id, input.local_lease_id)
&& input.access.local_byok
{
profiles.extend(load_local_profiles(pool, workspace_id, user_id, lease_id).await?);
}
profiles.extend(load_managed_profiles(
config,
input.slot,
input.built_in_route_id,
input.access.managed_tier,
input.managed_target_id,
)?);
Ok(profiles)
}
async fn load_server_profiles(pool: &PgPool, workspace_id: &str) -> RuntimeResult<Vec<AuthorizedProfileRef>> {
let rows = sqlx::query_as::<_, ServerProfileRow>(
r#"
SELECT id, workspace_id, provider, encrypted_api_key, definition, sort_order
FROM ai_workspace_byok_configs
WHERE workspace_id = $1 AND enabled = TRUE
ORDER BY sort_order ASC, created_at ASC
"#,
)
.bind(workspace_id)
.fetch_all(pool)
.await
.map_err(|error| RuntimeError::database("load authorized BYOK profiles failed", error))?;
rows
.into_iter()
.map(|row| {
let definition: ByokProfileDefinition = serde_json::from_value(row.definition)
.map_err(|error| RuntimeError::json("invalid stored BYOK definition", error))?;
let aad = server_aad(
&row.workspace_id,
&row.id,
&row.provider,
definition.endpoint_identity(),
);
Ok(AuthorizedProfileRef {
profile_id: row.id,
source: ProfileSource::Server,
provider: row.provider,
definition,
sort_order: row.sort_order,
credential_ref: CredentialRef::Envelope {
encrypted: row.encrypted_api_key,
aad,
},
})
})
.collect()
}
async fn load_local_profiles(
pool: &PgPool,
workspace_id: &str,
user_id: &str,
lease_id: &str,
) -> RuntimeResult<Vec<AuthorizedProfileRef>> {
let payload = sqlx::query(
r#"
SELECT payload
FROM runtime_states
WHERE purpose = 'copilot_byok_local_lease' AND token_hash = $1
AND consumed_at IS NULL AND expires_at > clock_timestamp()
"#,
)
.bind(token_hash(lease_id))
.fetch_optional(pool)
.await
.map_err(|error| RuntimeError::database("load BYOK local lease failed", error))?
.map(|row| row.get::<serde_json::Value, _>("payload"));
let Some(payload) = payload else {
return Ok(Vec::new());
};
let payload: LocalLeasePayload =
serde_json::from_value(payload).map_err(|error| RuntimeError::json("invalid BYOK local lease", error))?;
if payload.version != 1 || payload.workspace_id != workspace_id || payload.user_id != user_id {
return Ok(Vec::new());
}
Ok(
payload
.providers
.into_iter()
.enumerate()
.filter(|(_, provider)| provider.enabled)
.map(|(index, provider)| {
let aad = local_aad(
workspace_id,
user_id,
lease_id,
index,
&provider.provider,
provider.definition.endpoint_identity(),
);
AuthorizedProfileRef {
profile_id: format!("{lease_id}:{index}"),
source: ProfileSource::Local,
provider: provider.provider,
definition: provider.definition,
sort_order: index as i32,
credential_ref: CredentialRef::Envelope {
encrypted: provider.encrypted_credential,
aad,
},
}
})
.collect(),
)
}
fn load_managed_profiles(
config: &CopilotRuntimeConfig,
slot: &CatalogSlot,
built_in_route_id: Option<&str>,
managed_tier: route::CopilotManagedTier,
managed_target_id: Option<&str>,
) -> RuntimeResult<Vec<AuthorizedProfileRef>> {
let targets = if let Some(target_id) = managed_target_id {
vec![
route::managed_selected_target(built_in_route_id, target_id, managed_tier)
.ok_or_else(|| RuntimeError::invalid_input("managed_target_unavailable"))?,
]
} else if let Some(targets) = route::managed_targets(slot, built_in_route_id, managed_tier) {
targets
} else {
return Ok(Vec::new());
};
targets
.iter()
.enumerate()
.map(|(index, model_id)| {
let matches = config
.providers
.profiles
.iter()
.filter(|profile| profile.enabled && profile.models.iter().any(|model| model == model_id))
.collect::<Vec<_>>();
let [profile] = matches.as_slice() else {
return Err(RuntimeError::invalid_state(if matches.is_empty() {
"built-in managed route model is unavailable"
} else {
"built-in managed route model matches multiple profiles"
}));
};
let capabilities = provider_default_capability_upper_bound(&profile.provider, model_id)
.ok_or_else(|| RuntimeError::invalid_state("built-in managed route model is incompatible with its profile"))?;
Ok(AuthorizedProfileRef {
profile_id: profile.id.clone(),
source: ProfileSource::Managed,
provider: profile.provider.clone(),
definition: ByokProfileDefinition {
version: 1,
endpoint: managed_endpoint(profile)?,
models: vec![ByokModelDeclaration {
model_id: model_id.clone(),
enabled: true,
capabilities,
}],
},
sort_order: index as i32,
credential_ref: CredentialRef::Managed {
profile_id: profile.id.clone(),
},
})
})
.collect()
}
fn managed_endpoint(profile: &CopilotManagedProfileConfig) -> RuntimeResult<ByokEndpoint> {
if let Some(base_url) = profile.config.get("baseURL").and_then(serde_json::Value::as_str) {
return Ok(ByokEndpoint::Custom {
url: llm_adapter::target::canonicalize_endpoint(base_url)
.map_err(|error| RuntimeError::invalid_state(error.to_string()))?,
});
}
let endpoint = match profile.provider.as_str() {
"geminiVertex" | "anthropicVertex" => {
let location = required_config_text(profile, "location")?;
let project = required_config_text(profile, "project")?;
let publisher = if profile.provider == "geminiVertex" {
"google"
} else {
"anthropic"
};
format!(
"https://{location}-aiplatform.googleapis.com/v1/projects/{project}/locations/{location}/publishers/{publisher}"
)
}
"cloudflareWorkersAi" => format!(
"https://api.cloudflare.com/client/v4/accounts/{}/ai",
required_config_text(profile, "accountId")?
),
_ => return Ok(ByokEndpoint::ProviderDefault),
};
Ok(ByokEndpoint::Custom { url: endpoint })
}
pub(super) fn managed_profile<'a>(
config: &'a CopilotRuntimeConfig,
profile_id: &str,
) -> RuntimeResult<&'a CopilotManagedProfileConfig> {
config
.providers
.profiles
.iter()
.find(|profile| profile.id == profile_id && profile.enabled)
.ok_or_else(|| RuntimeError::invalid_state("managed copilot credential unavailable"))
}
pub(super) fn required_config_text<'a>(
profile: &'a CopilotManagedProfileConfig,
field: &'static str,
) -> RuntimeResult<&'a str> {
profile
.config
.get(field)
.and_then(serde_json::Value::as_str)
.filter(|value| !value.trim().is_empty())
.ok_or_else(|| RuntimeError::invalid_state(format!("managed copilot profile requires {field}")))
}
packages/backend/native/src/runtime/backend_runtime/copilot/dispatch.rs
+407
View File
@@ -0,0 +1,407 @@
use std::{
collections::{HashMap, HashSet},
sync::Arc,
};
use gcp_auth::{CustomServiceAccount, TokenProvider};
use llm_adapter::{
backend::{BackendError, DefaultHttpClient},
capability::{AttachmentKind, AttachmentSource},
core::{CoreContent, ImageInput, ImageRequest},
router::{ExecutablePreparedRoute, ExecutableProtocol, ExecutableRequest, ExecutableResponse},
target::{
BackendCredential, BackendEndpoint, BackendOperation, BackendProtocol, BackendProvider, BackendTargetInput,
EgressPolicy, compile_backend_target,
},
};
use llm_runtime::{CompiledPlan, CompiledRoute, RuntimeRouteEvent, RuntimeUsage, dispatch_compiled_plan};
use serde::Serialize;
use uuid::Uuid;
use zeroize::Zeroizing;
use super::{COPILOT_REQUEST_TIMEOUT, RuntimeError, RuntimeResult, context};
use crate::{
llm::{
byok::{ByokEndpoint, CredentialEnvelopeKey},
route::{
AuthorizedProfileRef, AuthorizedTargetRef, CatalogSlot, CredentialRef, RouteOperation, with_request_requirements,
},
},
runtime::{BackendRuntimeConfig, CopilotManagedProfileConfig},
};
#[derive(Serialize)]
#[serde(rename_all = "snake_case", tag = "type")]
pub(super) enum ProductEvent {
RouteSelected { route: RouteIdentity },
RouteFailed { route: RouteIdentity, error_kind: String },
Usage { route: RouteIdentity, usage: ProductUsage },
}
#[derive(Clone, Serialize)]
#[serde(rename_all = "camelCase")]
pub(super) struct RouteIdentity {
profile_id: String,
source: &'static str,
provider: String,
model: String,
}
#[derive(Serialize)]
#[serde(untagged)]
pub(super) enum ProductUsage {
Tokens(llm_adapter::core::CoreUsage),
Image(llm_adapter::core::ImageUsage),
}
#[derive(Serialize)]
#[serde(rename_all = "camelCase")]
pub(super) struct CopilotExecutionResult {
events: Vec<ProductEvent>,
result: serde_json::Value,
}
pub(super) struct CompiledExecution {
pub(super) plan: CompiledPlan,
identities: HashMap<String, RouteIdentity>,
}
impl CompiledExecution {
pub(super) fn project(&self, event: RuntimeRouteEvent) -> RuntimeResult<ProductEvent> {
project_event(event, &self.identities)
}
}
pub(super) fn request_and_slot(
slot: CatalogSlot,
request: serde_json::Value,
) -> RuntimeResult<(CatalogSlot, ExecutableRequest)> {
let executable = match slot.operation {
RouteOperation::Chat => ExecutableRequest::Chat(parse_request(request)?),
RouteOperation::Structured | RouteOperation::Transcription => {
ExecutableRequest::Structured(parse_request(request)?)
}
RouteOperation::Embedding => ExecutableRequest::Embedding(parse_request(request)?),
RouteOperation::Rerank => ExecutableRequest::Rerank(parse_request(request)?),
RouteOperation::Image => ExecutableRequest::Image(Box::new(parse_request(request)?)),
};
let (needs_tools, attachment_kinds, attachment_sources) = request_requirements(&executable);
Ok((
with_request_requirements(slot, needs_tools, attachment_kinds, attachment_sources),
executable,
))
}
pub(super) fn execute(
config: Arc<BackendRuntimeConfig>,
slot: CatalogSlot,
request: ExecutableRequest,
profiles: Vec<AuthorizedProfileRef>,
candidates: Vec<AuthorizedTargetRef>,
managed_credentials: HashMap<String, Zeroizing<String>>,
) -> RuntimeResult<CopilotExecutionResult> {
let execution = compile_execution(&config, slot, request, &profiles, &candidates, &managed_credentials)?;
let mut runtime_events = Vec::new();
let response = dispatch_compiled_plan(&DefaultHttpClient::default(), &execution.plan, |event| {
runtime_events.push(event)
})
.map_err(|error| RuntimeError::invalid_state(error.to_string()))?;
let events = runtime_events
.into_iter()
.map(|event| execution.project(event))
.collect::<RuntimeResult<Vec<_>>>()?;
Ok(CopilotExecutionResult {
events,
result: response_value(response)?,
})
}
pub(super) fn compile_execution(
config: &BackendRuntimeConfig,
slot: CatalogSlot,
request: ExecutableRequest,
profiles: &[AuthorizedProfileRef],
candidates: &[AuthorizedTargetRef],
managed_credentials: &HashMap<String, Zeroizing<String>>,
) -> RuntimeResult<CompiledExecution> {
let key = CredentialEnvelopeKey::derive(config.private_key.as_bytes())
.map_err(|_| RuntimeError::invalid_state("credential_unavailable"))?;
let mut identities = HashMap::new();
let mut routes = Vec::with_capacity(candidates.len());
for candidate in candidates {
let profile = profiles
.get(candidate.profile_index)
.ok_or_else(|| RuntimeError::invalid_state("invalid authorized route profile"))?;
let model = profile
.definition
.models
.get(candidate.model_index)
.ok_or_else(|| RuntimeError::invalid_state("invalid authorized route model"))?;
let credential = resolve_credential(&key, profile, managed_credentials)?;
let target = compile_backend_target(BackendTargetInput {
provider: provider(&profile.provider)?,
operation: operation(slot.operation),
endpoint: endpoint(&profile.provider, &profile.definition.endpoint),
model: model.model_id.clone(),
credential: BackendCredential::new(credential),
timeout_ms: Some(COPILOT_REQUEST_TIMEOUT.as_millis() as u64),
egress_policy: if profile.source != crate::llm::route::ProfileSource::Managed
&& config.copilot.byok.allow_private_endpoint
{
EgressPolicy::AllowPrivate
} else {
EgressPolicy::PublicOnly
},
})
.map_err(|error| RuntimeError::invalid_state(error.to_string()))?;
let route_id = Uuid::new_v4().to_string();
identities.insert(
route_id.clone(),
RouteIdentity {
profile_id: profile.profile_id.clone(),
source: match profile.source {
crate::llm::route::ProfileSource::Server => "server",
crate::llm::route::ProfileSource::Local => "local",
crate::llm::route::ProfileSource::Managed => "affine_cloud",
},
provider: profile.provider.clone(),
model: target.model.clone(),
},
);
let protocol = protocol(target.protocol);
let route = ExecutablePreparedRoute::new(protocol, target.model, target.config, request.clone())
.map_err(|error| RuntimeError::invalid_input(error.to_string()))?;
routes.push(CompiledRoute::new(route_id, route));
}
Ok(CompiledExecution {
plan: CompiledPlan::new(routes).map_err(|error| RuntimeError::invalid_state(error.to_string()))?,
identities,
})
}
fn parse_request<T: serde::de::DeserializeOwned>(value: serde_json::Value) -> RuntimeResult<T> {
serde_json::from_value(value).map_err(|error| RuntimeError::json("invalid copilot execution request", error))
}
fn request_requirements(request: &ExecutableRequest) -> (bool, Vec<AttachmentKind>, Vec<AttachmentSource>) {
let mut kinds = HashSet::new();
let mut sources = HashSet::new();
let needs_tools = match request {
ExecutableRequest::Chat(request) => {
collect_message_attachments(&request.messages, &mut kinds, &mut sources);
!request.tools.is_empty()
}
ExecutableRequest::Structured(request) => {
collect_message_attachments(&request.messages, &mut kinds, &mut sources);
false
}
ExecutableRequest::Image(request) => {
if let ImageRequest::Edit(request) = request.as_ref() {
kinds.insert(AttachmentKind::Image);
for image in &request.images {
sources.insert(match image {
ImageInput::Url { .. } => AttachmentSource::Url,
ImageInput::Data { .. } => AttachmentSource::Data,
ImageInput::Bytes { .. } => AttachmentSource::Bytes,
});
}
}
false
}
ExecutableRequest::Embedding(_) | ExecutableRequest::Rerank(_) => false,
};
(needs_tools, kinds.into_iter().collect(), sources.into_iter().collect())
}
fn collect_message_attachments(
messages: &[llm_adapter::core::CoreMessage],
kinds: &mut HashSet<AttachmentKind>,
sources: &mut HashSet<AttachmentSource>,
) {
for content in messages.iter().flat_map(|message| &message.content) {
let source = match content {
CoreContent::Image { source } => {
kinds.insert(AttachmentKind::Image);
source
}
CoreContent::Audio { source } => {
kinds.insert(AttachmentKind::Audio);
source
}
CoreContent::File { source } => {
kinds.insert(AttachmentKind::File);
source
}
_ => continue,
};
let source_kind = if source.get("url").is_some() {
AttachmentSource::Url
} else if source.get("data").is_some() || source.get("data_base64").is_some() {
AttachmentSource::Data
} else if source.get("bytes").is_some() {
AttachmentSource::Bytes
} else {
AttachmentSource::FileHandle
};
sources.insert(source_kind);
}
}
fn resolve_credential(
key: &CredentialEnvelopeKey,
profile: &AuthorizedProfileRef,
managed_credentials: &HashMap<String, Zeroizing<String>>,
) -> RuntimeResult<String> {
match &profile.credential_ref {
CredentialRef::Envelope { encrypted, aad } => key
.decrypt(encrypted, aad)
.map_err(|_| RuntimeError::invalid_state("credential_unavailable"))
.and_then(|credential| {
String::from_utf8(credential.expose().to_vec())
.map_err(|_| RuntimeError::invalid_state("credential_unavailable"))
}),
CredentialRef::Managed { profile_id } => managed_credentials
.get(profile_id)
.map(|credential| credential.as_str().to_string())
.ok_or_else(|| RuntimeError::invalid_state("managed copilot credential unavailable")),
}
}
pub(super) async fn managed_credential(
profile: &CopilotManagedProfileConfig,
token_provider: Option<Arc<dyn TokenProvider>>,
) -> RuntimeResult<String> {
if matches!(profile.provider.as_str(), "geminiVertex" | "anthropicVertex") {
return token_provider
.ok_or_else(|| RuntimeError::invalid_state("managed Vertex credential unavailable"))?
.token(&["https://www.googleapis.com/auth/cloud-platform"])
.await
.map(|token| token.as_str().to_string())
.map_err(|_| RuntimeError::invalid_state("managed Vertex credential unavailable"));
}
let field = if profile.provider == "cloudflareWorkersAi" {
"apiToken"
} else {
"apiKey"
};
Ok(context::required_config_text(profile, field)?.to_string())
}
pub(super) async fn create_vertex_token_provider(
profile: &CopilotManagedProfileConfig,
) -> RuntimeResult<Arc<dyn TokenProvider>> {
if let Some(credentials) = profile.config.pointer("/googleAuthOptions/credentials") {
let project = context::required_config_text(profile, "project")?.to_string();
let mut credentials = credentials.clone();
let object = credentials
.as_object_mut()
.ok_or_else(|| RuntimeError::invalid_state("managed Vertex credentials must be an object"))?;
object
.entry("type")
.or_insert_with(|| serde_json::Value::String("service_account".to_string()));
object
.entry("project_id")
.or_insert_with(|| serde_json::Value::String(project));
object
.entry("token_uri")
.or_insert_with(|| serde_json::Value::String("https://oauth2.googleapis.com/token".to_string()));
let json = serde_json::to_string(&credentials)
.map_err(|error| RuntimeError::json("serialize managed Vertex credentials failed", error))?;
return CustomServiceAccount::from_json(&json)
.map(|provider| Arc::new(provider) as Arc<dyn TokenProvider>)
.map_err(|_| RuntimeError::invalid_state("managed Vertex credential unavailable"));
}
gcp_auth::provider()
.await
.map_err(|_| RuntimeError::invalid_state("managed Vertex credential unavailable"))
}
pub(in crate::runtime::backend_runtime) fn provider(value: &str) -> RuntimeResult<BackendProvider> {
match value {
"openai" => Ok(BackendProvider::OpenAi),
"anthropic" => Ok(BackendProvider::Anthropic),
"anthropicVertex" => Ok(BackendProvider::AnthropicVertex),
"gemini" => Ok(BackendProvider::Gemini),
"geminiVertex" => Ok(BackendProvider::GeminiVertex),
"cloudflareWorkersAi" => Ok(BackendProvider::CloudflareWorkersAi),
"fal" => Ok(BackendProvider::Fal),
_ => Err(RuntimeError::invalid_state("unsupported copilot provider")),
}
}
fn operation(value: RouteOperation) -> BackendOperation {
match value {
RouteOperation::Chat => BackendOperation::Chat,
RouteOperation::Structured | RouteOperation::Transcription => BackendOperation::Structured,
RouteOperation::Embedding => BackendOperation::Embedding,
RouteOperation::Rerank => BackendOperation::Rerank,
RouteOperation::Image => BackendOperation::Image,
}
}
pub(in crate::runtime::backend_runtime) fn endpoint(provider: &str, value: &ByokEndpoint) -> BackendEndpoint {
match (provider, value) {
("anthropic", ByokEndpoint::ProviderDefault) => BackendEndpoint::Custom("https://api.anthropic.com".to_string()),
("openai" | "anthropic", ByokEndpoint::Custom { url }) => {
BackendEndpoint::Custom(url.strip_suffix("/v1").unwrap_or(url).to_string())
}
(_, ByokEndpoint::ProviderDefault) => BackendEndpoint::ProviderDefault,
(_, ByokEndpoint::Custom { url }) => BackendEndpoint::Custom(url.clone()),
}
}
pub(in crate::runtime::backend_runtime) fn protocol(value: BackendProtocol) -> ExecutableProtocol {
match value {
BackendProtocol::Chat(value) => ExecutableProtocol::Chat(value),
BackendProtocol::Structured(value) => ExecutableProtocol::Structured(value),
BackendProtocol::Embedding(value) => ExecutableProtocol::Embedding(value),
BackendProtocol::Rerank(value) => ExecutableProtocol::Rerank(value),
BackendProtocol::Image(value) => ExecutableProtocol::Image(value),
}
}
fn project_event(event: RuntimeRouteEvent, identities: &HashMap<String, RouteIdentity>) -> RuntimeResult<ProductEvent> {
match event {
RuntimeRouteEvent::Selected { route_id } => Ok(ProductEvent::RouteSelected {
route: identity(identities, &route_id)?,
}),
RuntimeRouteEvent::Failed { route_id, error_kind } => Ok(ProductEvent::RouteFailed {
route: identity(identities, &route_id)?,
error_kind,
}),
RuntimeRouteEvent::Usage { route_id, usage } => Ok(ProductEvent::Usage {
route: identity(identities, &route_id)?,
usage: match usage {
RuntimeUsage::Tokens(usage) => ProductUsage::Tokens(usage),
RuntimeUsage::Image(usage) => ProductUsage::Image(usage),
},
}),
}
}
fn identity(identities: &HashMap<String, RouteIdentity>, route_id: &str) -> RuntimeResult<RouteIdentity> {
identities
.get(route_id)
.cloned()
.ok_or_else(|| RuntimeError::invalid_state("runtime emitted an unknown route id"))
}
fn response_value(response: ExecutableResponse) -> RuntimeResult<serde_json::Value> {
match response {
ExecutableResponse::Chat(response) => serialize_response(response),
ExecutableResponse::Structured(response) => serialize_response(response),
ExecutableResponse::Embedding(response) => serialize_response(response),
ExecutableResponse::Rerank(response) => serialize_response(response),
ExecutableResponse::Image(response) => serialize_response(response),
}
}
fn serialize_response(value: impl Serialize) -> RuntimeResult<serde_json::Value> {
serde_json::to_value(value).map_err(|error| RuntimeError::json("serialize copilot response failed", error))
}
impl From<BackendError> for RuntimeError {
fn from(error: BackendError) -> Self {
RuntimeError::invalid_state(error.to_string())
}
}
packages/backend/native/src/runtime/backend_runtime/copilot/mod.rs
+256
View File
@@ -0,0 +1,256 @@
mod context;
mod dispatch;
mod stream;
use std::{
collections::HashMap,
sync::{Arc, RwLock},
time::Duration,
};
use gcp_auth::TokenProvider;
use sha2::{Digest, Sha256};
use tokio::sync::OnceCell;
use zeroize::Zeroizing;
pub(in crate::runtime::backend_runtime) use dispatch::{
endpoint as byok_endpoint, protocol as executable_protocol, provider as backend_provider,
};
use super::{BackendRuntime, RuntimeError, RuntimeResult, to_napi_error};
use crate::{
llm::{
CopilotExecuteInput, CopilotRouteCheckInput,
route::{self, AuthorizedProfileRef, AuthorizedTargetRef, CredentialRef},
},
runtime::{BackendRuntimeConfig, CopilotManagedProfileConfig},
};
pub(super) type ManagedTokenProviderCache = RwLock<HashMap<String, Arc<OnceCell<Arc<dyn TokenProvider>>>>>;
pub(super) const COPILOT_REQUEST_TIMEOUT: Duration = Duration::from_secs(30 * 60);
struct AuthorizedCopilotRoute {
config: std::sync::Arc<BackendRuntimeConfig>,
slot: route::CatalogSlot,
profiles: Vec<route::AuthorizedProfileRef>,
candidates: Vec<route::AuthorizedTargetRef>,
}
#[napi_derive::napi]
impl BackendRuntime {
#[napi]
pub async fn execute_copilot(&self, input: CopilotExecuteInput) -> napi::Result<String> {
self.execute_copilot_inner(input).await.map_err(to_napi_error)
}
#[napi]
pub async fn assert_copilot_route(&self, input: CopilotRouteCheckInput) -> napi::Result<()> {
let slot = route::slot(&input.slot).ok_or_else(|| RuntimeError::invalid_input("unknown copilot route slot"));
self
.authorize_copilot_route(input, slot.map_err(to_napi_error)?)
.await
.map(|_| ())
.map_err(to_napi_error)
}
}
impl BackendRuntime {
async fn execute_copilot_inner(&self, input: CopilotExecuteInput) -> RuntimeResult<String> {
let (config, slot, request, profiles, candidates, managed_credentials) = self.prepare_copilot(input).await?;
let output = tokio::task::spawn_blocking(move || {
dispatch::execute(config, slot, request, profiles, candidates, managed_credentials)
})
.await
.map_err(|error| RuntimeError::invalid_state(format!("copilot execution task failed: {error}")))??;
serde_json::to_string(&output).map_err(|error| RuntimeError::json("serialize copilot execution failed", error))
}
pub(super) async fn prepare_copilot(
&self,
input: CopilotExecuteInput,
) -> RuntimeResult<stream::PreparedCopilotExecution> {
let CopilotExecuteInput {
slot,
built_in_route_id,
workspace_id,
user_id,
local_lease_id,
access,
managed_target_id,
target_override,
request,
} = input;
let base_slot = route::slot(&slot).ok_or_else(|| RuntimeError::invalid_input("unknown copilot route slot"))?;
let (slot, request) = dispatch::request_and_slot(base_slot, request)?;
let authorized = self
.authorize_copilot_route(
CopilotRouteCheckInput {
slot: slot.id.to_string(),
built_in_route_id,
workspace_id,
user_id,
local_lease_id,
access,
managed_target_id,
target_override,
},
slot,
)
.await?;
let managed_credentials = self
.resolve_managed_credentials(&authorized.config, &authorized.profiles, &authorized.candidates)
.await?;
Ok((
authorized.config,
authorized.slot,
request,
authorized.profiles,
authorized.candidates,
managed_credentials,
))
}
async fn resolve_managed_credentials(
&self,
config: &BackendRuntimeConfig,
profiles: &[AuthorizedProfileRef],
candidates: &[AuthorizedTargetRef],
) -> RuntimeResult<HashMap<String, Zeroizing<String>>> {
let mut credentials = HashMap::new();
for candidate in candidates {
let profile = profiles
.get(candidate.profile_index)
.ok_or_else(|| RuntimeError::invalid_state("invalid authorized route profile"))?;
let CredentialRef::Managed { profile_id } = &profile.credential_ref else {
continue;
};
if credentials.contains_key(profile_id) {
continue;
}
let managed = context::managed_profile(&config.copilot, profile_id)?;
let token_provider = if matches!(managed.provider.as_str(), "geminiVertex" | "anthropicVertex") {
Some(self.managed_token_provider(managed).await?)
} else {
None
};
credentials.insert(
profile_id.clone(),
Zeroizing::new(dispatch::managed_credential(managed, token_provider).await?),
);
}
Ok(credentials)
}
async fn managed_token_provider(
&self,
profile: &CopilotManagedProfileConfig,
) -> RuntimeResult<Arc<dyn TokenProvider>> {
let config = serde_json::to_vec(&profile.config)
.map_err(|error| RuntimeError::json("serialize managed Vertex profile failed", error))?;
let cache_key = format!(
"{}:{}:{}",
profile.id,
profile.provider,
hex::encode(Sha256::digest(config))
);
let cell = {
let mut providers = self
.managed_token_providers
.write()
.map_err(|_| RuntimeError::invalid_state("managed token provider cache lock poisoned"))?;
Arc::clone(providers.entry(cache_key).or_insert_with(|| Arc::new(OnceCell::new())))
};
cell
.get_or_try_init(|| dispatch::create_vertex_token_provider(profile))
.await
.map(Arc::clone)
}
async fn authorize_copilot_route(
&self,
input: CopilotRouteCheckInput,
slot: route::CatalogSlot,
) -> RuntimeResult<AuthorizedCopilotRoute> {
let config = self.config()?;
if !config.copilot.enabled {
return Err(RuntimeError::invalid_state("copilot_disabled"));
}
let deployment = if std::env::var("DEPLOYMENT_TYPE").as_deref() == Ok("selfhosted") {
route::Deployment::SelfHosted
} else {
route::Deployment::Cloud
};
let profiles = context::load_profiles(
&self.pool().await?,
&config.copilot,
context::ProfileLoadInput {
slot: &slot,
built_in_route_id: input.built_in_route_id.as_deref(),
workspace_id: input.workspace_id.as_deref(),
user_id: input.user_id.as_deref(),
local_lease_id: input.local_lease_id.as_deref(),
access: &input.access,
managed_target_id: input.managed_target_id.as_deref(),
},
)
.await?;
if input.managed_target_id.is_some() && input.target_override.is_some() {
return Err(RuntimeError::invalid_input("multiple_target_selections"));
}
let target_override = if let Some(target_id) = input.managed_target_id.as_deref() {
let model_id =
route::managed_selected_target(input.built_in_route_id.as_deref(), target_id, input.access.managed_tier)
.ok_or_else(|| RuntimeError::invalid_input("managed_target_unavailable"))?;
let profile = profiles
.iter()
.find(|profile| {
profile.source == route::ProfileSource::Managed
&& profile.definition.models.iter().any(|model| model.model_id == model_id)
})
.ok_or_else(|| RuntimeError::invalid_state("managed_target_unavailable"))?;
Some(route::TargetOverride {
profile_id: profile.profile_id.clone(),
model_id,
})
} else {
input.target_override.map(|target| route::TargetOverride {
profile_id: target.profile_id,
model_id: target.model_id,
})
};
let candidates = match route::decide(route::RoutePolicyInput {
slot: &slot,
deployment,
byok_enabled: config.copilot.byok.enabled,
access_available: input.access.route_allowed
|| route::quota_policy(&slot, input.built_in_route_id.as_deref()) != route::QuotaPolicy::Metered,
profiles: &profiles,
target_override: target_override.as_ref(),
target_override_managed: input.managed_target_id.is_some(),
}) {
route::RouteDecision::Ready(candidates) => candidates,
route::RouteDecision::Denied(reason) => {
return Err(RuntimeError::invalid_input(reason_name(reason)));
}
route::RouteDecision::NoRoute(reason) => {
return Err(RuntimeError::invalid_state(reason_name(reason)));
}
};
Ok(AuthorizedCopilotRoute {
config,
slot,
profiles,
candidates,
})
}
}
fn reason_name(reason: route::RouteDecisionReason) -> &'static str {
match reason {
route::RouteDecisionReason::ByokDisabled => "byok_disabled",
route::RouteDecisionReason::AccessUnavailable => "access_unavailable",
route::RouteDecisionReason::ExplicitTargetUnavailable => "explicit_target_unavailable",
route::RouteDecisionReason::NoCompatibleTarget => "no_compatible_target",
route::RouteDecisionReason::ManagedPresetUnavailable => "managed_preset_unavailable",
}
}
packages/backend/native/src/runtime/backend_runtime/copilot/stream.rs
+261
View File
@@ -0,0 +1,261 @@
use std::{
collections::HashMap,
sync::{
Arc, Mutex,
atomic::{AtomicBool, Ordering},
mpsc,
},
time::{Duration, Instant},
};
use llm_adapter::{
backend::{BackendError, DefaultHttpClient},
core::CoreMessage,
router::ExecutableRequest,
};
use llm_runtime::{
AccumulatedToolCall, RuntimeRouteEvent, ToolCallbackRequest, ToolCallbackResponse, ToolExecutionResult,
ToolLoopEvent, dispatch_compiled_round, run_tool_loop,
};
use napi::{
JsValue, Result, Status,
bindgen_prelude::{CallbackContext, PromiseRaw, Unknown},
threadsafe_function::{ThreadsafeFunction, ThreadsafeFunctionCallMode},
};
use zeroize::Zeroizing;
use super::{BackendRuntime, COPILOT_REQUEST_TIMEOUT, RuntimeError, dispatch, to_napi_error};
use crate::{
llm::{
CopilotExecuteInput,
route::{AuthorizedProfileRef, AuthorizedTargetRef, CatalogSlot},
},
runtime::BackendRuntimeConfig,
};
pub(super) type PreparedCopilotExecution = (
Arc<BackendRuntimeConfig>,
CatalogSlot,
ExecutableRequest,
Vec<AuthorizedProfileRef>,
Vec<AuthorizedTargetRef>,
HashMap<String, Zeroizing<String>>,
);
const STREAM_END: &str = "__AFFINE_COPILOT_STREAM_END__";
const TOOL_CALLBACK_TIMEOUT: Duration = Duration::from_secs(5 * 60);
const TOOL_CALLBACK_POLL_INTERVAL: Duration = Duration::from_millis(100);
#[napi_derive::napi]
pub struct CopilotStreamHandle {
aborted: Arc<AtomicBool>,
}
#[napi_derive::napi]
impl CopilotStreamHandle {
#[napi]
pub fn abort(&self) {
self.aborted.store(true, Ordering::Relaxed);
}
}
#[napi_derive::napi]
impl BackendRuntime {
#[napi]
pub async fn execute_copilot_stream(
&self,
input: CopilotExecuteInput,
max_steps: u32,
callback: ThreadsafeFunction<String, ()>,
tool_callback: ThreadsafeFunction<String, PromiseRaw<'static, String>>,
) -> Result<CopilotStreamHandle> {
let (config, slot, request, profiles, candidates, managed_credentials) =
self.prepare_copilot(input).await.map_err(to_napi_error)?;
let messages = match &request {
ExecutableRequest::Chat(request) => request.messages.clone(),
_ => {
return Err(to_napi_error(RuntimeError::invalid_input(
"copilot stream requires a chat slot",
)));
}
};
let mut execution =
dispatch::compile_execution(&config, slot, request, &profiles, &candidates, &managed_credentials)
.map_err(to_napi_error)?;
let aborted = Arc::new(AtomicBool::new(false));
let worker_aborted = aborted.clone();
tokio::task::spawn_blocking(move || {
let deadline = Instant::now() + COPILOT_REQUEST_TIMEOUT;
let result = run_stream(
&mut execution,
messages,
max_steps.max(1) as usize,
&callback,
&tool_callback,
&worker_aborted,
deadline,
);
if let Err(message) = result
&& !worker_aborted.load(Ordering::Relaxed)
{
let _ = emit_json(
&callback,
&serde_json::json!({
"type": "error",
"errorKind": "dispatch",
"message": message,
}),
);
}
let _ = callback.call(Ok(STREAM_END.to_string()), ThreadsafeFunctionCallMode::Blocking);
});
Ok(CopilotStreamHandle { aborted })
}
}
fn run_stream(
execution: &mut dispatch::CompiledExecution,
mut messages: Vec<CoreMessage>,
max_steps: usize,
callback: &ThreadsafeFunction<String, ()>,
tool_callback: &ThreadsafeFunction<String, PromiseRaw<'static, String>>,
aborted: &AtomicBool,
deadline: Instant,
) -> std::result::Result<(), String> {
let result = run_tool_loop(
&mut messages,
max_steps,
|messages| {
let mut route_events = Vec::new();
let result = dispatch_compiled_round(
&DefaultHttpClient::default(),
&mut execution.plan,
messages,
|| aborted.load(Ordering::Relaxed) || Instant::now() >= deadline,
|event| emit_json(callback, event).map_err(transport_error),
|event: RuntimeRouteEvent| route_events.push(event),
);
for event in route_events {
let event = execution.project(event).map_err(|error| error.to_string())?;
emit_json(callback, &event)?;
}
result.map_err(|error| error.to_string())
},
|call: &AccumulatedToolCall| execute_tool(tool_callback, call, aborted, deadline),
|event: &ToolLoopEvent| emit_json(callback, event),
|| "tool loop reached max steps".to_string(),
);
if !aborted.load(Ordering::Relaxed) && Instant::now() >= deadline {
Err("copilot stream deadline exceeded".to_string())
} else {
result
}
}
fn emit_json(
callback: &ThreadsafeFunction<String, ()>,
value: &impl serde::Serialize,
) -> std::result::Result<(), String> {
let value = serde_json::to_string(value).map_err(|error| error.to_string())?;
let status = callback.call(Ok(value), ThreadsafeFunctionCallMode::Blocking);
if status == Status::Ok {
Ok(())
} else {
Err(format!("copilot stream callback failed: {status}"))
}
}
fn transport_error(message: String) -> BackendError {
BackendError::Transport { message }
}
fn execute_tool(
callback: &ThreadsafeFunction<String, PromiseRaw<'static, String>>,
call: &AccumulatedToolCall,
aborted: &AtomicBool,
stream_deadline: Instant,
) -> std::result::Result<ToolExecutionResult, String> {
let request = serde_json::to_string(&ToolCallbackRequest {
call_id: call.id.clone(),
name: call.name.clone(),
args: call.args.clone(),
raw_arguments_text: call.raw_arguments_text.clone(),
argument_parse_error: call.argument_parse_error.clone(),
})
.map_err(|error| error.to_string())?;
let (sender, receiver) = mpsc::sync_channel(1);
let sender = Arc::new(Mutex::new(Some(sender)));
let callback_sender = sender.clone();
let status = callback.call_with_return_value(
Ok(request),
ThreadsafeFunctionCallMode::NonBlocking,
move |promise, _env| {
match promise {
Ok(promise) => {
let success_sender = callback_sender.clone();
let failure_sender = callback_sender.clone();
match promise.then(move |ctx| {
send_tool_result(
&success_sender,
serde_json::from_str::<ToolCallbackResponse>(&ctx.value).map_err(|error| error.to_string()),
);
Ok(())
}) {
Ok(promise) => {
if let Err(error) = promise.catch(move |ctx: CallbackContext<Unknown>| {
let message = ctx.value.coerce_to_string()?.into_utf8()?.as_str()?.to_string();
send_tool_result(&failure_sender, Err(message));
Ok(())
}) {
send_tool_result(&callback_sender, Err(error.to_string()));
}
}
Err(error) => send_tool_result(&callback_sender, Err(error.to_string())),
}
}
Err(error) => send_tool_result(&callback_sender, Err(error.to_string())),
}
Ok(())
},
);
if status != Status::Ok {
return Err(format!("copilot tool callback failed: {status}"));
}
let tool_deadline = std::cmp::min(stream_deadline, Instant::now() + TOOL_CALLBACK_TIMEOUT);
let response = loop {
if aborted.load(Ordering::Relaxed) {
return Err("copilot stream aborted".to_string());
}
let now = Instant::now();
if now >= tool_deadline {
return Err("copilot tool callback deadline exceeded".to_string());
}
match receiver.recv_timeout(std::cmp::min(TOOL_CALLBACK_POLL_INTERVAL, tool_deadline - now)) {
Ok(response) => break response?,
Err(mpsc::RecvTimeoutError::Timeout) => continue,
Err(mpsc::RecvTimeoutError::Disconnected) => {
return Err("copilot tool callback closed before completion".to_string());
}
}
};
if !response.args.is_object() {
return Err("copilot tool callback args must be an object".to_string());
}
Ok(ToolExecutionResult {
call_id: response.call_id,
name: response.name,
arguments: response.args,
arguments_text: response.raw_arguments_text,
arguments_error: response.argument_parse_error,
output: response.output,
is_error: response.is_error,
})
}
type ToolResultSender = Arc<Mutex<Option<mpsc::SyncSender<std::result::Result<ToolCallbackResponse, String>>>>>;
fn send_tool_result(sender: &ToolResultSender, result: std::result::Result<ToolCallbackResponse, String>) {
if let Some(sender) = sender.lock().expect("tool callback sender poisoned").take() {
let _ = sender.send(result);
}
}
packages/backend/native/src/runtime/backend_runtime/mod.rs
+134 -9
View File
@@ -1,5 +1,7 @@
mod byok;
mod constants;
mod coordination_lease;
mod copilot;
mod doc_compactor;
mod doc_storage;
mod gate;
@@ -9,8 +11,13 @@ mod runtime_state;
#[cfg(test)]
mod tests;
mod workspace_stats;
use std::{sync::RwLock, time::Duration};
use std::{
sync::{Arc, RwLock},
time::Duration,
};
use byok::LocalLeasePayload;
use copilot::{backend_provider, byok_endpoint, executable_protocol};
use napi::Result;
use sha2::{Digest, Sha256};
use sqlx::{PgPool, Row, postgres::PgPoolOptions};
@@ -22,6 +29,11 @@ pub(super) use super::{
BackendRuntimeConfig, InviteQuotaConfig, RuntimeError, RuntimeResult, migrations::migrate_runtime_tables, napi_error,
to_napi_error,
};
use crate::llm::{
ByokLocalLeaseOutput, ByokProbeResultOutput, ByokProfileOutput, CreateByokLocalLeaseInput, CreateByokProfileInput,
ProbeByokDraftInput, ProbeByokProfileInput, ReorderByokProfilesInput, ReplaceByokProfileInput,
RotateByokCredentialInput,
};
pub(super) fn token_hash(token: &str) -> String {
hex::encode(Sha256::digest(token.as_bytes()))
@@ -29,17 +41,20 @@ pub(super) fn token_hash(token: &str) -> String {
#[napi_derive::napi]
pub struct BackendRuntime {
config: RwLock<BackendRuntimeConfig>,
config: RwLock<Arc<BackendRuntimeConfig>>,
pool: Mutex<Option<PgPool>>,
managed_token_providers: copilot::ManagedTokenProviderCache,
}
#[napi_derive::napi]
impl BackendRuntime {
#[napi(constructor)]
pub fn new() -> Result<Self> {
pub fn new(private_key: Option<String>) -> Result<Self> {
let config = BackendRuntimeConfig::from_config_files(private_key).map_err(to_napi_error)?;
Ok(Self {
config: RwLock::new(BackendRuntimeConfig::from_config_files().map_err(to_napi_error)?),
config: RwLock::new(Arc::new(config)),
pool: Mutex::new(None),
managed_token_providers: Default::default(),
})
}
@@ -54,11 +69,12 @@ impl BackendRuntime {
return Ok(());
}
let database_url = self.config()?.database_url;
let config = self.config()?;
let database_url = &config.database_url;
let pool = PgPoolOptions::new()
.max_connections(5)
.acquire_timeout(Duration::from_secs(5))
.connect(&database_url)
.connect(database_url)
.await
.map_err(|err| RuntimeError::database("BackendRuntime failed to connect postgres", err))?;
@@ -83,6 +99,18 @@ impl BackendRuntime {
Ok(())
}
#[napi]
pub async fn reload_config(&self, private_key: Option<String>) -> Result<()> {
let pool = self.pool().await.map_err(to_napi_error)?;
let active_private_key = self.config().map_err(to_napi_error)?.private_key.to_string();
let config = BackendRuntimeConfig::from_config_files(private_key.or(Some(active_private_key)))
.map_err(to_napi_error)?
.with_db_overrides(&pool)
.await
.map_err(to_napi_error)?;
self.update_config(config).map_err(to_napi_error)
}
#[napi]
pub async fn health(&self) -> Result<BackendRuntimeHealth> {
let pool = self.pool.lock().await.as_ref().cloned();
@@ -107,6 +135,98 @@ impl BackendRuntime {
migrate_runtime_tables(&pool).await.map_err(to_napi_error)
}
#[napi]
pub async fn list_byok_profiles(&self, workspace_id: String) -> Result<Vec<ByokProfileOutput>> {
byok::list(&self.pool().await?, &workspace_id)
.await
.map_err(to_napi_error)
}
#[napi]
pub async fn create_byok_profile(&self, input: CreateByokProfileInput) -> Result<ByokProfileOutput> {
let config = self.config()?;
byok::create(
&self.pool().await?,
config.private_key.as_bytes(),
&config.copilot.byok,
input,
)
.await
.map_err(to_napi_error)
}
#[napi]
pub async fn replace_byok_profile(&self, input: ReplaceByokProfileInput) -> Result<ByokProfileOutput> {
let config = self.config()?;
byok::replace(
&self.pool().await?,
config.private_key.as_bytes(),
&config.copilot.byok,
input,
)
.await
.map_err(to_napi_error)
}
#[napi]
pub async fn rotate_byok_credential(&self, input: RotateByokCredentialInput) -> Result<ByokProfileOutput> {
let config = self.config()?;
byok::rotate(&self.pool().await?, config.private_key.as_bytes(), input)
.await
.map_err(to_napi_error)
}
#[napi]
pub async fn probe_byok_profile(&self, input: ProbeByokProfileInput) -> Result<ByokProbeResultOutput> {
let config = self.config()?;
byok::probe_profile(
&self.pool().await?,
config.private_key.as_bytes(),
&config.copilot.byok,
input,
)
.await
.map_err(to_napi_error)
}
#[napi]
pub async fn probe_byok_draft(&self, input: ProbeByokDraftInput) -> Result<ByokProbeResultOutput> {
let config = self.config()?;
byok::probe_draft(
&self.pool().await?,
config.private_key.as_bytes(),
&config.copilot.byok,
input,
)
.await
.map_err(to_napi_error)
}
#[napi]
pub async fn delete_byok_profile(&self, workspace_id: String, profile_id: String) -> Result<bool> {
byok::delete(&self.pool().await?, &workspace_id, &profile_id)
.await
.map_err(to_napi_error)
}
#[napi]
pub async fn reorder_byok_profiles(&self, input: ReorderByokProfilesInput) -> Result<Vec<ByokProfileOutput>> {
byok::reorder(&self.pool().await?, input).await.map_err(to_napi_error)
}
#[napi]
pub async fn create_byok_local_lease(&self, input: CreateByokLocalLeaseInput) -> Result<ByokLocalLeaseOutput> {
let config = self.config()?;
byok::create_local_lease(
&self.pool().await?,
config.private_key.as_bytes(),
&config.copilot.byok,
input,
)
.await
.map_err(to_napi_error)
}
pub(crate) async fn pool(&self) -> RuntimeResult<PgPool> {
self
.pool
@@ -117,19 +237,24 @@ impl BackendRuntime {
.ok_or_else(|| RuntimeError::invalid_state("BackendRuntime must be started before using postgres operations"))
}
pub(crate) fn config(&self) -> RuntimeResult<BackendRuntimeConfig> {
pub(crate) fn config(&self) -> RuntimeResult<Arc<BackendRuntimeConfig>> {
self
.config
.read()
.map(|config| config.clone())
.map(|config| Arc::clone(&config))
.map_err(|_| RuntimeError::invalid_state("BackendRuntime config lock poisoned"))
}
fn update_config(&self, config: BackendRuntimeConfig) -> RuntimeResult<()> {
self
.managed_token_providers
.write()
.map_err(|_| RuntimeError::invalid_state("managed token provider cache lock poisoned"))?
.clear();
*self
.config
.write()
.map_err(|_| RuntimeError::invalid_state("BackendRuntime config lock poisoned"))? = config;
.map_err(|_| RuntimeError::invalid_state("BackendRuntime config lock poisoned"))? = Arc::new(config);
Ok(())
}
}
packages/backend/native/src/runtime/backend_runtime/rolling_quota/mod.rs
+4 -3
View File
@@ -75,8 +75,9 @@ impl BackendRuntime {
&self,
input: RuntimeMailDeliveryQuotaInput,
) -> Result<RuntimeMailDeliveryQuotaDecision> {
let config = self.config()?.invite_quota;
let Some(class) = mail_class(&input.mail_name, &config) else {
let runtime_config = self.config()?;
let config = &runtime_config.invite_quota;
let Some(class) = mail_class(&input.mail_name, config) else {
return Ok(RuntimeMailDeliveryQuotaDecision {
allowed: false,
reservation_id: None,
@@ -91,7 +92,7 @@ impl BackendRuntime {
});
};
let pool = self.pool().await?;
let scopes = build_mail_scopes(&input, class, &config);
let scopes = build_mail_scopes(&input, class, config);
match reserve_scopes(&pool, "mail_delivery", input.request_id.as_deref(), scopes).await? {
Ok(reservation) => Ok(RuntimeMailDeliveryQuotaDecision {
allowed: true,
packages/backend/native/src/runtime/backend_runtime/rolling_quota/workspace_invite.rs
+6 -5
View File
@@ -293,14 +293,15 @@ impl BackendRuntime {
if input.target_count <= 0 {
return Err(napi_error("target_count must be positive"));
}
let config = self.config()?.invite_quota;
let runtime_config = self.config()?;
let config = &runtime_config.invite_quota;
let pool = self.pool().await?;
let now: DateTime<Utc> = sqlx::query_scalar("SELECT clock_timestamp()")
.fetch_one(&pool)
.await
.map_err(|err| RuntimeError::database("failed to read database clock", err))?;
let actor = load_actor(&pool, &input.actor_user_id).await?;
let actor_subject = subject_hash(&actor.email, &config);
let actor_subject = subject_hash(&actor.email, config);
if let Some(status) = active_subject_status(&pool, &actor_subject).await?
&& matches!(status.as_str(), "banned" | "quarantined")
{
@@ -387,14 +388,14 @@ impl BackendRuntime {
action_required: None,
});
}
if let Some(abuse_decision) = high_confidence_invite_abuse(&input, &actor, &config) {
if let Some(abuse_decision) = high_confidence_invite_abuse(&input, &actor, config) {
let reason = abuse_decision.reason;
let scope_key = match abuse_decision.subject_kind {
"workspace" => format!("invite:workspace_subject:{}", abuse_decision.subject_key),
"source_prefix_domain" => format!("invite:source_cohort_subject:{}", abuse_decision.subject_key),
_ => format!("invite:actor_subject:{}", abuse_decision.subject_key),
};
let action_required = record_invite_abuse_action(&pool, &input, &actor, abuse_decision, &config).await?;
let action_required = record_invite_abuse_action(&pool, &input, &actor, abuse_decision, config).await?;
return Ok(RuntimeWorkspaceInviteQuotaDecision {
allowed: false,
reservation_id: None,
@@ -411,7 +412,7 @@ impl BackendRuntime {
let workspace = load_workspace(&pool, &input.workspace_id).await?;
let activity = load_invite_activity(&pool, &input.actor_user_id, &input.workspace_id).await?;
let scopes = build_invite_scopes(&input, &actor, &workspace, &quota, &activity, &config, now)?;
let scopes = build_invite_scopes(&input, &actor, &workspace, &quota, &activity, config, now)?;
match reserve_scopes(&pool, "workspace_invite", input.request_id.as_deref(), scopes).await? {
Ok(reservation) => Ok(RuntimeWorkspaceInviteQuotaDecision {
allowed: true,
packages/backend/native/src/runtime/backend_runtime/runtime_state/byok_local_lease.rs
-128
View File
@@ -1,128 +0,0 @@
use super::{
BYOK_LOCAL_LEASE_ACTIVE_PURPOSE, BYOK_LOCAL_LEASE_PURPOSE, Result, RuntimeByokLocalLeaseRecord, RuntimeError,
dto::{RuntimeStateInsertPayload, RuntimeStatePayloadRow, RuntimeStateRows},
};
pub(super) async fn get(rows: &RuntimeStateRows, lease_id: String) -> Result<Option<RuntimeByokLocalLeaseRecord>> {
get_lease_by_id(rows, &lease_id).await
}
pub(super) async fn create(
rows: &RuntimeStateRows,
active_key: String,
lease_id: String,
payload: serde_json::Value,
ttl_ms: i64,
) -> Result<RuntimeByokLocalLeaseRecord> {
if ttl_ms <= 0 {
return Err(RuntimeError::invalid_input("BYOK local lease ttl must be positive"));
}
let mut tx = rows.begin("RuntimeState BYOK local lease").await?;
sqlx::query("SELECT pg_advisory_xact_lock(hashtextextended($1, 0))")
.bind(&active_key)
.execute(&mut *tx)
.await
.map_err(|err| RuntimeError::database("RuntimeState BYOK local lease active lock failed", err))?;
if let Some(active) = rows
.active_payload_with_expires_for_update_in_tx(
&mut tx,
BYOK_LOCAL_LEASE_ACTIVE_PURPOSE,
&active_key,
"RuntimeState BYOK local lease active get",
)
.await?
{
let existing_lease = match active.payload.get("leaseId").and_then(serde_json::Value::as_str) {
Some(existing_lease_id) => get_lease_by_id_in_tx(rows, &mut tx, existing_lease_id).await?,
None => None,
};
if let Some(lease) = existing_lease {
tx.commit()
.await
.map_err(|err| RuntimeError::database("RuntimeState BYOK local lease transaction commit failed", err))?;
return Ok(lease);
}
rows
.delete_by_key_in_tx(
&mut tx,
BYOK_LOCAL_LEASE_ACTIVE_PURPOSE,
&active_key,
"RuntimeState BYOK local lease stale active delete",
)
.await?;
}
let expires_at_ms = rows
.insert_payload_returning_expires_in_tx(
&mut tx,
RuntimeStateInsertPayload {
purpose: BYOK_LOCAL_LEASE_PURPOSE,
token: &lease_id,
lookup_key: &active_key,
payload: &payload,
ttl_ms,
context: "RuntimeState BYOK local lease create",
},
)
.await?;
let active_payload = serde_json::json!({ "leaseId": lease_id });
rows
.insert_payload_returning_expires_in_tx(
&mut tx,
RuntimeStateInsertPayload {
purpose: BYOK_LOCAL_LEASE_ACTIVE_PURPOSE,
token: &active_key,
lookup_key: &active_key,
payload: &active_payload,
ttl_ms,
context: "RuntimeState BYOK local lease active create",
},
)
.await?;
tx.commit()
.await
.map_err(|err| RuntimeError::database("RuntimeState BYOK local lease transaction commit failed", err))?;
Ok(RuntimeByokLocalLeaseRecord {
lease_id,
payload,
expires_at_ms,
})
}
async fn get_lease_by_id(rows: &RuntimeStateRows, lease_id: &str) -> Result<Option<RuntimeByokLocalLeaseRecord>> {
rows
.active_payload_with_expires(BYOK_LOCAL_LEASE_PURPOSE, lease_id, "RuntimeState BYOK local lease get")
.await?
.map(|row| record_from_row(lease_id, row))
.transpose()
}
async fn get_lease_by_id_in_tx(
rows: &RuntimeStateRows,
tx: &mut sqlx::Transaction<'_, sqlx::Postgres>,
lease_id: &str,
) -> Result<Option<RuntimeByokLocalLeaseRecord>> {
rows
.active_payload_with_expires_for_update_in_tx(
tx,
BYOK_LOCAL_LEASE_PURPOSE,
lease_id,
"RuntimeState BYOK local lease get",
)
.await?
.map(|row| record_from_row(lease_id, row))
.transpose()
}
fn record_from_row(lease_id: &str, row: RuntimeStatePayloadRow) -> Result<RuntimeByokLocalLeaseRecord> {
Ok(RuntimeByokLocalLeaseRecord {
lease_id: lease_id.to_string(),
payload: row.payload,
expires_at_ms: row.expires_at_ms,
})
}
packages/backend/native/src/runtime/backend_runtime/runtime_state/mod.rs
+3 -29
View File
@@ -1,18 +1,14 @@
use super::{BackendRuntime, RuntimeError, RuntimeResult, napi_error};
pub(super) use super::{
constants::{
BYOK_LOCAL_LEASE_ACTIVE_PURPOSE, BYOK_LOCAL_LEASE_PURPOSE, MAGIC_LINK_OTP_PURPOSE, MAX_MAGIC_LINK_OTP_ATTEMPTS,
WORKSPACE_INVITE_LINK_ID_PURPOSE, WORKSPACE_INVITE_LINK_WORKSPACE_PURPOSE,
MAGIC_LINK_OTP_PURPOSE, MAX_MAGIC_LINK_OTP_ATTEMPTS, WORKSPACE_INVITE_LINK_ID_PURPOSE,
WORKSPACE_INVITE_LINK_WORKSPACE_PURPOSE,
},
token_hash,
types::{
RuntimeByokLocalLeaseRecord, RuntimeMagicLinkOtpConsumeResult, RuntimeVerificationTokenRecord,
RuntimeWorkspaceInviteLinkRecord,
},
types::{RuntimeMagicLinkOtpConsumeResult, RuntimeVerificationTokenRecord, RuntimeWorkspaceInviteLinkRecord},
};
mod auth_challenge;
mod byok_local_lease;
mod dto;
mod invite_link;
mod magic_link_otp;
@@ -197,28 +193,6 @@ impl BackendRuntime {
.map_err(napi::Error::from)
}
#[napi]
pub async fn create_byok_local_lease(
&self,
active_key: String,
lease_id: String,
payload: serde_json::Value,
ttl_ms: i64,
) -> napi::Result<RuntimeByokLocalLeaseRecord> {
RuntimeStateStore::new(self.pool().await?)
.create_byok_local_lease(active_key, lease_id, payload, ttl_ms)
.await
.map_err(napi::Error::from)
}
#[napi]
pub async fn get_byok_local_lease(&self, lease_id: String) -> napi::Result<Option<RuntimeByokLocalLeaseRecord>> {
RuntimeStateStore::new(self.pool().await?)
.get_byok_local_lease(lease_id)
.await
.map_err(napi::Error::from)
}
#[napi]
pub async fn cleanup_expired_runtime_states(&self, limit: i64) -> napi::Result<i64> {
if limit <= 0 {
packages/backend/native/src/runtime/backend_runtime/runtime_state/store.rs
+2 -17
View File
@@ -1,9 +1,8 @@
use sqlx::PgPool;
use super::{
Result, RuntimeByokLocalLeaseRecord, RuntimeMagicLinkOtpConsumeResult, RuntimeVerificationTokenRecord,
RuntimeWorkspaceInviteLinkRecord, auth_challenge, byok_local_lease, dto::RuntimeStateRows, invite_link,
magic_link_otp, verification_token,
Result, RuntimeMagicLinkOtpConsumeResult, RuntimeVerificationTokenRecord, RuntimeWorkspaceInviteLinkRecord,
auth_challenge, dto::RuntimeStateRows, invite_link, magic_link_otp, verification_token,
};
pub(super) struct RuntimeStateStore {
@@ -121,18 +120,4 @@ impl RuntimeStateStore {
pub(super) async fn revoke_workspace_invite_link(&self, workspace_id: String) -> Result<bool> {
invite_link::revoke(&self.rows, workspace_id).await
}
pub(super) async fn create_byok_local_lease(
&self,
active_key: String,
lease_id: String,
payload: serde_json::Value,
ttl_ms: i64,
) -> Result<RuntimeByokLocalLeaseRecord> {
byok_local_lease::create(&self.rows, active_key, lease_id, payload, ttl_ms).await
}
pub(super) async fn get_byok_local_lease(&self, lease_id: String) -> Result<Option<RuntimeByokLocalLeaseRecord>> {
byok_local_lease::get(&self.rows, lease_id).await
}
}
packages/backend/native/src/runtime/backend_runtime/tests.rs
+8 -2
View File
@@ -97,11 +97,14 @@ async fn runtime_from_database_url() -> AnyResult<Option<BackendRuntime>> {
.context("cleanup invite abuse subjects for backend runtime tests")?;
Ok(Some(BackendRuntime {
config: std::sync::RwLock::new(BackendRuntimeConfig {
config: std::sync::RwLock::new(std::sync::Arc::new(BackendRuntimeConfig {
database_url,
invite_quota: Default::default(),
}),
private_key: std::sync::Arc::new(zeroize::Zeroizing::new("test-private-key".to_string())),
copilot: Default::default(),
})),
pool: Mutex::new(Some(pool)),
managed_token_providers: Default::default(),
}))
}
@@ -248,6 +251,7 @@ async fn runtime_gate_sql_semantics_are_atomic_and_ttl_bound() {
let runtime = BackendRuntime {
config: std::sync::RwLock::new(runtime.config().unwrap()),
pool: Mutex::new(Some(runtime.pool().await.unwrap())),
managed_token_providers: Default::default(),
};
tasks.push(tokio::spawn(async move {
runtime
@@ -579,6 +583,7 @@ async fn coordination_lease_sql_semantics_are_fenced_and_ttl_bound() {
let runtime = BackendRuntime {
config: std::sync::RwLock::new(runtime.config().unwrap()),
pool: Mutex::new(Some(runtime.pool().await.unwrap())),
managed_token_providers: Default::default(),
};
tasks.push(tokio::spawn(async move {
runtime
@@ -783,6 +788,7 @@ async fn verification_token_sql_state_machine_handles_keep_verify_and_cleanup()
let runtime = BackendRuntime {
config: std::sync::RwLock::new(runtime.config().unwrap()),
pool: Mutex::new(Some(runtime.pool().await.unwrap())),
managed_token_providers: Default::default(),
};
let token = concurrent_token.clone();
tasks.push(tokio::spawn(async move {
packages/backend/native/src/runtime/config.rs
+296 -35
View File
@@ -2,18 +2,71 @@ use std::{
collections::BTreeMap,
env, fs,
path::{Path, PathBuf},
sync::Arc,
};
use llm_adapter::capability::provider_default_capability_upper_bound;
use serde::Deserialize;
use serde_json::Map;
use sqlx::{PgPool, Row};
use zeroize::Zeroizing;
use super::{RuntimeError, RuntimeResult};
#[derive(Clone, Debug)]
pub(crate) struct BackendRuntimeConfig {
pub(crate) database_url: String,
pub(crate) invite_quota: InviteQuotaConfig,
pub(crate) private_key: Arc<Zeroizing<String>>,
pub(crate) copilot: CopilotRuntimeConfig,
}
#[derive(Clone, Default, Deserialize)]
#[serde(rename_all = "camelCase", default)]
pub(crate) struct CopilotRuntimeConfig {
pub(crate) enabled: bool,
pub(crate) byok: CopilotByokRuntimeConfig,
pub(crate) providers: CopilotProvidersRuntimeConfig,
}
#[derive(Clone, Deserialize)]
#[serde(rename_all = "camelCase", default)]
pub(crate) struct CopilotByokRuntimeConfig {
pub(crate) enabled: bool,
pub(crate) allow_custom_endpoint: bool,
pub(crate) allow_private_endpoint: bool,
}
impl Default for CopilotByokRuntimeConfig {
fn default() -> Self {
Self {
enabled: true,
allow_custom_endpoint: false,
allow_private_endpoint: false,
}
}
}
#[derive(Clone, Default, Deserialize)]
#[serde(rename_all = "camelCase", default)]
pub(crate) struct CopilotProvidersRuntimeConfig {
pub(crate) profiles: Vec<CopilotManagedProfileConfig>,
}
#[derive(Clone, Deserialize)]
#[serde(rename_all = "camelCase")]
pub(crate) struct CopilotManagedProfileConfig {
pub(crate) id: String,
#[serde(rename = "type")]
pub(crate) provider: String,
#[serde(default = "enabled_by_default")]
pub(crate) enabled: bool,
#[serde(default)]
pub(crate) models: Vec<String>,
pub(crate) config: serde_json::Value,
}
fn enabled_by_default() -> bool {
true
}
#[derive(Clone, Debug)]
@@ -45,35 +98,114 @@ impl Default for InviteQuotaConfig {
}
impl BackendRuntimeConfig {
pub(crate) fn from_config_files() -> RuntimeResult<Self> {
pub(crate) fn from_config_files(private_key: Option<String>) -> RuntimeResult<Self> {
let app_config = app_config_from_config_files()?;
let database_url = database_url_from_env()
.or(app_config.database_url())
.unwrap_or_else(|| "postgresql://localhost:5432/affine".to_string());
Ok(Self {
Self {
database_url,
invite_quota: app_config.invite_quota_config(),
})
private_key: Arc::new(Zeroizing::new(
private_key
.filter(|key| !key.trim().is_empty())
.or_else(private_key_from_env)
.or_else(|| app_config.crypto.as_ref().and_then(|crypto| crypto.private_key.clone()))
.unwrap_or_default(),
)),
copilot: app_config.copilot.unwrap_or_default(),
}
.validated()
}
pub(crate) async fn with_db_overrides(&self, pool: &PgPool) -> RuntimeResult<Self> {
let mut app_config = app_config_from_config_files()?;
app_config.apply_file_config(load_app_config_overrides_from_db(pool).await?);
Ok(Self {
let app_config_value = app_config_value_from_config_files()?;
let db_overrides = load_app_config_overrides_from_db(pool).await?;
self.apply_db_overrides(app_config_value, db_overrides)
}
fn apply_db_overrides(
&self,
mut app_config_value: serde_json::Value,
db_overrides: serde_json::Value,
) -> RuntimeResult<Self> {
let db_private_key = db_overrides
.pointer("/crypto/privateKey")
.and_then(serde_json::Value::as_str)
.map(str::to_string)
.and_then(non_empty_string);
merge_config_value(&mut app_config_value, db_overrides);
let app_config = deserialize_app_config(app_config_value)?;
Self {
// The DB override is loaded after this connection already exists, so it
// must not rewrite the active datasource URL.
database_url: self.database_url.clone(),
invite_quota: app_config.invite_quota_config(),
})
private_key: db_private_key
.map(|key| Arc::new(Zeroizing::new(key)))
.unwrap_or_else(|| Arc::clone(&self.private_key)),
copilot: app_config.copilot.unwrap_or_else(|| self.copilot.clone()),
}
.validated()
}
fn validated(self) -> RuntimeResult<Self> {
if self.copilot.enabled && self.copilot.byok.enabled && self.private_key.is_empty() {
return Err(RuntimeError::invalid_state(
"stable crypto.privateKey is required when persistent BYOK is enabled",
));
}
validate_copilot_config(&self.copilot)?;
Ok(self)
}
}
#[derive(Debug, Default, Deserialize)]
struct AppConfigFile {
db: Option<DbConfigFile>,
fn validate_copilot_config(config: &CopilotRuntimeConfig) -> RuntimeResult<()> {
let mut profile_ids = std::collections::HashSet::new();
for profile in &config.providers.profiles {
if profile.id.trim().is_empty() || !profile_ids.insert(profile.id.as_str()) {
return Err(RuntimeError::invalid_state(
"managed copilot profile ids must be non-empty and unique",
));
}
if profile.provider.trim().is_empty() {
return Err(RuntimeError::invalid_state(
"managed copilot profile provider is required",
));
}
if profile.models.is_empty() {
return Err(RuntimeError::invalid_state(
"managed copilot profile models must be non-empty",
));
}
let mut models = std::collections::HashSet::new();
for model in &profile.models {
if model.trim().is_empty() || !models.insert(model.as_str()) {
return Err(RuntimeError::invalid_state(
"managed copilot profile models must be non-empty and unique",
));
}
provider_default_capability_upper_bound(&profile.provider, model)
.ok_or_else(|| RuntimeError::invalid_state("managed copilot profile model is unsupported"))?;
}
}
Ok(())
}
#[derive(Debug, Default, Deserialize)]
#[derive(Default, Deserialize)]
struct AppConfigFile {
db: Option<DbConfigFile>,
crypto: Option<CryptoConfigFile>,
copilot: Option<CopilotRuntimeConfig>,
}
#[derive(Default, Deserialize)]
#[serde(rename_all = "camelCase")]
struct CryptoConfigFile {
private_key: Option<String>,
}
#[derive(Default, Deserialize)]
#[serde(rename_all = "camelCase")]
struct DbConfigFile {
datasource_url: Option<String>,
@@ -97,31 +229,67 @@ fn database_url_from_env() -> Option<String> {
env::var("DATABASE_URL").ok().and_then(non_empty_string)
}
fn private_key_from_env() -> Option<String> {
env::var("AFFINE_PRIVATE_KEY").ok().and_then(non_empty_string)
}
fn non_empty_string(value: String) -> Option<String> {
if value.trim().is_empty() { None } else { Some(value) }
}
fn app_config_from_config_files() -> RuntimeResult<AppConfigFile> {
let mut merged = AppConfigFile::default();
deserialize_app_config(app_config_value_from_config_files()?)
}
fn app_config_value_from_config_files() -> RuntimeResult<serde_json::Value> {
let mut merged = serde_json::Value::Object(Map::new());
for path in config_json_paths() {
if !path.exists() {
continue;
}
let raw = fs::read_to_string(&path).map_err(|err| RuntimeError::io("failed to read config file", err))?;
let config: AppConfigFile =
serde_json::from_str(&raw).map_err(|err| RuntimeError::json("failed to parse config file", err))?;
merged.apply_file_config(config);
let value = serde_json::from_str(&raw).map_err(|err| RuntimeError::json("failed to parse config file", err))?;
merge_config_value(&mut merged, expand_module_config_paths(value));
}
Ok(merged)
}
impl AppConfigFile {
fn apply_file_config(&mut self, config: AppConfigFile) {
if config.db.is_some() {
self.db = config.db;
fn expand_module_config_paths(mut value: serde_json::Value) -> serde_json::Value {
if let Some(root) = value.as_object_mut() {
for module in root.values_mut().filter_map(serde_json::Value::as_object_mut) {
let entries = std::mem::take(module);
for (path, value) in entries {
insert_flat_override(module, &path, value);
}
}
}
value
}
#[cfg(test)]
fn app_config_from_module_json(value: serde_json::Value) -> RuntimeResult<AppConfigFile> {
deserialize_app_config(expand_module_config_paths(value))
}
fn deserialize_app_config(value: serde_json::Value) -> RuntimeResult<AppConfigFile> {
serde_json::from_value(value).map_err(|err| RuntimeError::json("failed to parse config file", err))
}
fn merge_config_value(base: &mut serde_json::Value, overrides: serde_json::Value) {
match (base, overrides) {
(serde_json::Value::Object(base), serde_json::Value::Object(overrides)) => {
for (key, value) in overrides {
if let Some(existing) = base.get_mut(&key) {
merge_config_value(existing, value);
} else {
base.insert(key, value);
}
}
}
(base, overrides) => *base = overrides,
}
}
fn default_mail_class_mapping() -> BTreeMap<String, String> {
@@ -161,40 +329,60 @@ fn default_mail_class_mapping() -> BTreeMap<String, String> {
.collect()
}
async fn load_app_config_overrides_from_db(pool: &PgPool) -> RuntimeResult<AppConfigFile> {
async fn load_app_config_overrides_from_db(pool: &PgPool) -> RuntimeResult<serde_json::Value> {
let rows = match sqlx::query("SELECT id, value FROM app_configs").fetch_all(pool).await {
Ok(rows) => rows,
Err(sqlx::Error::Database(err)) if err.code().as_deref() == Some("42P01") => return Ok(AppConfigFile::default()),
Err(sqlx::Error::Database(err)) if err.code().as_deref() == Some("42P01") => {
return Ok(serde_json::Value::Object(Map::new()));
}
Err(err) => return Err(RuntimeError::database("failed to load app config overrides", err)),
};
app_config_from_flat_overrides(rows.into_iter().map(|row| {
Ok(app_config_value_from_flat_overrides(rows.into_iter().map(|row| {
let id: String = row.get("id");
let value: serde_json::Value = row.get("value");
(id, value)
}))
})))
}
#[cfg(test)]
fn app_config_from_flat_overrides<I, S>(rows: I) -> RuntimeResult<AppConfigFile>
where
I: IntoIterator<Item = (S, serde_json::Value)>,
S: AsRef<str>,
{
deserialize_app_config(app_config_value_from_flat_overrides(rows))
}
fn app_config_value_from_flat_overrides<I, S>(rows: I) -> serde_json::Value
where
I: IntoIterator<Item = (S, serde_json::Value)>,
S: AsRef<str>,
{
let mut root = Map::new();
for (path, value) in rows {
let Some((module, key)) = path.as_ref().split_once('.') else {
continue;
};
root
.entry(module.to_string())
.or_insert_with(|| serde_json::Value::Object(Map::new()));
if let Some(serde_json::Value::Object(module_object)) = root.get_mut(module) {
module_object.insert(key.to_string(), value);
}
insert_flat_override(&mut root, path.as_ref(), value);
}
serde_json::from_value(serde_json::Value::Object(root))
.map_err(|err| RuntimeError::json("invalid app config overrides", err))
serde_json::Value::Object(root)
}
fn insert_flat_override(root: &mut Map<String, serde_json::Value>, path: &str, value: serde_json::Value) {
let mut parts = path.split('.').peekable();
let mut current = root;
while let Some(part) = parts.next() {
if parts.peek().is_none() {
current.insert(part.to_string(), value);
return;
}
let entry = current
.entry(part.to_string())
.or_insert_with(|| serde_json::Value::Object(Map::new()));
if !entry.is_object() {
*entry = serde_json::Value::Object(Map::new());
}
current = entry.as_object_mut().expect("override node must be an object");
}
}
pub(super) fn config_json_paths() -> Vec<PathBuf> {
@@ -273,6 +461,79 @@ mod tests {
);
}
#[test]
fn expands_module_config_paths_from_json_files() {
let app_config = app_config_from_module_json(serde_json::json!({
"copilot": {
"enabled": true,
"byok.enabled": false,
"providers.profiles": [{
"id": "managed-openai",
"type": "openai",
"models": ["gpt-5.6-luna"],
"config": { "apiKey": "test" }
}]
}
}))
.unwrap();
let copilot = app_config.copilot.unwrap();
assert!(copilot.enabled);
assert!(!copilot.byok.enabled);
assert_eq!(copilot.providers.profiles.len(), 1);
assert_eq!(copilot.providers.profiles[0].id, "managed-openai");
}
#[test]
fn partial_database_config_preserves_file_config_siblings() {
let mut file_config = expand_module_config_paths(serde_json::json!({
"copilot": {
"enabled": true,
"byok": { "enabled": true, "allowCustomEndpoint": true },
"providers": {
"profiles": [{
"id": "managed-openai",
"type": "openai",
"models": ["gpt-5.6-luna"],
"config": { "apiKey": "test" }
}]
}
}
}));
let database_config = app_config_value_from_flat_overrides([("copilot.byok.enabled", serde_json::json!(false))]);
merge_config_value(&mut file_config, database_config);
let copilot = deserialize_app_config(file_config).unwrap().copilot.unwrap();
assert!(copilot.enabled);
assert!(!copilot.byok.enabled);
assert!(copilot.byok.allow_custom_endpoint);
assert_eq!(copilot.providers.profiles.len(), 1);
assert_eq!(copilot.providers.profiles[0].id, "managed-openai");
}
#[test]
fn database_config_only_replaces_an_active_private_key_explicitly() {
let active = BackendRuntimeConfig {
database_url: "postgresql://active".to_string(),
invite_quota: InviteQuotaConfig::default(),
private_key: Arc::new(Zeroizing::new("active-private-key".to_string())),
copilot: CopilotRuntimeConfig::default(),
};
let empty = serde_json::Value::Object(Map::new());
let unchanged = active.apply_db_overrides(empty.clone(), empty.clone()).unwrap();
assert_eq!(unchanged.private_key.as_str(), "active-private-key");
let overridden = active
.apply_db_overrides(
empty,
app_config_value_from_flat_overrides([("crypto.privateKey", serde_json::json!("database-private-key"))]),
)
.unwrap();
assert_eq!(overridden.private_key.as_str(), "database-private-key");
}
#[test]
fn invite_quota_policy_is_internal_not_app_configurable() {
let app_config = app_config_from_flat_overrides([
packages/backend/native/src/runtime/mod.rs
+1 -1
View File
@@ -6,5 +6,5 @@ pub(crate) mod error;
pub(crate) mod migrations;
pub(crate) mod types;
pub(crate) use config::{BackendRuntimeConfig, InviteQuotaConfig};
pub(crate) use config::{BackendRuntimeConfig, CopilotManagedProfileConfig, CopilotRuntimeConfig, InviteQuotaConfig};
pub(crate) use error::{RuntimeError, RuntimeResult, napi_error, to_napi_error};
packages/backend/native/src/runtime/types.rs
-7
View File
@@ -138,13 +138,6 @@ pub struct RuntimeWorkspaceInviteLinkRecord {
pub expires_at_ms: i64,
}
#[napi_derive::napi(object)]
pub struct RuntimeByokLocalLeaseRecord {
pub lease_id: String,
pub payload: serde_json::Value,
pub expires_at_ms: i64,
}
#[napi_derive::napi(object)]
pub struct RuntimeDocHistoryInput {
pub workspace_id: String,
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 =

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