diff --git a/.docker/selfhost/schema.json b/.docker/selfhost/schema.json index b3b4bc7991..682a12a7f4 100644 --- a/.docker/selfhost/schema.json +++ b/.docker/selfhost/schema.json @@ -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\":\"\"}", diff --git a/Cargo.lock b/Cargo.lock index 66dee1c13f..e663066ed4 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -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" diff --git a/Cargo.toml b/Cargo.toml index 815ea8fd21..28e25f8101 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -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] diff --git a/packages/backend/native/Cargo.toml b/packages/backend/native/Cargo.toml index 191e07034b..faedfec8f6 100644 --- a/packages/backend/native/Cargo.toml +++ b/packages/backend/native/Cargo.toml @@ -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 } diff --git a/packages/backend/native/index.d.ts b/packages/backend/native/index.d.ts index d7c1691672..8919f43e40 100644 --- a/packages/backend/native/index.d.ts +++ b/packages/backend/native/index.d.ts @@ -8,6 +8,9 @@ export declare class BackendRuntime { acquireCoordinationLease(key: string, owner: string, ttlMs: number): Promise releaseCoordinationLease(key: string, owner: string, fencingToken: bigint | number): Promise renewCoordinationLease(key: string, owner: string, fencingToken: bigint | number, ttlMs: number): Promise + executeCopilotStream(input: CopilotExecuteInput, maxSteps: number, callback: ((err: Error | null, arg: string) => void), toolCallback: ((err: Error | null, arg: string) => Promise)): Promise + executeCopilot(input: CopilotExecuteInput): Promise + assertCopilotRoute(input: CopilotRouteCheckInput): Promise /** * Merge pending doc updates with y-octo and persist the merged snapshot. * @@ -51,21 +54,29 @@ export declare class BackendRuntime { getWorkspaceInviteLink(workspaceId: string): Promise getWorkspaceInviteLinkById(inviteId: string): Promise revokeWorkspaceInviteLink(workspaceId: string): Promise - createByokLocalLease(activeKey: string, leaseId: string, payload: any, ttlMs: number): Promise - getByokLocalLease(leaseId: string): Promise cleanupExpiredRuntimeStates(limit: number): Promise refreshWorkspaceAdminStatsDirty(batchLimit: number, owner: string, leaseTtlMs: number): Promise recalibrateWorkspaceAdminStats(lastSid: number, batchLimit: number, owner: string, leaseTtlMs: number): Promise writeWorkspaceAdminStatsDailySnapshot(owner: string, leaseTtlMs: number): Promise recalibrateWorkspaceAdminStatsDaily(batchLimit: number, owner: string, leaseTtlMs: number, lockRetryTimes: number, lockRetryDelayMs: number): Promise - constructor() + constructor(privateKey?: string | undefined | null) start(): Promise stop(): Promise + reloadConfig(privateKey?: string | undefined | null): Promise health(): Promise runMigrations(): Promise + listByokProfiles(workspaceId: string): Promise> + createByokProfile(input: CreateByokProfileInput): Promise + replaceByokProfile(input: ReplaceByokProfileInput): Promise + rotateByokCredential(input: RotateByokCredentialInput): Promise + probeByokProfile(input: ProbeByokProfileInput): Promise + probeByokDraft(input: ProbeByokDraftInput): Promise + deleteByokProfile(workspaceId: string, profileId: string): Promise + reorderByokProfiles(input: ReorderByokProfilesInput): Promise> + createByokLocalLease(input: CreateByokLocalLeaseInput): Promise } -export declare class LlmStreamHandle { +export declare class CopilotStreamHandle { abort(): void } @@ -107,46 +118,6 @@ export declare class Tokenizer { count(content: string, allowedSpecial?: Array | 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 - errorCode?: string -} - export declare function activateLicense(request: LicenseKeyRequest): Promise /** @@ -194,6 +165,15 @@ export interface BackendRuntimeHealth { export declare function buildPublicRootDoc(rootDocBin: Buffer, docMetas: Array): Buffer +export interface BuiltInManagedTarget { + id: string + displayName: string + minimumTier: BuiltInManagedTargetTier +} + +export type BuiltInManagedTargetTier = 'Standard'| +'Premium'; + export interface BuiltInPromptRenderContract { name: string renderParams: Record @@ -203,20 +183,124 @@ export interface BuiltInPromptSessionContract { name: string turns: Array renderParams: Record - maxTokenSize: number } export interface BuiltInPromptSpec { name: string action?: string - model: string - optionalModels?: Array config?: any params?: Record builtins?: Array messages: Array } +export interface BuiltInRouteOptions { + routeId: string + standardDefaultTargetId?: string + premiumDefaultTargetId?: string + choices: Array +} + +export interface ByokCapabilityInput { + input: Array + output: Array + features: Array + attachmentKinds: Array + attachmentSources: Array +} + +export interface ByokCatalogModelOutput { + modelId: string + displayName: string + recommended: boolean + capabilities: Array +} + +export interface ByokCatalogOutput { + version: string + providers: Array +} + +export interface ByokCatalogProviderOutput { + provider: string + models: Array +} + +export interface ByokEndpointInput { + kind: string + url?: string +} + +export interface ByokLocalLeaseOutput { + leaseId: string + expiresAtMs: number +} + +export interface ByokModelDeclarationInput { + modelId: string + enabled: boolean + capabilities: Array +} + +export interface ByokModelProbeCheckOutput { + operation: string + status: ByokProbeStatusOutput +} + +export interface ByokModelProbeOutput { + modelId: string + checks: Array +} + +export interface ByokProbeCheckInput { + modelId: string + operation: string +} + +export interface ByokProbeResultOutput { + definitionFingerprint: string + stale: boolean + connection: ByokProbeStatusOutput + models: Array +} + +export interface ByokProbeStatusOutput { + kind: string + testedAtMs?: number + errorKind?: string +} + +export interface ByokProfileDefinitionInput { + version: number + endpoint: ByokEndpointInput + models: Array +} + +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 +} + export interface CanonicalChatRequestContract { model: string messages: Array @@ -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 +} + +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 } -export declare function llmCountPromptTokens(request: PromptTokenCountContract): PromptTokenCountResult - -export declare function llmDispatchPrepared(routesJson: string): Promise - -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)): LlmStreamHandle - -export declare function llmDispatchToolLoopStreamPrepared(routesJson: string, maxSteps: number, callback: ((err: Error | null, arg: string) => void), toolCallback: ((err: Error | null, arg: string) => Promise)): 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)): LlmStreamHandle - -export declare function llmEmbeddingDispatch(protocol: string, backendConfigJson: string, requestJson: string): Promise - -export declare function llmEmbeddingDispatchPrepared(routesJson: string): Promise - export interface LlmEmbeddingRequestContract { model: string inputs: Array @@ -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 +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 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 @@ -537,10 +660,6 @@ export interface LlmRequestContract { middleware?: any } -export declare function llmRerankDispatch(protocol: string, backendConfigJson: string, requestJson: string): Promise - -export declare function llmRerankDispatchPrepared(routesJson: string): Promise - 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 - -export declare function llmStructuredDispatchPrepared(routesJson: string): Promise - export interface LlmStructuredRequestContract { model: string messages: Array @@ -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 +} + +export interface ProbeByokProfileInput { + workspaceId: string + profileId: string + checks: Array +} + export declare function processImage(input: Buffer, maxEdge: number, keepExif: boolean): Promise 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 -} - -export interface PromptMetadataResult { - paramKeys: Array - templateParams: Record -} - export interface PromptParamSpec { default?: string enumValues?: Array } -export interface PromptRenderContract { - messages: Array - templateParams: Record - renderParams: Record -} - export interface PromptRenderResult { messages: Array warnings: Array } -export interface PromptSessionContract { - prompt: PromptSessionPrompt - turns: Array - renderParams: Record - maxTokenSize: number -} - -export interface PromptSessionPrompt { - action?: string - model?: string - promptTokens: number - templateParams: Record - messages: Array -} - export interface PromptSessionResult { messages: Array warnings: Array @@ -775,15 +868,6 @@ export interface PromptStructuredResponseContract { strict?: boolean } -export interface PromptTokenCountContract { - model?: string - messages: Array -} - -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 - optionalModels: Array - requestedModelId?: string - defaultModel?: string +export interface ReorderByokProfilesInput { + workspaceId: string + profiles: Array + 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 } -export interface RuntimeByokLocalLeaseRecord { - leaseId: string - payload: any - expiresAtMs: number -} - export interface RuntimeDocBlobRefsResult { scannedDocs: number parsedDocs: number diff --git a/packages/backend/native/src/llm/action/catalog.rs b/packages/backend/native/src/llm/action/catalog.rs index 87c2a120d4..348a9cb7da 100644 --- a/packages/backend/native/src/llm/action/catalog.rs +++ b/packages/backend/native/src/llm/action/catalog.rs @@ -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) -> 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 { - 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 { - 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> { - 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) -> Result { + 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) -> 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\"")); + } } } diff --git a/packages/backend/native/src/llm/action/contract.rs b/packages/backend/native/src/llm/action/contract.rs index 5a9b0b77c6..86065b6e8f 100644 --- a/packages/backend/native/src/llm/action/contract.rs +++ b/packages/backend/native/src/llm/action/contract.rs @@ -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, -} - -#[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, - #[serde(default)] - pub state_patch: Option, -} - -#[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, - #[serde(skip_serializing_if = "Option::is_none")] - pub status: Option, - #[serde(skip_serializing_if = "Option::is_none")] - pub attachment: Option, - #[serde(skip_serializing_if = "Option::is_none")] - pub result: Option, - #[serde(skip_serializing_if = "Option::is_none")] - pub error_code: Option, - #[serde(skip_serializing_if = "Option::is_none")] - pub error_message: Option, - #[serde(skip_serializing_if = "Option::is_none")] - pub trace: Option, -} - -#[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, - #[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, - pub state: Value, - pub steps: Vec, - pub trace: ActionTrace, - pub events: Vec, -} - -#[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, - #[serde(skip_serializing_if = "Option::is_none")] - pub state_patch: Option, - #[serde(skip_serializing_if = "Option::is_none")] - pub error: Option, -} - -#[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, - #[serde(skip_serializing_if = "Option::is_none")] - pub error_code: Option, -} - #[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>, #[serde(skip_serializing_if = "Option::is_none")] pub slice_manifest: Option>, - #[serde(skip_serializing_if = "Option::is_none")] - pub prepared_routes: Option, } #[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, - #[schemars(required)] - pub provider_meta: Option, } #[derive(Clone, Debug, Deserialize, JsonSchema, PartialEq, Serialize)] @@ -253,8 +104,5 @@ pub struct TranscriptResult { pub normalized_transcript: String, #[schemars(required)] pub summary_json: Option, - #[schemars(required)] - pub provider_meta: Option, pub version: String, - pub strategy: String, } diff --git a/packages/backend/native/src/llm/action/mod.rs b/packages/backend/native/src/llm/action/mod.rs index 10dc13fba1..7186e05b9e 100644 --- a/packages/backend/native/src/llm/action/mod.rs +++ b/packages/backend/native/src/llm/action/mod.rs @@ -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, -) -> Result { - 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::(); - 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; diff --git a/packages/backend/native/src/llm/action/runtime.rs b/packages/backend/native/src/llm/action/runtime.rs deleted file mode 100644 index 195bc05964..0000000000 --- a/packages/backend/native/src/llm/action/runtime.rs +++ /dev/null @@ -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>, - pub event_sender: Option>, - #[cfg(test)] - pub abort_after_events: Option, - #[cfg(test)] - pub mock_output: Option, -} - -#[derive(Clone, Debug)] -pub struct ActionRuntimeState { - pub status: ActionRunStatus, - pub result: Value, - pub action_state: Value, - pub steps: Vec, - pub events: Vec, - pub trace: ActionTrace, - pub error_code: Option, -} - -fn invalid_input(message: impl Into) -> Error { - Error::new(Status::InvalidArg, message.into()) -} - -pub fn run_action_recipe_prepared_with_control( - input: ActionRuntimeInput, - control: ActionRuntimeControl, -) -> Result { - 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 { - 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 { - validate_value("input", &recipe.input_schema, &input.input)?; - run_recipe(recipe, input, control) -} - -fn run_recipe( - recipe: ActionRecipe, - input: ActionRuntimeInput, - control: ActionRuntimeControl, -) -> Result { - 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 { - 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::>(); - 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, - lightweight: Vec, - step_patches: std::collections::HashMap>, - ) -> 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>>, -) -> Vec { - 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 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>>, -} - -impl<'a> AffineActionStepExecutor<'a> { - fn new(_control: &'a ActionRuntimeControl, attachments: Arc>>) -> 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, - ) -> std::result::Result { - 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, - ) -> std::result::Result { - 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::>(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, state: &Value) -> std::result::Result { - 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, - state: &Value, - ) -> std::result::Result { - 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) -> Option { - 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)) -} diff --git a/packages/backend/native/src/llm/action/slides_outline.rs b/packages/backend/native/src/llm/action/slides_outline.rs deleted file mode 100644 index f15fedca06..0000000000 --- a/packages/backend/native/src/llm/action/slides_outline.rs +++ /dev/null @@ -1,240 +0,0 @@ -use serde_json::{Map, Value}; - -pub(super) fn project_slides_outline_markdown(value: &Value) -> Result { - 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::(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 { - 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::>() - .join("\n"), - ) - } else { - Some(format!(" - {content}")) - } - } - _ => None, - } -} - -fn render_slide_item(item: &Value) -> Result { - 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 { - 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 { - 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::, _>>()? - .into_iter() - .flatten() - .collect::>() - } else { - render_slide_object(content)? - }; - - Ok( - std::iter::once(format!("- {title}")) - .chain(rendered_sections) - .collect::>() - .join("\n"), - ) -} - -fn parse_labeled_segments(text: &str) -> std::collections::HashMap { - 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, 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) -> Result, 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, keys: &[&str]) -> Option { - keys - .iter() - .find_map(|key| object.get(*key).and_then(value_to_optional_string)) -} - -fn required_string_prop(object: &Map, keys: &[&str], name: &str) -> Result { - string_prop(object, keys) - .filter(|value| !value.is_empty()) - .ok_or_else(|| format!("slidesOutlineMarkdown requires {name}")) -} - -fn value_to_optional_string(value: &Value) -> Option { - 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::>() - .join(", "); - Some(joined) - } - _ => None, - } -} diff --git a/packages/backend/native/src/llm/action/tests.rs b/packages/backend/native/src/llm/action/tests.rs deleted file mode 100644 index f221cc620a..0000000000 --- a/packages/backend/native/src/llm/action/tests.rs +++ /dev/null @@ -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![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![ - 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) -> 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) - ); -} diff --git a/packages/backend/native/src/llm/assets/prompts/built-in.json b/packages/backend/native/src/llm/assets/prompts/built-in.json index cb52ec8ec0..9fd3425da8 100644 --- a/packages/backend/native/src/llm/assets/prompts/built-in.json +++ b/packages/backend/native/src/llm/assets/prompts/built-in.json @@ -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": [] } ] diff --git a/packages/backend/native/src/llm/byok/catalog.rs b/packages/backend/native/src/llm/byok/catalog.rs new file mode 100644 index 0000000000..2832728720 --- /dev/null +++ b/packages/backend/native/src/llm/byok/catalog.rs @@ -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, +} + +#[derive(Clone, Serialize)] +#[serde(rename_all = "camelCase")] +#[napi_derive::napi(object)] +pub struct ByokCatalogProviderOutput { + pub provider: String, + pub models: Vec, +} + +#[derive(Clone, Serialize)] +#[serde(rename_all = "camelCase")] +#[napi_derive::napi(object)] +pub struct ByokCatalogOutput { + pub version: String, + pub providers: Vec, +} + +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::>>(); + + 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::>(); + 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()); + } + } + } +} diff --git a/packages/backend/native/src/llm/byok/contract.rs b/packages/backend/native/src/llm/byok/contract.rs new file mode 100644 index 0000000000..8827363294 --- /dev/null +++ b/packages/backend/native/src/llm/byok/contract.rs @@ -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, + pub output: Vec, + pub features: Vec, + pub attachment_kinds: Vec, + pub attachment_sources: Vec, +} + +#[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, +} + +#[derive(Clone, Deserialize, Serialize)] +#[serde(rename_all = "camelCase", deny_unknown_fields)] +#[napi_derive::napi(object)] +pub struct ByokEndpointInput { + pub kind: String, + pub url: Option, +} + +#[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, +} + +#[derive(Clone)] +#[napi_derive::napi(object)] +pub struct CreateByokProfileInput { + pub workspace_id: String, + pub provider: String, + pub name: String, + pub description: Option, + 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, + pub definition: ByokProfileDefinitionInput, + pub credential: Option, + 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, + 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, +} + +#[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, + pub profile_id: Option, + pub expected_revision: Option, + pub definition: ByokProfileDefinitionInput, + pub checks: Vec, +} + +#[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, + 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, +} + +#[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, + pub definition: ByokProfileDefinitionInput, + pub enabled: bool, + pub sort_order: i32, + pub revision: i32, + pub validation: Option, +} + +#[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, + pub error_kind: Option, +} + +#[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, +} + +#[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, +} + +#[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, +} + +#[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, +} + +#[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, +} + +#[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 { + 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::, _>>()?; + 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 { + 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(values: Vec, parse: impl Fn(&str) -> Option) -> Result, 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 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) -> 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()); + } + } +} diff --git a/packages/backend/native/src/llm/byok/envelope.rs b/packages/backend/native/src/llm/byok/envelope.rs new file mode 100644 index 0000000000..585abc5d81 --- /dev/null +++ b/packages/backend/native/src/llm/byok/envelope.rs @@ -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>); + +impl SensitiveCredential { + pub(crate) fn new(value: impl Into>) -> 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 { + if root_secret.is_empty() { + return Err(CredentialEnvelopeError::Unavailable); + } + let mut key = Zeroizing::new([0_u8; 32]); + Hkdf::::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 { + 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 { + 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 { + ["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 { + [ + "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()); + } +} diff --git a/packages/backend/native/src/llm/byok/mod.rs b/packages/backend/native/src/llm/byok/mod.rs new file mode 100644 index 0000000000..f54dfd77c0 --- /dev/null +++ b/packages/backend/native/src/llm/byok/mod.rs @@ -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}; diff --git a/packages/backend/native/src/llm/byok/validation.rs b/packages/backend/native/src/llm/byok/validation.rs new file mode 100644 index 0000000000..735aeb8287 --- /dev/null +++ b/packages/backend/native/src/llm/byok/validation.rs @@ -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, + old_definition: &ByokProfileDefinition, + definition: &ByokProfileDefinition, + credential_generation: i32, + credential_changed: bool, +) -> Option { + 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)); + } +} diff --git a/packages/backend/native/src/llm/contract_schema.rs b/packages/backend/native/src/llm/contract_schema.rs index e9369deaa6..b4d00cf099 100644 --- a/packages/backend/native/src/llm/contract_schema.rs +++ b/packages/backend/native/src/llm/contract_schema.rs @@ -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::(); - 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::(); - 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 { match name { - // runtime-owned temporary native facade - "executionPlan" => Some(generated_schema_for::()), - // adapter-owned temporary native facade - "preparedRoutes" => Some(generated_schema_for::< - Vec, - >()), // AFFiNE-native-owned N-API projection over adapter model registry/matcher "capabilityMatchRequest" => Some(generated_schema_for::()), "capabilityMatchResponse" => Some(generated_schema_for::()), @@ -133,11 +125,6 @@ fn schema_by_name(name: &str) -> Option { "modelRegistryResolveRequest" => Some(generated_schema_for::()), "modelRegistryResolveResponse" => Some(generated_schema_for::()), "providerDriverSpec" => Some(generated_schema_for::()), - // AFFiNE-native-owned prompt facade over adapter prompt DTOs/catalog - "promptRenderContract" => Some(generated_schema_for::()), - "promptSessionContract" => Some(generated_schema_for::()), - "requestedModelMatchRequest" => Some(generated_schema_for::()), - "requestedModelMatchResponse" => Some(generated_schema_for::()), // runtime-owned "toolCallbackRequest" => Some(generated_schema_for::()), "toolCallbackResponse" => Some(generated_schema_for::()), @@ -176,23 +163,6 @@ pub fn llm_validate_contract(name: String, value: Value) -> Result { ))) } -#[napi(catch_unwind)] -pub fn llm_compile_execution_plan(value: Value) -> Result { - 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 { - 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")); - } } diff --git a/packages/backend/native/src/llm/core/capability.rs b/packages/backend/native/src/llm/core/capability.rs index 42da4dcb65..4a6d9fd4d2 100644 --- a/packages/backend/native/src/llm/core/capability.rs +++ b/packages/backend/native/src/llm/core/capability.rs @@ -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 { @@ -14,25 +12,7 @@ pub fn llm_match_model_capabilities(payload: CapabilityMatchRequest) -> Result Result { - 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)?, }) } diff --git a/packages/backend/native/src/llm/core/contracts/mod.rs b/packages/backend/native/src/llm/core/contracts/mod.rs index f17b91fab2..d8ecc2a999 100644 --- a/packages/backend/native/src/llm/core/contracts/mod.rs +++ b/packages/backend/native/src/llm/core/contracts/mod.rs @@ -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, - #[napi(ts_type = "Record")] - pub template_params: Value, - #[napi(ts_type = "Record")] - 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, - pub messages: Vec, -} - -#[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, -} - -#[napi(object)] #[derive(Debug, Clone, Deserialize, Serialize, PartialEq)] #[serde(rename_all = "camelCase")] pub struct PromptMetadataResult { pub param_keys: Vec, - #[napi(ts_type = "Record")] 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, - #[napi(ts_type = "Record")] - 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, - #[serde(skip_serializing_if = "Option::is_none")] - pub model: Option, - pub prompt_tokens: u32, - #[napi(ts_type = "Record")] - pub template_params: Value, - pub messages: Vec, -} - #[napi(object)] #[derive(Debug, Clone, Deserialize, Serialize, PartialEq)] #[serde(rename_all = "camelCase")] @@ -112,7 +48,6 @@ pub struct BuiltInPromptSessionContract { pub turns: Vec, #[napi(ts_type = "Record")] pub render_params: Value, - pub max_token_size: u32, } #[napi(object)] @@ -289,29 +224,6 @@ pub struct CapabilityMatchResponse { pub model_id: Option, } -#[napi(object)] -#[derive(Debug, Clone, Deserialize, JsonSchema, Serialize, PartialEq)] -#[serde(rename_all = "camelCase")] -#[serde(deny_unknown_fields)] -pub struct RequestedModelMatchRequest { - pub provider_ids: Vec, - pub optional_models: Vec, - #[serde(skip_serializing_if = "Option::is_none")] - pub requested_model_id: Option, - #[serde(skip_serializing_if = "Option::is_none")] - pub default_model: Option, -} - -#[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, - 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, #[serde(skip_serializing_if = "Option::is_none")] pub options: Option, @@ -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() { diff --git a/packages/backend/native/src/llm/core/model_registry.rs b/packages/backend/native/src/llm/core/model_registry.rs index 60032fa74d..260bb08802 100644 --- a/packages/backend/native/src/llm/core/model_registry.rs +++ b/packages/backend/native/src/llm/core/model_registry.rs @@ -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 Result, messages: &[PromptMessageContract]) -> u32 { - let content = messages - .iter() - .map(|message| message.content.as_str()) - .collect::(); - prompt_tokenizer(model) - .map(|tokenizer| tokenizer.count(content, None)) - .unwrap_or(0) -} - -fn prompt_tokenizer(model: Option<&str>) -> Option { - 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 { - 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 { - let content = request - .messages - .iter() - .map(|message| message.content.as_str()) - .collect::(); - 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 { 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 { - 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 { - 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 { 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> { pub fn llm_get_built_in_prompt_spec(name: String) -> Result> { 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::(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::(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::(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::(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::(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::(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::(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"] - } - }), - ); - } -} diff --git a/packages/backend/native/src/llm/core/prompt/render.rs b/packages/backend/native/src/llm/core/prompt/render.rs index 6efb918ca5..0338c4a25f 100644 --- a/packages/backend/native/src/llm/core/prompt/render.rs +++ b/packages/backend/native/src/llm/core/prompt/render.rs @@ -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, ¶ms).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(), ¶ms).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" + })]) + ); + } +} diff --git a/packages/backend/native/src/llm/core/prompt/session.rs b/packages/backend/native/src/llm/core/prompt/session.rs index ea4e2b8669..1f5d0697d3 100644 --- a/packages/backend/native/src/llm/core/prompt/session.rs +++ b/packages/backend/native/src/llm/core/prompt/session.rs @@ -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, params: &Map, ) -> std::result::Result { - 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, + params: &Map, + history_input_bytes: usize, +) -> std::result::Result { + 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::>(); 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 { - 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, 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 { 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::>(), + ["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::>(), + ["action", "latest"] + ); + } +} diff --git a/packages/backend/native/src/llm/core/request_builder/mod.rs b/packages/backend/native/src/llm/core/request_builder/mod.rs index 3020939c12..b2baa0156e 100644 --- a/packages/backend/native/src/llm/core/request_builder/mod.rs +++ b/packages/backend/native/src/llm/core/request_builder/mod.rs @@ -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 } pub(crate) fn build_image_request_from_messages(request: LlmImageRequestBuildContract) -> Result { - 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) -> Result { 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" })) ); } diff --git a/packages/backend/native/src/llm/core/request_builder/types.rs b/packages/backend/native/src/llm/core/request_builder/types.rs index 42f7eec64c..8360c9d2a9 100644 --- a/packages/backend/native/src/llm/core/request_builder/types.rs +++ b/packages/backend/native/src/llm/core/request_builder/types.rs @@ -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, }; diff --git a/packages/backend/native/src/llm/ffi/mod.rs b/packages/backend/native/src/llm/ffi/mod.rs index 727f9aa5cc..236cfe4549 100644 --- a/packages/backend/native/src/llm/ffi/mod.rs +++ b/packages/backend/native/src/llm/ffi/mod.rs @@ -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, }; diff --git a/packages/backend/native/src/llm/ffi/payload.rs b/packages/backend/native/src/llm/ffi/payload.rs index a6b9530e88..0060834c3d 100644 --- a/packages/backend/native/src/llm/ffi/payload.rs +++ b/packages/backend/native/src/llm/ffi/payload.rs @@ -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 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 for LlmRerankDispatchPayload { } } } - -pub(crate) type LlmPreparedImageDispatchRoutePayload = SerializablePreparedRoute; - -#[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::>>(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::>>(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::>>(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::>(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"); - } -} diff --git a/packages/backend/native/src/llm/mod.rs b/packages/backend/native/src/llm/mod.rs index d00d7f32ce..731332d6b6 100644 --- a/packages/backend/native/src/llm/mod.rs +++ b/packages/backend/native/src/llm/mod.rs @@ -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) -> 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()) +} diff --git a/packages/backend/native/src/llm/prompt_catalog.rs b/packages/backend/native/src/llm/prompt_catalog.rs index 669ad152e4..8d2de16999 100644 --- a/packages/backend/native/src/llm/prompt_catalog.rs +++ b/packages/backend/native/src/llm/prompt_catalog.rs @@ -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, - pub model: String, - #[serde(skip_serializing_if = "Option::is_none")] - pub optional_models: Option>, #[serde(skip_serializing_if = "Option::is_none")] pub config: Option, #[serde(skip_serializing_if = "Option::is_none")] @@ -64,6 +61,69 @@ pub struct BuiltInPromptSpec { pub messages: Vec, } +#[derive(Clone, Deserialize)] +#[serde(rename_all = "camelCase", deny_unknown_fields)] +struct PromptCatalogSpec { + name: String, + #[serde(default)] + action: Option, + #[serde(default)] + managed_route: Option, + #[serde(default)] + config: Option, + #[serde(default)] + params: Option>, + #[serde(default)] + builtins: Option>, + messages: Vec, +} + +#[derive(Clone, Deserialize)] +#[serde(rename_all = "camelCase", deny_unknown_fields)] +struct BuiltInManagedRouteSpec { + targets: Vec, + #[serde(default)] + premium_targets: Option>, + #[serde(default)] + selectable_targets: Vec, +} + +#[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, + pub premium_default_target_id: Option, + pub choices: Vec, +} + #[derive(Debug, Clone, Serialize)] #[serde(rename_all = "camelCase")] pub(crate) struct BuiltInPromptMessage { @@ -73,20 +133,29 @@ pub(crate) struct BuiltInPromptMessage { pub(crate) params: Option>, } -#[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, - pub(crate) model: String, - #[serde(skip_serializing_if = "Option::is_none")] - pub(crate) optional_models: Option>, + pub(crate) managed_targets: Vec, + pub(crate) managed_premium_targets: Option>, + #[serde(skip)] + pub(crate) managed_selectable_targets: Vec, #[serde(skip_serializing_if = "Option::is_none")] pub(crate) config: Option, pub(crate) messages: Vec, } +#[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, prompts: Vec, @@ -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 { + 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 { let partials: BTreeMap = serde_json::from_str(PROMPT_PARTIALS_SOURCE).map_err(|error| format!("invalid prompt partials JSON: {error}"))?; - let specs: Vec = + let catalog_specs: Vec = 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::, _>>()?; + 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::>(); Ok(Self { specs_by_name: specs @@ -140,7 +282,17 @@ impl PromptCatalog { } } -fn compile_prompt_spec(spec: &BuiltInPromptSpec, partials: &BTreeMap) -> Result { +fn compile_prompt_spec(spec: &PromptCatalogSpec, partials: &BTreeMap) -> Result { + 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 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::>(); + let models = route + .selectable_targets + .iter() + .map(|target| target.model_id.as_str()) + .collect::>(); + 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) 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::>()), - 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::>()), + 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::>()), + Some(vec!["gemini-3.6-flash"]) ); } } diff --git a/packages/backend/native/src/llm/route/catalog.rs b/packages/backend/native/src/llm/route/catalog.rs new file mode 100644 index 0000000000..79dd7b75b9 --- /dev/null +++ b/packages/backend/native/src/llm/route/catalog.rs @@ -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> { + 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 { + 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 { + 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, + attachment_sources: Vec, +) -> 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 + )); + } +} diff --git a/packages/backend/native/src/llm/route/contract.rs b/packages/backend/native/src/llm/route/contract.rs new file mode 100644 index 0000000000..27ce0c4438 --- /dev/null +++ b/packages/backend/native/src/llm/route/contract.rs @@ -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, + pub workspace_id: Option, + pub user_id: Option, + pub local_lease_id: Option, + pub access: CopilotAccessProjection, + pub managed_target_id: Option, + pub target_override: Option, +} + +#[derive(Clone)] +#[napi_derive::napi(object)] +pub struct CopilotExecuteInput { + pub slot: String, + pub built_in_route_id: Option, + pub workspace_id: Option, + pub user_id: Option, + pub local_lease_id: Option, + pub access: CopilotAccessProjection, + pub managed_target_id: Option, + pub target_override: Option, + #[napi(ts_type = "unknown")] + pub request: serde_json::Value, +} diff --git a/packages/backend/native/src/llm/route/mod.rs b/packages/backend/native/src/llm/route/mod.rs new file mode 100644 index 0000000000..b72b92c0ad --- /dev/null +++ b/packages/backend/native/src/llm/route/mod.rs @@ -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, +}; diff --git a/packages/backend/native/src/llm/route/policy.rs b/packages/backend/native/src/llm/route/policy.rs new file mode 100644 index 0000000000..d9699bc104 --- /dev/null +++ b/packages/backend/native/src/llm/route/policy.rs @@ -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 }, + 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), + 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 { + let mut profiles = input + .profiles + .iter() + .enumerate() + .filter(|(_, profile)| (profile.source == ProfileSource::Managed) == managed) + .collect::>(); + 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) + )); + } +} diff --git a/packages/backend/native/src/runtime/backend_runtime/byok/admission.rs b/packages/backend/native/src/runtime/backend_runtime/byok/admission.rs new file mode 100644 index 0000000000..058ce167a3 --- /dev/null +++ b/packages/backend/native/src/runtime/backend_runtime/byok/admission.rs @@ -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))) + } + } +} diff --git a/packages/backend/native/src/runtime/backend_runtime/byok/local.rs b/packages/backend/native/src/runtime/backend_runtime/byok/local.rs new file mode 100644 index 0000000000..876c95a9e3 --- /dev/null +++ b/packages/backend/native/src/runtime/backend_runtime/byok/local.rs @@ -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, +} + +#[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, + 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 { + 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::::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, + }) +} diff --git a/packages/backend/native/src/runtime/backend_runtime/byok/mod.rs b/packages/backend/native/src/runtime/backend_runtime/byok/mod.rs new file mode 100644 index 0000000000..b4f0c034ca --- /dev/null +++ b/packages/backend/native/src/runtime/backend_runtime/byok/mod.rs @@ -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}; diff --git a/packages/backend/native/src/runtime/backend_runtime/byok/probe.rs b/packages/backend/native/src/runtime/backend_runtime/byok/probe.rs new file mode 100644 index 0000000000..a6c8adc655 --- /dev/null +++ b/packages/backend/native/src/runtime/backend_runtime/byok/probe.rs @@ -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, +) -> RuntimeResult { + 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::>(); + 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 { + 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) { + 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::(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"}"# + )); + } +} diff --git a/packages/backend/native/src/runtime/backend_runtime/byok/profile.rs b/packages/backend/native/src/runtime/backend_runtime/byok/profile.rs new file mode 100644 index 0000000000..812c723020 --- /dev/null +++ b/packages/backend/native/src/runtime/backend_runtime/byok/profile.rs @@ -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, + encrypted_api_key: String, + definition: serde_json::Value, + sort_order: i32, + enabled: bool, + revision: i32, + credential_generation: i32, + validation: Option, +} + +#[derive(FromRow)] +struct ProfileAdmissionRow { + provider: String, + revision: i32, +} + +pub(in super::super) async fn list(pool: &PgPool, workspace_id: &str) -> RuntimeResult> { + 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 { + 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 { + 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 { + 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 { + 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 { + 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> { + 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::>(); + 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::>(); + if current_set != unique { + return Err(RuntimeError::invalid_input( + "BYOK profile order must contain every server profile", + )); + } + let revisions = current.into_iter().collect::>(); + 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 { + 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 { + 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 { + 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::derive(root_secret) + .map_err(|_| RuntimeError::invalid_state("stable crypto.privateKey is required for persistent BYOK")) +} + +fn parse_definition(value: serde_json::Value) -> RuntimeResult { + serde_json::from_value(value).map_err(|error| RuntimeError::json("invalid stored BYOK definition", error)) +} + +fn profile_output(row: ProfileRow) -> RuntimeResult { + 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) -> RuntimeResult> { + 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(()) + } +} diff --git a/packages/backend/native/src/runtime/backend_runtime/constants.rs b/packages/backend/native/src/runtime/backend_runtime/constants.rs index 29c3f474e1..3f2215a446 100644 --- a/packages/backend/native/src/runtime/backend_runtime/constants.rs +++ b/packages/backend/native/src/runtime/backend_runtime/constants.rs @@ -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"; diff --git a/packages/backend/native/src/runtime/backend_runtime/copilot/context.rs b/packages/backend/native/src/runtime/backend_runtime/copilot/context.rs new file mode 100644 index 0000000000..a4d6e3fad8 --- /dev/null +++ b/packages/backend/native/src/runtime/backend_runtime/copilot/context.rs @@ -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> { + 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> { + 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> { + 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::("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> { + 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::>(); + 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 { + 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}"))) +} diff --git a/packages/backend/native/src/runtime/backend_runtime/copilot/dispatch.rs b/packages/backend/native/src/runtime/backend_runtime/copilot/dispatch.rs new file mode 100644 index 0000000000..d59ea88aaa --- /dev/null +++ b/packages/backend/native/src/runtime/backend_runtime/copilot/dispatch.rs @@ -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, + result: serde_json::Value, +} + +pub(super) struct CompiledExecution { + pub(super) plan: CompiledPlan, + identities: HashMap, +} + +impl CompiledExecution { + pub(super) fn project(&self, event: RuntimeRouteEvent) -> RuntimeResult { + 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, + slot: CatalogSlot, + request: ExecutableRequest, + profiles: Vec, + candidates: Vec, + managed_credentials: HashMap>, +) -> RuntimeResult { + 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::>>()?; + 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>, +) -> RuntimeResult { + 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(value: serde_json::Value) -> RuntimeResult { + serde_json::from_value(value).map_err(|error| RuntimeError::json("invalid copilot execution request", error)) +} + +fn request_requirements(request: &ExecutableRequest) -> (bool, Vec, Vec) { + 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, + sources: &mut HashSet, +) { + 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>, +) -> RuntimeResult { + 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>, +) -> RuntimeResult { + 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> { + 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) + .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 { + 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) -> RuntimeResult { + 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, route_id: &str) -> RuntimeResult { + identities + .get(route_id) + .cloned() + .ok_or_else(|| RuntimeError::invalid_state("runtime emitted an unknown route id")) +} + +fn response_value(response: ExecutableResponse) -> RuntimeResult { + 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::to_value(value).map_err(|error| RuntimeError::json("serialize copilot response failed", error)) +} + +impl From for RuntimeError { + fn from(error: BackendError) -> Self { + RuntimeError::invalid_state(error.to_string()) + } +} diff --git a/packages/backend/native/src/runtime/backend_runtime/copilot/mod.rs b/packages/backend/native/src/runtime/backend_runtime/copilot/mod.rs new file mode 100644 index 0000000000..6ca7528dff --- /dev/null +++ b/packages/backend/native/src/runtime/backend_runtime/copilot/mod.rs @@ -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>>>>; +pub(super) const COPILOT_REQUEST_TIMEOUT: Duration = Duration::from_secs(30 * 60); + +struct AuthorizedCopilotRoute { + config: std::sync::Arc, + slot: route::CatalogSlot, + profiles: Vec, + candidates: Vec, +} + +#[napi_derive::napi] +impl BackendRuntime { + #[napi] + pub async fn execute_copilot(&self, input: CopilotExecuteInput) -> napi::Result { + 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 { + 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 { + 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>> { + 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> { + 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 { + 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", + } +} diff --git a/packages/backend/native/src/runtime/backend_runtime/copilot/stream.rs b/packages/backend/native/src/runtime/backend_runtime/copilot/stream.rs new file mode 100644 index 0000000000..fca4dc6238 --- /dev/null +++ b/packages/backend/native/src/runtime/backend_runtime/copilot/stream.rs @@ -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, + CatalogSlot, + ExecutableRequest, + Vec, + Vec, + HashMap>, +); + +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, +} + +#[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, + tool_callback: ThreadsafeFunction>, + ) -> Result { + 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, + max_steps: usize, + callback: &ThreadsafeFunction, + tool_callback: &ThreadsafeFunction>, + 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, + 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>, + call: &AccumulatedToolCall, + aborted: &AtomicBool, + stream_deadline: Instant, +) -> std::result::Result { + 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::(&ctx.value).map_err(|error| error.to_string()), + ); + Ok(()) + }) { + Ok(promise) => { + if let Err(error) = promise.catch(move |ctx: CallbackContext| { + 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>>>>; + +fn send_tool_result(sender: &ToolResultSender, result: std::result::Result) { + if let Some(sender) = sender.lock().expect("tool callback sender poisoned").take() { + let _ = sender.send(result); + } +} diff --git a/packages/backend/native/src/runtime/backend_runtime/mod.rs b/packages/backend/native/src/runtime/backend_runtime/mod.rs index 75bd1f1fb4..0bb08469d9 100644 --- a/packages/backend/native/src/runtime/backend_runtime/mod.rs +++ b/packages/backend/native/src/runtime/backend_runtime/mod.rs @@ -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, + config: RwLock>, pool: Mutex>, + managed_token_providers: copilot::ManagedTokenProviderCache, } #[napi_derive::napi] impl BackendRuntime { #[napi(constructor)] - pub fn new() -> Result { + pub fn new(private_key: Option) -> Result { + 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) -> 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 { 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> { + byok::list(&self.pool().await?, &workspace_id) + .await + .map_err(to_napi_error) + } + + #[napi] + pub async fn create_byok_profile(&self, input: CreateByokProfileInput) -> Result { + 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 { + 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 { + 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 { + 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 { + 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 { + 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> { + byok::reorder(&self.pool().await?, input).await.map_err(to_napi_error) + } + + #[napi] + pub async fn create_byok_local_lease(&self, input: CreateByokLocalLeaseInput) -> Result { + 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 { 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 { + pub(crate) fn config(&self) -> RuntimeResult> { 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(()) } } diff --git a/packages/backend/native/src/runtime/backend_runtime/rolling_quota/mod.rs b/packages/backend/native/src/runtime/backend_runtime/rolling_quota/mod.rs index 912d232859..7eac481007 100644 --- a/packages/backend/native/src/runtime/backend_runtime/rolling_quota/mod.rs +++ b/packages/backend/native/src/runtime/backend_runtime/rolling_quota/mod.rs @@ -75,8 +75,9 @@ impl BackendRuntime { &self, input: RuntimeMailDeliveryQuotaInput, ) -> Result { - 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, diff --git a/packages/backend/native/src/runtime/backend_runtime/rolling_quota/workspace_invite.rs b/packages/backend/native/src/runtime/backend_runtime/rolling_quota/workspace_invite.rs index bfde78f11c..560006d032 100644 --- a/packages/backend/native/src/runtime/backend_runtime/rolling_quota/workspace_invite.rs +++ b/packages/backend/native/src/runtime/backend_runtime/rolling_quota/workspace_invite.rs @@ -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 = 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, "a, &activity, &config, now)?; + let scopes = build_invite_scopes(&input, &actor, &workspace, "a, &activity, config, now)?; match reserve_scopes(&pool, "workspace_invite", input.request_id.as_deref(), scopes).await? { Ok(reservation) => Ok(RuntimeWorkspaceInviteQuotaDecision { allowed: true, diff --git a/packages/backend/native/src/runtime/backend_runtime/runtime_state/byok_local_lease.rs b/packages/backend/native/src/runtime/backend_runtime/runtime_state/byok_local_lease.rs deleted file mode 100644 index 1318bd392a..0000000000 --- a/packages/backend/native/src/runtime/backend_runtime/runtime_state/byok_local_lease.rs +++ /dev/null @@ -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> { - 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 { - 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> { - 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> { - 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 { - Ok(RuntimeByokLocalLeaseRecord { - lease_id: lease_id.to_string(), - payload: row.payload, - expires_at_ms: row.expires_at_ms, - }) -} diff --git a/packages/backend/native/src/runtime/backend_runtime/runtime_state/mod.rs b/packages/backend/native/src/runtime/backend_runtime/runtime_state/mod.rs index bee1c15e7e..902634f5d6 100644 --- a/packages/backend/native/src/runtime/backend_runtime/runtime_state/mod.rs +++ b/packages/backend/native/src/runtime/backend_runtime/runtime_state/mod.rs @@ -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 { - 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> { - 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 { if limit <= 0 { diff --git a/packages/backend/native/src/runtime/backend_runtime/runtime_state/store.rs b/packages/backend/native/src/runtime/backend_runtime/runtime_state/store.rs index 6eba25f63f..db23105ec8 100644 --- a/packages/backend/native/src/runtime/backend_runtime/runtime_state/store.rs +++ b/packages/backend/native/src/runtime/backend_runtime/runtime_state/store.rs @@ -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 { 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 { - 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> { - byok_local_lease::get(&self.rows, lease_id).await - } } diff --git a/packages/backend/native/src/runtime/backend_runtime/tests.rs b/packages/backend/native/src/runtime/backend_runtime/tests.rs index 885c0f14f4..48a60a25bd 100644 --- a/packages/backend/native/src/runtime/backend_runtime/tests.rs +++ b/packages/backend/native/src/runtime/backend_runtime/tests.rs @@ -97,11 +97,14 @@ async fn runtime_from_database_url() -> AnyResult> { .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 { diff --git a/packages/backend/native/src/runtime/config.rs b/packages/backend/native/src/runtime/config.rs index 529dfdbe49..854c4cc7f6 100644 --- a/packages/backend/native/src/runtime/config.rs +++ b/packages/backend/native/src/runtime/config.rs @@ -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>, + 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, +} + +#[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, + 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 { + pub(crate) fn from_config_files(private_key: Option) -> RuntimeResult { 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 { - 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 { + 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 { + 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, +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, + crypto: Option, + copilot: Option, +} + +#[derive(Default, Deserialize)] +#[serde(rename_all = "camelCase")] +struct CryptoConfigFile { + private_key: Option, +} + +#[derive(Default, Deserialize)] #[serde(rename_all = "camelCase")] struct DbConfigFile { datasource_url: Option, @@ -97,31 +229,67 @@ fn database_url_from_env() -> Option { env::var("DATABASE_URL").ok().and_then(non_empty_string) } +fn private_key_from_env() -> Option { + env::var("AFFINE_PRIVATE_KEY").ok().and_then(non_empty_string) +} + fn non_empty_string(value: String) -> Option { if value.trim().is_empty() { None } else { Some(value) } } fn app_config_from_config_files() -> RuntimeResult { - let mut merged = AppConfigFile::default(); + deserialize_app_config(app_config_value_from_config_files()?) +} + +fn app_config_value_from_config_files() -> RuntimeResult { + 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 { + deserialize_app_config(expand_module_config_paths(value)) +} + +fn deserialize_app_config(value: serde_json::Value) -> RuntimeResult { + 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 { @@ -161,40 +329,60 @@ fn default_mail_class_mapping() -> BTreeMap { .collect() } -async fn load_app_config_overrides_from_db(pool: &PgPool) -> RuntimeResult { +async fn load_app_config_overrides_from_db(pool: &PgPool) -> RuntimeResult { 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(rows: I) -> RuntimeResult +where + I: IntoIterator, + S: AsRef, +{ + deserialize_app_config(app_config_value_from_flat_overrides(rows)) +} + +fn app_config_value_from_flat_overrides(rows: I) -> serde_json::Value where I: IntoIterator, S: AsRef, { 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, 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 { @@ -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([ diff --git a/packages/backend/native/src/runtime/mod.rs b/packages/backend/native/src/runtime/mod.rs index 6449b56fb6..881808dbe5 100644 --- a/packages/backend/native/src/runtime/mod.rs +++ b/packages/backend/native/src/runtime/mod.rs @@ -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}; diff --git a/packages/backend/native/src/runtime/types.rs b/packages/backend/native/src/runtime/types.rs index ae926bf50d..2c0d4c9181 100644 --- a/packages/backend/native/src/runtime/types.rs +++ b/packages/backend/native/src/runtime/types.rs @@ -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, diff --git a/packages/backend/server/migrations/20260803095500_converge_copilot_runtime/migration.sql b/packages/backend/server/migrations/20260803095500_converge_copilot_runtime/migration.sql new file mode 100644 index 0000000000..d8c461ee6d --- /dev/null +++ b/packages/backend/server/migrations/20260803095500_converge_copilot_runtime/migration.sql @@ -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"; diff --git a/packages/backend/server/schema.prisma b/packages/backend/server/schema.prisma index cf2332bf10..0718c83aad 100644 --- a/packages/backend/server/schema.prisma +++ b/packages/backend/server/schema.prisma @@ -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 diff --git a/packages/backend/server/src/__tests__/copilot/__snapshots__/copilot.e2e.ts.md b/packages/backend/server/src/__tests__/copilot/__snapshots__/copilot.e2e.ts.md deleted file mode 100644 index 2bf58ee451..0000000000 --- a/packages/backend/server/src/__tests__/copilot/__snapshots__/copilot.e2e.ts.md +++ /dev/null @@ -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, - }, - ] diff --git a/packages/backend/server/src/__tests__/copilot/__snapshots__/copilot.spec.ts.md b/packages/backend/server/src/__tests__/copilot/__snapshots__/copilot.spec.ts.md deleted file mode 100644 index 1650e7f294..0000000000 --- a/packages/backend/server/src/__tests__/copilot/__snapshots__/copilot.spec.ts.md +++ /dev/null @@ -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' diff --git a/packages/backend/server/src/__tests__/copilot/__snapshots__/native-provider.spec.ts.md b/packages/backend/server/src/__tests__/copilot/__snapshots__/native-provider.spec.ts.md deleted file mode 100644 index 62f6c47dd2..0000000000 --- a/packages/backend/server/src/__tests__/copilot/__snapshots__/native-provider.spec.ts.md +++ /dev/null @@ -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, - ], - } diff --git a/packages/backend/server/src/__tests__/copilot/__snapshots__/provider-native.spec.ts.md b/packages/backend/server/src/__tests__/copilot/__snapshots__/provider-native.spec.ts.md deleted file mode 100644 index 9dc47b3221..0000000000 --- a/packages/backend/server/src/__tests__/copilot/__snapshots__/provider-native.spec.ts.md +++ /dev/null @@ -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: [], - }, - }, - ] diff --git a/packages/backend/server/src/__tests__/copilot/byok-probe.spec.ts b/packages/backend/server/src/__tests__/copilot/byok-probe.spec.ts deleted file mode 100644 index 70a4e4db5f..0000000000 --- a/packages/backend/server/src/__tests__/copilot/byok-probe.spec.ts +++ /dev/null @@ -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, - ReturnType - >().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, - ReturnType - >().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); -}); diff --git a/packages/backend/server/src/__tests__/copilot/byok.spec.ts b/packages/backend/server/src/__tests__/copilot/byok.spec.ts index 10f8634509..14583e503f 100644 --- a/packages/backend/server/src/__tests__/copilot/byok.spec.ts +++ b/packages/backend/server/src/__tests__/copilot/byok.spec.ts @@ -1,56 +1,66 @@ -import { createHash, randomUUID } from 'node:crypto'; +import { generateKeyPairSync, randomUUID } from 'node:crypto'; -import { PrismaClient, WorkspaceMemberStatus } from '@prisma/client'; -import ava, { type ExecutionContext, type TestFn } from 'ava'; -import Sinon from 'sinon'; +import { PrismaClient } from '@prisma/client'; +import type { TestFn } from 'ava'; +import ava from 'ava'; -import { Cache, ConfigFactory, CryptoHelper } from '../../base'; -import { EntitlementService } from '../../core/entitlement'; -import { Models, WorkspaceRole } from '../../models'; -import { CopilotAccessPolicy } from '../../plugins/copilot/access'; -import { ByokService } from '../../plugins/copilot/byok'; -import { - type ByokFeatureKind, - ByokKeyStorage, - ByokKeyTestStatus, - ByokProvider, -} from '../../plugins/copilot/byok/types'; -import { - SubscriptionPlan, - SubscriptionRecurring, - SubscriptionStatus, -} from '../../plugins/payment/types'; +import { Config } from '../../base'; +import type { CurrentUser } from '../../core/auth'; +import { BackendRuntimeProvider } from '../../core/backend-runtime'; +import type { WorkspaceType } from '../../core/workspaces'; +import { Models } from '../../models'; +import { WorkspaceByokResolver } from '../../plugins/copilot/byok/resolver'; import { createTestingModule, type TestingModule } from '../utils'; -interface Context { +type Context = { module: TestingModule; - models: Models; db: PrismaClient; - access: CopilotAccessPolicy; - byok: ByokService; - crypto: CryptoHelper; - cache: Cache; - entitlement: EntitlementService; -} + models: Models; + runtime: BackendRuntimeProvider; + config: Config; + resolver: WorkspaceByokResolver; +}; const test = ava.serial as TestFn; -const originalNamespace = globalThis.env.NAMESPACE; -const originalDeploymentType = globalThis.env.DEPLOYMENT_TYPE; +const previousKey = process.env.AFFINE_PRIVATE_KEY; +const { privateKey } = generateKeyPairSync('ec', { + namedCurve: 'P-256', +}); +const testPrivateKey = privateKey + .export({ + format: 'pem', + type: 'pkcs8', + }) + .toString(); + +const definition = { + version: 1, + endpoint: { kind: 'provider_default' }, + models: [ + { + modelId: 'gpt-4o-mini', + enabled: true, + capabilities: [ + { + input: ['text'], + output: ['text'], + features: [], + attachmentKinds: [], + attachmentSources: [], + }, + ], + }, + ], +}; test.before(async t => { - Object.assign(globalThis.env, { - NAMESPACE: 'dev', - DEPLOYMENT_TYPE: 'affine', - }); - const module = await createTestingModule(); - t.context.module = module; - t.context.models = module.get(Models); - t.context.db = module.get(PrismaClient); - t.context.access = module.get(CopilotAccessPolicy); - t.context.byok = module.get(ByokService); - t.context.crypto = module.get(CryptoHelper); - t.context.cache = module.get(Cache); - t.context.entitlement = module.get(EntitlementService); + process.env.AFFINE_PRIVATE_KEY = testPrivateKey; + t.context.module = await createTestingModule(); + t.context.db = t.context.module.get(PrismaClient); + t.context.models = t.context.module.get(Models); + t.context.runtime = t.context.module.get(BackendRuntimeProvider); + t.context.config = t.context.module.get(Config); + t.context.resolver = t.context.module.get(WorkspaceByokResolver); }); test.beforeEach(async t => { @@ -58,1221 +68,180 @@ test.beforeEach(async t => { }); test.after.always(async t => { - await t.context.module.close(); - Object.assign(globalThis.env, { - NAMESPACE: originalNamespace, - DEPLOYMENT_TYPE: originalDeploymentType, - }); + await t.context.module?.close(); + if (previousKey === undefined) delete process.env.AFFINE_PRIVATE_KEY; + else process.env.AFFINE_PRIVATE_KEY = previousKey; }); -async function createUserWorkspace(t: ExecutionContext) { +test('BYOK settings expose the configured custom endpoint policy', async t => { const user = await t.context.models.user.create({ email: `${randomUUID()}@affine.pro`, }); const workspace = await t.context.models.workspace.create(user.id); - return { user, workspace }; -} + const previous = t.context.config.copilot.byok.allowCustomEndpoint; + t.context.config.copilot.byok.allowCustomEndpoint = true; -function workspaceHash(workspaceId: string) { - return createHash('sha256').update(workspaceId).digest('hex').slice(0, 12); -} - -async function grantUserPlan( - t: ExecutionContext, - userId: string, - feature: ByokUserPlanFeature = 'pro_plan_v1' -) { - if (feature === 'unlimited_copilot') { - await t.context.entitlement.upsertFromCloudSubscription({ - targetId: userId, - plan: SubscriptionPlan.AI, - recurring: SubscriptionRecurring.Monthly, - status: SubscriptionStatus.Active, - }); - return; - } - - await t.context.entitlement.upsertFromCloudSubscription({ - targetId: userId, - plan: SubscriptionPlan.Pro, - recurring: - feature === 'lifetime_pro_plan_v1' - ? SubscriptionRecurring.Lifetime - : SubscriptionRecurring.Monthly, - status: SubscriptionStatus.Active, - }); -} - -async function revokeUserPlan( - t: ExecutionContext, - userId: string, - feature: ByokUserPlanFeature = 'pro_plan_v1' -) { - if (feature === 'unlimited_copilot') { - await t.context.entitlement.revokeCloudSubscription({ - targetId: userId, - plan: SubscriptionPlan.AI, - }); - return; - } - - await t.context.entitlement.revokeCloudSubscription({ - targetId: userId, - plan: SubscriptionPlan.Pro, - }); -} - -async function grantTeamPlan( - t: ExecutionContext, - workspaceId: string -) { - await t.context.entitlement.upsertFromCloudSubscription({ - targetId: workspaceId, - plan: SubscriptionPlan.Team, - recurring: SubscriptionRecurring.Yearly, - status: SubscriptionStatus.Active, - }); -} - -async function revokeTeamPlan( - t: ExecutionContext, - workspaceId: string -) { - await t.context.entitlement.revokeCloudSubscription({ - targetId: workspaceId, - plan: SubscriptionPlan.Team, - }); -} - -type ByokMatrixCase = { - name: string; - role: WorkspaceRole; - team?: boolean; - ownerPlan?: boolean; - ownerPlanFeature?: ByokUserPlanFeature; - actorPlan?: boolean; - actorPlanFeature?: ByokUserPlanFeature; - settings: { - entitled: boolean; - serverEntitled: boolean; - localEntitled: boolean; - }; - canConfigureServer: boolean; - canConfigureLocal: boolean; -}; - -type ByokUserPlanFeature = - | 'pro_plan_v1' - | 'lifetime_pro_plan_v1' - | 'unlimited_copilot'; - -async function createByokMatrixWorkspace( - t: ExecutionContext, - input: Pick< - ByokMatrixCase, - | 'role' - | 'team' - | 'ownerPlan' - | 'ownerPlanFeature' - | 'actorPlan' - | 'actorPlanFeature' - > -) { - const { user: owner, workspace } = await createUserWorkspace(t); - const actor = - input.role === WorkspaceRole.Owner - ? owner - : await t.context.models.user.create({ - email: `${randomUUID()}@affine.pro`, - }); - - if (input.role !== WorkspaceRole.Owner) { - await t.context.models.workspaceUser.set( - workspace.id, - actor.id, - input.role, - { status: WorkspaceMemberStatus.Accepted } + try { + const settings = await t.context.resolver.settings( + { + id: user.id, + email: user.email, + avatarUrl: user.avatarUrl, + name: user.name, + disabled: user.disabled, + hasPassword: null, + emailVerified: true, + } satisfies CurrentUser, + { id: workspace.id } as WorkspaceType ); + t.true(settings.customEndpointSupported); + } finally { + t.context.config.copilot.byok.allowCustomEndpoint = previous; } - if (input.team) { - await grantTeamPlan(t, workspace.id); - } - if (input.ownerPlan) { - await grantUserPlan(t, owner.id, input.ownerPlanFeature); - } - if (input.actorPlan && actor.id !== owner.id) { - await grantUserPlan(t, actor.id, input.actorPlanFeature); - } - - return { owner, actor, workspace }; -} - -const byokManagementMatrix: ByokMatrixCase[] = [ - { - name: 'owner without plan in a personal workspace', - role: WorkspaceRole.Owner, - settings: { entitled: false, serverEntitled: false, localEntitled: false }, - canConfigureServer: false, - canConfigureLocal: false, - }, - { - name: 'owner with BYOK plan in a personal workspace', - role: WorkspaceRole.Owner, - ownerPlan: true, - settings: { entitled: true, serverEntitled: true, localEntitled: true }, - canConfigureServer: true, - canConfigureLocal: true, - }, - { - name: 'owner in team workspace without user plan', - role: WorkspaceRole.Owner, - team: true, - settings: { entitled: true, serverEntitled: true, localEntitled: true }, - canConfigureServer: true, - canConfigureLocal: true, - }, - { - name: 'admin in believer owner-backed personal workspace', - role: WorkspaceRole.Admin, - ownerPlan: true, - ownerPlanFeature: 'unlimited_copilot', - settings: { entitled: true, serverEntitled: true, localEntitled: true }, - canConfigureServer: true, - canConfigureLocal: true, - }, - { - name: 'admin with own lifetime plan but no owner-backed server entitlement', - role: WorkspaceRole.Admin, - actorPlan: true, - actorPlanFeature: 'lifetime_pro_plan_v1', - settings: { entitled: true, serverEntitled: false, localEntitled: true }, - canConfigureServer: false, - canConfigureLocal: true, - }, - { - name: 'admin without plan in non-entitled personal workspace', - role: WorkspaceRole.Admin, - settings: { entitled: false, serverEntitled: false, localEntitled: false }, - canConfigureServer: false, - canConfigureLocal: false, - }, - { - name: 'admin in team workspace without user plan', - role: WorkspaceRole.Admin, - team: true, - settings: { entitled: true, serverEntitled: true, localEntitled: true }, - canConfigureServer: true, - canConfigureLocal: true, - }, - { - name: 'ordinary member in team workspace without user plan', - role: WorkspaceRole.Collaborator, - team: true, - settings: { entitled: false, serverEntitled: false, localEntitled: false }, - canConfigureServer: false, - canConfigureLocal: false, - }, - { - name: 'ordinary member with own plan', - role: WorkspaceRole.Collaborator, - ownerPlan: true, - actorPlan: true, - settings: { entitled: false, serverEntitled: false, localEntitled: false }, - canConfigureServer: false, - canConfigureLocal: false, - }, -]; - -for (const matrixCase of byokManagementMatrix) { - test(`BYOK management entitlement: ${matrixCase.name}`, async t => { - const { actor, workspace } = await createByokMatrixWorkspace(t, matrixCase); - const settings = await t.context.byok.getSettings(workspace.id, actor.id); - - t.like(settings, matrixCase.settings); - - const serverConfig = t.context.byok.upsertConfig({ - workspaceId: workspace.id, - userId: actor.id, - provider: ByokProvider.openai, - storage: ByokKeyStorage.server, - name: 'Server', - apiKey: 'sk-server', - }); - if (matrixCase.canConfigureServer) { - t.truthy(await serverConfig); - } else { - await t.throwsAsync(serverConfig); - } - - const localLease = t.context.byok.createLocalLease({ - workspaceId: workspace.id, - userId: actor.id, - providers: [ - { - provider: ByokProvider.openai, - name: 'Local', - apiKey: 'sk-local', - }, - ], - }); - if (matrixCase.canConfigureLocal) { - t.truthy(await localLease); - } else { - await t.throwsAsync(localLease); - } - }); -} - -test('byok service persists encrypted server keys and never returns plaintext', async t => { - const { user, workspace } = await createUserWorkspace(t); - await grantUserPlan(t, user.id); - - const primary = await t.context.byok.upsertConfig({ - workspaceId: workspace.id, - userId: user.id, - provider: ByokProvider.openai, - storage: ByokKeyStorage.server, - name: 'Primary', - description: 'Team key', - apiKey: 'sk-test-primary', - sortOrder: 1, - }); - const backup = await t.context.byok.upsertConfig({ - workspaceId: workspace.id, - userId: user.id, - provider: ByokProvider.openai, - storage: ByokKeyStorage.server, - name: 'Backup', - apiKey: 'sk-test-backup', - sortOrder: 2, - }); - await t.throwsAsync( - t.context.byok.upsertConfig({ - workspaceId: workspace.id, - userId: user.id, - provider: ByokProvider.openai, - storage: ByokKeyStorage.server, - name: 'Primary', - apiKey: 'sk-test-duplicate', - }) - ); - - t.true(primary.configured); - t.is(primary.storage, ByokKeyStorage.server); - t.false(JSON.stringify(primary).includes('sk-test-primary')); - - const row = await t.context.models.copilotWorkspaceByokConfig.get(primary.id); - t.truthy(row); - if (!row) { - return; - } - t.not(row.encryptedApiKey, 'sk-test-primary'); - t.is(t.context.crypto.decrypt(row.encryptedApiKey), 'sk-test-primary'); - - const reordered = await t.context.byok.reorderConfigs({ - workspaceId: workspace.id, - userId: user.id, - storage: ByokKeyStorage.server, - ids: [backup.id, primary.id], - }); - t.deepEqual( - reordered.map(key => [key.id, key.sortOrder]), - [ - [backup.id, 0], - [primary.id, 1], - ] - ); - - const profiles = await t.context.byok.getProfiles({ - workspaceId: workspace.id, - userId: user.id, - }); - t.deepEqual( - profiles.map(profile => profile.id), - [ - `byok-${workspaceHash(workspace.id)}-openai-${backup.id}`, - `byok-${workspaceHash(workspace.id)}-openai-${primary.id}`, - ] - ); }); -test('byok service preserves server key fields during partial updates', async t => { - const { user, workspace } = await createUserWorkspace(t); - await grantUserPlan(t, user.id); +test('native BYOK runtime owns multi-model profile CAS, ordering, and credential rotation', async t => { + t.is(typeof BackendRuntimeProvider.prototype.probeByokProfile, 'function'); + t.false('runProviderProbe' in BackendRuntimeProvider.prototype); - const key = await t.context.byok.upsertConfig({ - workspaceId: workspace.id, - userId: user.id, - provider: ByokProvider.openai, - storage: ByokKeyStorage.server, - name: 'Primary', - description: 'Team key', - apiKey: 'sk-test-primary', - sortOrder: 3, - enabled: false, - }); - - await t.context.byok.upsertConfig({ - id: key.id, - workspaceId: workspace.id, - userId: user.id, - provider: ByokProvider.openai, - storage: ByokKeyStorage.server, - name: 'Primary renamed', - apiKey: 'sk-test-primary-next', - }); - - const updated = await t.context.models.copilotWorkspaceByokConfig.get(key.id); - t.truthy(updated); - if (!updated) { - return; - } - t.is(updated.name, 'Primary renamed'); - t.is(updated.description, 'Team key'); - t.is( - t.context.crypto.decrypt(updated.encryptedApiKey), - 'sk-test-primary-next' - ); - t.is(updated.sortOrder, 3); - t.false(updated.enabled); - - await t.context.byok.upsertConfig({ - id: key.id, - workspaceId: workspace.id, - userId: user.id, - provider: ByokProvider.openai, - storage: ByokKeyStorage.server, - name: 'Primary renamed', - description: null, - }); - - const cleared = await t.context.models.copilotWorkspaceByokConfig.get(key.id); - t.is(cleared?.description, null); - t.is(cleared?.sortOrder, 3); - t.false(cleared?.enabled ?? true); -}); - -test('local leases are short lived and do not persist keys to server configs', async t => { - const { user, workspace } = await createUserWorkspace(t); - await grantUserPlan(t, user.id); - - const before = Date.now(); - const lease = await t.context.byok.createLocalLease({ - workspaceId: workspace.id, - userId: user.id, - providers: [ - { - provider: ByokProvider.openai, - name: 'Local', - apiKey: 'sk-local', - }, - ], - }); - const reusedLease = await t.context.byok.createLocalLease({ - workspaceId: workspace.id, - userId: user.id, - providers: [ - { - provider: ByokProvider.openai, - name: 'Local', - apiKey: 'sk-local', - }, - ], - }); - t.is(reusedLease.leaseId, lease.leaseId); - const cachedLease = await t.context.cache.get<{ - providers: Array<{ apiKey?: string; encryptedApiKey?: string }>; - }>(`copilot:byok:lease:${lease.leaseId}`); - t.truthy(cachedLease); - t.false(JSON.stringify(cachedLease).includes('sk-local')); - t.is(cachedLease?.providers[0].apiKey, undefined); - t.truthy(cachedLease?.providers[0].encryptedApiKey); - - const updatedLease = await t.context.byok.createLocalLease({ - workspaceId: workspace.id, - userId: user.id, - providers: [ - { - provider: ByokProvider.openai, - name: 'Local', - apiKey: 'sk-local-updated', - }, - ], - }); - t.not(updatedLease.leaseId, lease.leaseId); - - const lifetime = lease.expiresAt.getTime() - before; - t.true(lifetime >= 5 * 60 * 1000); - t.true(lifetime <= 15 * 60 * 1000); - t.deepEqual( - await t.context.models.copilotWorkspaceByokConfig.list(workspace.id), - [] - ); - - const profiles = await t.context.byok.getProfiles({ - workspaceId: workspace.id, - userId: user.id, - byokLeaseId: lease.leaseId, - }); - t.deepEqual( - profiles.map(profile => profile.type), - ['openai'] - ); - - const otherWorkspace = await t.context.models.workspace.create(user.id); - t.deepEqual( - await t.context.byok.getProfiles({ - workspaceId: otherWorkspace.id, - userId: user.id, - byokLeaseId: lease.leaseId, - }), - [] - ); - - await t.context.cache.delete(`copilot:byok:lease:${lease.leaseId}`); - t.deepEqual( - await t.context.byok.getProfiles({ - workspaceId: workspace.id, - userId: user.id, - byokLeaseId: lease.leaseId, - }), - [] - ); - const renewedLease = await t.context.byok.createLocalLease({ - workspaceId: workspace.id, - userId: user.id, - providers: [ - { - provider: ByokProvider.openai, - name: 'Local', - apiKey: 'sk-local', - }, - ], - }); - t.not(renewedLease.leaseId, lease.leaseId); -}); - -test('local leases persist normalized custom endpoints', async t => { - const config = t.context.module.get(ConfigFactory); - const deploymentType = globalThis.env.DEPLOYMENT_TYPE; - Object.assign(globalThis.env, { DEPLOYMENT_TYPE: 'selfhosted' }); - t.false(t.context.byok.customEndpointSupported); - config.override({ copilot: { byok: { allowCustomEndpoint: true } } }); - t.true(t.context.byok.customEndpointSupported); - t.teardown(() => { - config.override({ copilot: { byok: { allowCustomEndpoint: false } } }); - Object.assign(globalThis.env, { DEPLOYMENT_TYPE: deploymentType }); - }); - const { user, workspace } = await createUserWorkspace(t); - await grantUserPlan(t, user.id); - - const lease = await t.context.byok.createLocalLease({ - workspaceId: workspace.id, - userId: user.id, - providers: [ - { - provider: ByokProvider.openai, - name: 'Local', - apiKey: 'sk-local', - endpoint: 'https://api.openai.example/v1/', - }, - ], - }); - const reusedLease = await t.context.byok.createLocalLease({ - workspaceId: workspace.id, - userId: user.id, - providers: [ - { - provider: ByokProvider.openai, - name: 'Local', - apiKey: 'sk-local', - endpoint: 'https://api.openai.example/v1', - }, - ], - }); - - t.is(reusedLease.leaseId, lease.leaseId); - const cachedLease = await t.context.cache.get<{ - providers: Array<{ endpoint?: string | null }>; - }>(`copilot:byok:lease:${lease.leaseId}`); - t.is(cachedLease?.providers[0]?.endpoint, 'https://api.openai.example/v1'); - - const profiles = await t.context.byok.getProfiles({ - workspaceId: workspace.id, - userId: user.id, - byokLeaseId: lease.leaseId, - }); - t.is( - (profiles[0]!.config as { baseURL?: string }).baseURL, - 'https://api.openai.example/v1' - ); -}); - -type ByokProfileAvailabilityCase = { - name: string; - actorRole: WorkspaceRole; - ownerPlan?: boolean; - actorPlan?: boolean; - team?: boolean; - createServerKey?: boolean; - createActorLocalLease?: boolean; - createOwnerLocalLease?: boolean; - revokeOwnerPlan?: boolean; - revokeTeam?: boolean; - demoteActor?: boolean; - expectedSources: Array<'server' | 'local'>; -}; - -const byokProfileAvailabilityMatrix: ByokProfileAvailabilityCase[] = [ - { - name: 'ordinary members can use server BYOK while the owner is entitled', - actorRole: WorkspaceRole.Collaborator, - ownerPlan: true, - createServerKey: true, - expectedSources: ['server'], - }, - { - name: 'ordinary members can use server BYOK in team workspaces', - actorRole: WorkspaceRole.Collaborator, - team: true, - createServerKey: true, - expectedSources: ['server'], - }, - { - name: 'ordinary members cannot use another user local BYOK lease', - actorRole: WorkspaceRole.Collaborator, - ownerPlan: true, - createOwnerLocalLease: true, - expectedSources: [], - }, - { - name: 'owner-backed server and local BYOK stop after owner plan is removed', - actorRole: WorkspaceRole.Admin, - ownerPlan: true, - createServerKey: true, - createActorLocalLease: true, - revokeOwnerPlan: true, - expectedSources: [], - }, - { - name: 'admin local BYOK remains available after owner plan removal when admin is entitled', - actorRole: WorkspaceRole.Admin, - ownerPlan: true, - actorPlan: true, - createServerKey: true, - createActorLocalLease: true, - revokeOwnerPlan: true, - expectedSources: ['local'], - }, - { - name: 'team BYOK stops after team entitlement is removed without user plan', - actorRole: WorkspaceRole.Admin, - team: true, - createServerKey: true, - createActorLocalLease: true, - revokeTeam: true, - expectedSources: [], - }, - { - name: 'local BYOK lease stops after an admin is demoted', - actorRole: WorkspaceRole.Admin, - actorPlan: true, - createActorLocalLease: true, - demoteActor: true, - expectedSources: [], - }, -]; - -for (const matrixCase of byokProfileAvailabilityMatrix) { - test(`BYOK profile availability: ${matrixCase.name}`, async t => { - const { owner, actor, workspace } = await createByokMatrixWorkspace(t, { - role: matrixCase.actorRole, - team: matrixCase.team, - ownerPlan: matrixCase.ownerPlan, - actorPlan: matrixCase.actorPlan, - }); - - if (matrixCase.createServerKey) { - const creator = - matrixCase.actorRole === WorkspaceRole.Collaborator ? owner : actor; - await t.context.byok.upsertConfig({ - workspaceId: workspace.id, - userId: creator.id, - provider: ByokProvider.openai, - storage: ByokKeyStorage.server, - name: 'Server', - apiKey: 'sk-server', - }); - } - - let leaseId: string | undefined; - if (matrixCase.createActorLocalLease) { - leaseId = ( - await t.context.byok.createLocalLease({ - workspaceId: workspace.id, - userId: actor.id, - providers: [ - { - provider: ByokProvider.openai, - name: 'Local', - apiKey: 'sk-local', - }, - ], - }) - ).leaseId; - } else if (matrixCase.createOwnerLocalLease) { - leaseId = ( - await t.context.byok.createLocalLease({ - workspaceId: workspace.id, - userId: owner.id, - providers: [ - { - provider: ByokProvider.openai, - name: 'Local', - apiKey: 'sk-local', - }, - ], - }) - ).leaseId; - } - - if (matrixCase.revokeOwnerPlan) { - await revokeUserPlan(t, owner.id); - } - if (matrixCase.revokeTeam) { - await revokeTeamPlan(t, workspace.id); - } - if (matrixCase.demoteActor) { - await t.context.models.workspaceUser.set( - workspace.id, - actor.id, - WorkspaceRole.Collaborator, - { status: WorkspaceMemberStatus.Accepted } - ); - } - - const profiles = await t.context.byok.getProfiles({ - workspaceId: workspace.id, - userId: actor.id, - byokLeaseId: leaseId, - }); - - t.deepEqual( - profiles.map(profile => - profile.id.includes('-local-') ? 'local' : 'server' - ), - matrixCase.expectedSources - ); - }); -} - -test('BYOK profile availability: local-only workspace does not resolve BYOK profiles', async t => { const user = await t.context.models.user.create({ email: `${randomUUID()}@affine.pro`, }); - await grantUserPlan(t, user.id); - - const profiles = await t.context.byok.getProfiles({ - workspaceId: randomUUID(), - userId: user.id, - }); - - t.deepEqual(profiles, []); -}); - -test('test key failure disables a saved key and success restores it', async t => { - const { user, workspace } = await createUserWorkspace(t); - await grantUserPlan(t, user.id); - const key = await t.context.byok.upsertConfig({ + const workspace = await t.context.models.workspace.create(user.id); + const created = await t.context.runtime.createByokProfile({ workspaceId: workspace.id, - userId: user.id, - provider: ByokProvider.openai, - storage: ByokKeyStorage.server, - name: 'Primary', - apiKey: 'sk-test-primary', + provider: 'openai', + name: 'OpenAI', + description: undefined, + credential: 'first-secret', + definition, + enabled: true, + actorUserId: user.id, }); + t.is(created.definition.models[0].modelId, 'gpt-4o-mini'); + t.is(created.validation, undefined); + t.false(JSON.stringify(created).includes('first-secret')); + const stored = await t.context.db.aiWorkspaceByokConfig.findUniqueOrThrow({ + where: { id: created.profileId }, + }); + t.not(stored.encryptedApiKey, 'first-secret'); + t.false(stored.encryptedApiKey.includes('first-secret')); - const fetch = Sinon.stub(t.context.byok as any, 'probeFetch'); - fetch - .onFirstCall() - .resolves( - new Response('{"error":"invalid sk-test-primary"}', { status: 401 }) - ); - fetch.onSecondCall().resolves(new Response('{}', { status: 200 })); - t.teardown(() => fetch.restore()); - - const failed = await t.context.byok.testConfig({ + const multiModelDefinition = { + ...definition, + models: [ + definition.models[0], + { + modelId: 'text-embedding-3-small', + enabled: false, + capabilities: [ + { + input: ['text'], + output: ['embedding'], + features: [], + attachmentKinds: [], + attachmentSources: [], + }, + ], + }, + ], + }; + const replaced = await t.context.runtime.replaceByokProfile({ workspaceId: workspace.id, - userId: user.id, - provider: ByokProvider.openai, - storage: ByokKeyStorage.server, - apiKey: 'sk-test-primary', - configId: key.id, + profileId: created.profileId, + expectedRevision: created.revision, + name: 'OpenAI models', + definition: multiModelDefinition, + enabled: true, + actorUserId: user.id, }); - t.false(failed.ok); - t.is(failed.status, ByokKeyTestStatus.failed); - t.false(failed.message?.includes('sk-test-primary')); + t.is(replaced.revision, created.revision + 1); + t.is(replaced.definition.models.length, 2); + t.false(replaced.definition.models[1].enabled); - const disabled = await t.context.models.copilotWorkspaceByokConfig.get( - key.id + const conflict = await t.throwsAsync( + t.context.runtime.replaceByokProfile({ + workspaceId: workspace.id, + profileId: created.profileId, + expectedRevision: created.revision, + name: 'stale update', + definition, + enabled: true, + actorUserId: user.id, + }) ); - t.truthy(disabled); - if (!disabled) { - return; - } - t.false(disabled.enabled); - t.is(disabled.disabledReason, 'recent_failure'); + t.regex(conflict.message, /byok_revision_conflict/); - const passed = await t.context.byok.testConfig({ + const rotated = await t.context.runtime.rotateByokCredential({ workspaceId: workspace.id, - userId: user.id, - provider: ByokProvider.openai, - storage: ByokKeyStorage.server, - configId: key.id, + profileId: created.profileId, + expectedRevision: replaced.revision, + credential: 'second-secret', + actorUserId: user.id, }); - t.true(passed.ok); - const restored = await t.context.models.copilotWorkspaceByokConfig.get( - key.id - ); - t.truthy(restored); - if (!restored) { - return; - } - t.true(restored.enabled); - t.is(restored.disabledReason, null); + t.is(rotated.profileId, created.profileId); - const profiles = await t.context.byok.getProfiles({ + const second = await t.context.runtime.createByokProfile({ workspaceId: workspace.id, - userId: user.id, + provider: 'openai', + name: 'Fallback', + credential: 'fallback-secret', + definition, + enabled: true, + actorUserId: user.id, + }); + const reordered = await t.context.runtime.reorderByokProfiles({ + workspaceId: workspace.id, + profiles: [ + { profileId: second.profileId, expectedRevision: second.revision }, + { profileId: created.profileId, expectedRevision: rotated.revision }, + ], + actorUserId: user.id, }); t.deepEqual( - profiles.map(profile => profile.type), - ['openai'] + reordered.map(profile => profile.profileId), + [second.profileId, created.profileId] ); -}); - -test('local key test does not mutate saved server config', async t => { - const { user, workspace } = await createUserWorkspace(t); - await grantUserPlan(t, user.id); - const key = await t.context.byok.upsertConfig({ - workspaceId: workspace.id, - userId: user.id, - provider: ByokProvider.openai, - storage: ByokKeyStorage.server, - name: 'Server', - apiKey: 'sk-server', - }); - - const fetch = Sinon.stub(t.context.byok as any, 'probeFetch').resolves( - new Response('{"error":"invalid sk-local"}', { status: 401 }) - ); - t.teardown(() => fetch.restore()); - - const failed = await t.context.byok.testConfig({ - workspaceId: workspace.id, - userId: user.id, - provider: ByokProvider.openai, - storage: ByokKeyStorage.local, - apiKey: 'sk-local', - configId: key.id, - }); - t.false(failed.ok); - - const unchanged = await t.context.models.copilotWorkspaceByokConfig.get( - key.id - ); - t.truthy(unchanged); - if (!unchanged) { - return; - } - t.true(unchanged.enabled); - t.is(unchanged.disabledReason, null); - t.is(unchanged.lastValidationError, null); -}); - -test('Gemini key test sends key in header and returns safe failure message', async t => { - const { user, workspace } = await createUserWorkspace(t); - await grantUserPlan(t, user.id); - - const fetch = Sinon.stub(t.context.byok as any, 'probeFetch').resolves( - new Response( - 'failed https://generativelanguage.googleapis.com/v1beta/models?key=gemini-secret', - { status: 401 } - ) - ); - t.teardown(() => fetch.restore()); - - const result = await t.context.byok.testConfig({ - workspaceId: workspace.id, - userId: user.id, - provider: ByokProvider.gemini, - storage: ByokKeyStorage.server, - apiKey: 'gemini-secret', - }); - - t.false(result.ok); - t.is( - fetch.firstCall.args[0], - 'https://generativelanguage.googleapis.com/v1beta/models' - ); - t.is( - (fetch.firstCall.args[1]!.headers as Record)[ - 'x-goog-api-key' - ], - 'gemini-secret' - ); - t.deepEqual(fetch.firstCall.args[2]?.allowedHeaders, ['x-goog-api-key']); - t.false(result.message?.includes('gemini-secret')); - t.is(result.message, 'Provider rejected the BYOK key.'); -}); - -test('FAL key test uses read-only platform API probe endpoint', async t => { - const { user, workspace } = await createUserWorkspace(t); - await grantUserPlan(t, user.id); - - const fetch = Sinon.stub(t.context.byok as any, 'probeFetch').resolves( - new Response('{}', { status: 200 }) - ); - t.teardown(() => fetch.restore()); - - const result = await t.context.byok.testConfig({ - workspaceId: workspace.id, - userId: user.id, - provider: ByokProvider.fal, - storage: ByokKeyStorage.server, - apiKey: 'fal-secret', - }); - - t.true(result.ok); - t.is(fetch.firstCall.args[0], 'https://api.fal.ai/v1/models?limit=10'); - t.is( - (fetch.firstCall.args[1]!.headers as Record).Authorization, - 'Key fal-secret' - ); - t.deepEqual(fetch.firstCall.args[2]?.allowedHeaders, ['Authorization']); -}); - -test('provider test failures do not return raw provider response body', async t => { - const { user, workspace } = await createUserWorkspace(t); - await grantUserPlan(t, user.id); - const cases = [ - { - body: 'authorization: Bearer token=a+b%2F==', - status: 401, - message: 'Provider rejected the BYOK key.', - }, - { - body: 'failed https://example.com/models?token=tok+%2F==&limit=1', - status: 403, - message: 'Provider rejected the BYOK key permissions.', - }, - { - body: '{"api_key":"key+value==","accessToken":"tok%2Fvalue"}', - status: 429, - message: 'Provider rate limit exceeded while testing the key.', - }, - { - body: 'Key fal-key+value==', - status: 500, - message: 'Provider service is unavailable.', - }, - ]; - const fetch = Sinon.stub(t.context.byok as any, 'probeFetch'); - for (const [index, matrixCase] of cases.entries()) { - fetch - .onCall(index) - .resolves(new Response(matrixCase.body, { status: matrixCase.status })); - } - t.teardown(() => fetch.restore()); - - for (const matrixCase of cases) { - const result = await t.context.byok.testConfig({ + const reorderConflict = await t.throwsAsync( + t.context.runtime.reorderByokProfiles({ workspaceId: workspace.id, - userId: user.id, - provider: ByokProvider.openai, - storage: ByokKeyStorage.local, - apiKey: 'submitted-secret', - }); - - t.false(result.ok); - t.is(result.message, matrixCase.message); - t.false(result.message?.includes(matrixCase.body)); - } + profiles: [ + { profileId: created.profileId, expectedRevision: rotated.revision }, + { profileId: second.profileId, expectedRevision: second.revision }, + ], + actorUserId: user.id, + }) + ); + t.regex(reorderConflict.message, /byok_revision_conflict/); + t.deepEqual( + (await t.context.runtime.listByokProfiles(workspace.id)).map( + profile => profile.sortOrder + ), + [0, 1] + ); + t.true( + await t.context.runtime.deleteByokProfile(workspace.id, created.profileId) + ); }); -test('dispatch failure disables server BYOK key by provider id', async t => { - const { user, workspace } = await createUserWorkspace(t); - await grantUserPlan(t, user.id); - const key = await t.context.byok.upsertConfig({ - workspaceId: workspace.id, - userId: user.id, - provider: ByokProvider.openai, - storage: ByokKeyStorage.server, - name: 'Primary', - apiKey: 'sk-dispatch-primary', +test('native local lease requires explicit model declarations', async t => { + const user = await t.context.models.user.create({ + email: `${randomUUID()}@affine.pro`, }); - - await t.context.byok.recordProviderFailure({ - workspaceId: workspace.id, - providerId: `byok-${workspaceHash(workspace.id)}-openai-${key.id}`, - featureKind: 'chat', - error: new Error('401 invalid sk-dispatch-primary'), - }); - - const disabled = await t.context.models.copilotWorkspaceByokConfig.get( - key.id - ); - t.truthy(disabled); - if (!disabled) { - return; - } - t.false(disabled.enabled); - t.is(disabled.disabledReason, 'recent_failure'); - t.is(disabled.lastError, 'Provider request failed.'); -}); - -test('dispatch accounting ignores provider ids from another workspace hash', async t => { - const { user, workspace } = await createUserWorkspace(t); - await grantUserPlan(t, user.id); - const otherWorkspace = await t.context.models.workspace.create(user.id); - const key = await t.context.byok.upsertConfig({ - workspaceId: workspace.id, - userId: user.id, - provider: ByokProvider.openai, - storage: ByokKeyStorage.server, - name: 'Primary', - apiKey: 'sk-dispatch-primary', - }); - const mismatchedProviderId = `byok-${workspaceHash(otherWorkspace.id)}-openai-${key.id}`; - - await t.context.byok.recordProviderFailure({ - workspaceId: workspace.id, - providerId: mismatchedProviderId, - featureKind: 'chat', - error: new Error('401 invalid sk-dispatch-primary'), - }); - await t.context.byok.recordUsage({ - workspaceId: workspace.id, - userId: user.id, - providerId: mismatchedProviderId, - featureKind: 'chat', - usage: { total_tokens: 3 }, - }); - - const config = await t.context.models.copilotWorkspaceByokConfig.get(key.id); - t.truthy(config); - t.true(config?.enabled); - t.is(config?.lastError, null); - const usage = await t.context.byok.getUsage( - workspace.id, - new Date(Date.now() - 60_000), - new Date(Date.now() + 60_000) - ); - t.deepEqual(usage, []); -}); - -test('effective profiles use local lease before server keys and skip disabled keys', async t => { - const { user, workspace } = await createUserWorkspace(t); - await grantUserPlan(t, user.id); - const serverKey = await t.context.byok.upsertConfig({ - workspaceId: workspace.id, - userId: user.id, - provider: ByokProvider.openai, - storage: ByokKeyStorage.server, - name: 'Server', - apiKey: 'sk-server', - }); - await t.context.models.copilotWorkspaceByokConfig.markFailure( - workspace.id, - serverKey.id, - 'recent_failure' - ); - await t.context.byok.upsertConfig({ - workspaceId: workspace.id, - userId: user.id, - provider: ByokProvider.gemini, - storage: ByokKeyStorage.server, - name: 'Gemini', - apiKey: 'gemini-server', - }); - const lease = await t.context.byok.createLocalLease({ + const workspace = await t.context.models.workspace.create(user.id); + const lease = await t.context.runtime.createByokLocalLease({ workspaceId: workspace.id, userId: user.id, providers: [ { - provider: ByokProvider.openai, - name: 'Local', - apiKey: 'sk-local', + provider: 'openai', + name: 'Local OpenAI', + credential: 'local-secret', + definition, + enabled: true, }, ], }); - - const profiles = await t.context.byok.getProfiles({ - workspaceId: workspace.id, - userId: user.id, - byokLeaseId: lease.leaseId, - }); - - t.deepEqual( - profiles.map(profile => profile.id.includes('-local-')), - [true, false] - ); - t.deepEqual( - profiles.map(profile => profile.type), - ['openai', 'gemini'] - ); - - const serverOnlyFeatureKinds: ByokFeatureKind[] = [ - 'transcript', - 'embedding', - 'workspace_indexing', - 'rerank', - ]; - for (const featureKind of serverOnlyFeatureKinds) { - const featureProfiles = await t.context.access.getByokProfiles({ - workspaceId: workspace.id, - userId: user.id, - byokLeaseId: lease.leaseId, - featureKind, - }); - t.deepEqual( - featureProfiles.map(profile => profile.type), - ['gemini'] - ); - } -}); - -test('capability warnings match server Gemini background coverage', async t => { - const { user, workspace } = await createUserWorkspace(t); - await grantUserPlan(t, user.id); - - const emptySettings = await t.context.byok.getSettings(workspace.id, user.id); - t.deepEqual( - emptySettings.warnings.map(warning => warning.featureKind), - ['transcript', 'workspace_indexing'] - ); - - await t.context.byok.upsertConfig({ - workspaceId: workspace.id, - userId: user.id, - provider: ByokProvider.gemini, - storage: ByokKeyStorage.server, - name: 'Gemini', - apiKey: 'gemini-server', - }); - - const coveredSettings = await t.context.byok.getSettings( - workspace.id, - user.id - ); - t.deepEqual(coveredSettings.warnings, []); - t.deepEqual(coveredSettings.keys[0].capabilities, [ - 'Text', - 'Image input', - 'Actions', - 'Image generate', - 'Transcript', - 'Indexing', - ]); -}); - -test('usage query only returns byok sources', async t => { - const { user, workspace } = await createUserWorkspace(t); - await t.context.byok.recordUsage({ - workspaceId: workspace.id, - userId: user.id, - providerId: `byok-${workspaceHash(workspace.id)}-openai-server-key1`, - featureKind: 'chat', - model: 'gpt-5-mini', - usage: { prompt_tokens: 1, completion_tokens: 2, total_tokens: 3 }, - }); - await t.context.models.copilotUsage.create({ - workspaceId: workspace.id, - userId: user.id, - provider: 'openai', - providerSource: 'affine_plan', - featureKind: 'chat', - totalTokens: 99, - }); - - const usage = await t.context.byok.getUsage( - workspace.id, - new Date(Date.now() - 60_000), - new Date(Date.now() + 60_000) - ); - - t.is(usage.length, 1); - t.is(usage[0].featureKind, 'chat'); - t.is(usage[0].totalTokens, 3); -}); - -test('usage query aggregates BYOK usage by day and feature in the database', async t => { - const { user, workspace } = await createUserWorkspace(t); - const day = new Date('2026-01-02T08:30:00.000Z'); - await t.context.db.aiUsageEvent.createMany({ - data: [ - { - workspaceId: workspace.id, - userId: user.id, - provider: 'openai', - providerSource: 'byok_server', - featureKind: 'chat', - totalTokens: 3, - createdAt: day, - }, - { - workspaceId: workspace.id, - userId: user.id, - provider: 'openai', - providerSource: 'byok_server', - featureKind: 'chat', - totalTokens: 5, - createdAt: new Date('2026-01-02T20:10:00.000Z'), - }, - { - workspaceId: workspace.id, - userId: user.id, - provider: 'gemini', - providerSource: 'byok_local', - featureKind: 'transcript', - totalTokens: 7, - createdAt: new Date('2026-01-02T21:00:00.000Z'), - }, - { - workspaceId: workspace.id, - userId: user.id, - provider: 'openai', - providerSource: 'affine_plan', - featureKind: 'chat', - totalTokens: 99, - createdAt: day, - }, - ], - }); - - const usage = await t.context.byok.getUsage( - workspace.id, - new Date('2026-01-01T00:00:00.000Z'), - new Date('2026-01-03T00:00:00.000Z') - ); - - t.deepEqual( - usage.map(point => ({ - date: point.date.toISOString(), - featureKind: point.featureKind, - totalTokens: point.totalTokens, - })), - [ - { - date: '2026-01-02T00:00:00.000Z', - featureKind: 'chat', - totalTokens: 8, - }, - { - date: '2026-01-02T00:00:00.000Z', - featureKind: 'transcript', - totalTokens: 7, - }, - ] - ); + t.truthy(lease.leaseId); + t.true(lease.expiresAtMs > Date.now()); }); diff --git a/packages/backend/server/src/__tests__/copilot/capability-runtime.spec.ts b/packages/backend/server/src/__tests__/copilot/capability-runtime.spec.ts new file mode 100644 index 0000000000..c6cf1ba1bb --- /dev/null +++ b/packages/backend/server/src/__tests__/copilot/capability-runtime.spec.ts @@ -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(source: AsyncIterable) { + 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).polluted, undefined); +}); diff --git a/packages/backend/server/src/__tests__/copilot/utils.spec.ts b/packages/backend/server/src/__tests__/copilot/citation-formatter.spec.ts similarity index 100% rename from packages/backend/server/src/__tests__/copilot/utils.spec.ts rename to packages/backend/server/src/__tests__/copilot/citation-formatter.spec.ts diff --git a/packages/backend/server/src/__tests__/copilot/conversation-host.spec.ts b/packages/backend/server/src/__tests__/copilot/conversation-host.spec.ts new file mode 100644 index 0000000000..4aa94587a7 --- /dev/null +++ b/packages/backend/server/src/__tests__/copilot/conversation-host.spec.ts @@ -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(); + const accepted = new Map(); + 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); +}); diff --git a/packages/backend/server/src/__tests__/copilot/copilot-provider.spec.ts b/packages/backend/server/src/__tests__/copilot/copilot-provider.spec.ts index a4efdb98e6..97204aedc4 100644 --- a/packages/backend/server/src/__tests__/copilot/copilot-provider.spec.ts +++ b/packages/backend/server/src/__tests__/copilot/copilot-provider.spec.ts @@ -1,1108 +1,238 @@ import { randomUUID } from 'node:crypto'; -import { Global, Module } from '@nestjs/common'; -import type { Prisma } from '@prisma/client'; -import type { ExecutionContext, TestFn } from 'ava'; +import type { TestFn } from 'ava'; import ava from 'ava'; -import Sinon from 'sinon'; -import { z } from 'zod'; -import { AppModuleBuilder, FunctionalityModules } from '../../app.module'; -import { JobModule, JobQueue } from '../../base'; -import { ServerFeature, ServerService } from '../../core'; -import { AuthModule, AuthService } from '../../core/auth'; -import { QuotaModule } from '../../core/quota'; import { Models } from '../../models'; -import { llmImageDispatchPlan } from '../../native'; -import { CopilotModule } from '../../plugins/copilot'; -import { PromptService } from '../../plugins/copilot/prompt'; -import { - CopilotProviderFactory, - CopilotProviderType, - StreamObject, - StreamObjectSchema, -} from '../../plugins/copilot/providers'; -import { ActionStreamHost } from '../../plugins/copilot/runtime/hosts/action-stream-host'; -import { getProviderRuntimeHost } from '../../plugins/copilot/runtime/provider-runtime-context'; -import { ChatSession, ChatSessionService } from '../../plugins/copilot/session'; +import { CapabilityRuntime } from '../../plugins/copilot/runtime/capability-runtime'; +import { CopilotTranscriptionService } from '../../plugins/copilot/transcript'; import { TranscriptPayloadSchema } from '../../plugins/copilot/transcript/schema'; -import { CopilotTranscriptionService } from '../../plugins/copilot/transcript/service'; -import { TestingPromptService } from '../mocks/prompt-service.mock'; -import { MockJobQueue } from '../mocks/queue.mock'; -import { createTestingModule, TestingModule } from '../utils'; -import { TestAssets } from '../utils/copilot'; +import { createTestingApp, createWorkspace, type TestingApp } from '../utils'; import { - assistantPrompt, - promptMessages, - singleUserPromptMessages, - userPrompt, -} from './prompt-test-helper'; + chatWithActionStream, + chatWithImages, + chatWithStreamObject, + chatWithText, + createCopilotMessage, + createCopilotSession, + sse2array, +} from '../utils/copilot'; -type Tester = { - auth: AuthService; - module: TestingModule; +type Context = { + app: TestingApp; models: Models; - service: ServerService; - prompt: TestingPromptService; - factory: CopilotProviderFactory; - session: ChatSessionService; - actionStreams: ActionStreamHost; + runtime: CapabilityRuntime; transcript: CopilotTranscriptionService; }; -const test = ava as TestFn; +const test = ava.serial as TestFn; +const providerTest = + process.env.AFFINE_TEST_COPILOT_PROVIDER === 'true' ? test : test.skip; -@Global() -@Module({ - providers: [{ provide: JobQueue, useClass: MockJobQueue }], - exports: [JobQueue], -}) -class MockJobModule {} +test.before(async t => { + const app = await createTestingApp(); + t.context = { + app, + models: app.get(Models), + runtime: app.get(CapabilityRuntime), + transcript: app.get(CopilotTranscriptionService), + }; +}); -let isCopilotConfigured = false; -const runIfCopilotConfigured = test.macro( - async ( - t, - callback: (t: ExecutionContext) => Promise | void - ) => { - if (isCopilotConfigured) { - await callback(t); - } else { - t.log('Skip test because copilot is not configured'); - t.pass(); +test.beforeEach(async t => { + await t.context.app.initTestingDB(); +}); + +test.after.always(async t => { + await t.context.app?.close(); +}); + +async function assertManagedRoute( + runtime: CapabilityRuntime, + slot: string, + builtInRouteId?: string +) { + await runtime.assertRoute( + slot, + {}, + { + builtInRouteId, + quotaBackedRoutesAllowed: true, } + ); +} + +providerTest( + 'managed text and object routes satisfy the public SSE contract', + async t => { + const { app, runtime } = t.context; + await assertManagedRoute(runtime, 'chat.default', 'Chat With AFFiNE AI'); + await app.signupV1(); + const workspace = await createWorkspace(app); + + const textSession = await createCopilotSession( + app, + workspace.id, + randomUUID(), + 'Chat With AFFiNE AI' + ); + const textToken = await createCopilotMessage( + app, + textSession, + 'Explain AFFiNE in one sentence.' + ); + t.truthy((await chatWithText(app, textSession, textToken)).trim()); + + const objectSession = await createCopilotSession( + app, + workspace.id, + randomUUID(), + 'Chat With AFFiNE AI' + ); + const objectToken = await createCopilotMessage( + app, + objectSession, + 'Explain AFFiNE in one sentence.' + ); + const events = sse2array( + await chatWithStreamObject(app, objectSession, objectToken) + ); + t.false(events.some(event => event.event === 'error')); + t.true(events.some(event => event.event === 'message')); } ); -test.serial.before(async t => { - const appModule = new AppModuleBuilder() - .use( - ...FunctionalityModules.filter(module => { - const moduleType = 'module' in module ? module.module : module; - return moduleType !== JobModule; - }), - MockJobModule, - AuthModule, - QuotaModule, - CopilotModule - ) - .compile(); - const module = await createTestingModule({ - imports: [appModule], - tapModule: builder => { - builder.overrideProvider(PromptService).useClass(TestingPromptService); - }, - }); +providerTest( + 'managed embedding and rerank routes expose final result shapes', + async t => { + const { app, runtime } = t.context; + await assertManagedRoute(runtime, 'index.embedding'); + await assertManagedRoute(runtime, 'search.rerank'); + const user = await app.signupV1(); + const workspace = await createWorkspace(app); + const options = { user: user.id, workspace: workspace.id }; - const service = module.get(ServerService); - isCopilotConfigured = service.features.includes(ServerFeature.Copilot); - - const auth = module.get(AuthService); - const models = module.get(Models); - const prompt = module.get(PromptService) as TestingPromptService; - const factory = module.get(CopilotProviderFactory); - const session = module.get(ChatSessionService); - const actionStreams = module.get(ActionStreamHost); - const transcript = module.get(CopilotTranscriptionService); - - t.context.module = module; - t.context.auth = auth; - t.context.service = service; - t.context.models = models; - t.context.prompt = prompt; - t.context.factory = factory; - t.context.session = session; - t.context.actionStreams = actionStreams; - t.context.transcript = transcript; -}); - -test.serial.before(async t => { - const { prompt } = t.context; - - prompt.reset(); -}); - -test.after(async t => { - await t.context.module.close(); -}); - -const assertNotWrappedInCodeBlock = ( - t: ExecutionContext, - result: string -) => { - t.assert( - !result.replaceAll('\n', '').trim().startsWith('```') && - !result.replaceAll('\n', '').trim().endsWith('```'), - 'should not wrap in code block' - ); -}; - -const citationChecker = ( - t: ExecutionContext, - citations: { citationNumber: string; citationJson: string }[] -) => { - t.assert(citations.length > 0, 'should have citation'); - for (const { citationJson } of citations) { - t.notThrows(() => { - JSON.parse(citationJson); - }, `should be valid json: ${citationJson}`); - } -}; - -type CitationChecker = typeof citationChecker; - -const assertCitation = ( - t: ExecutionContext, - result: string, - citationCondition: CitationChecker = citationChecker -) => { - const regex = /\[\^(\d+)\]:\s*({.*})/g; - const citations = []; - let match; - while ((match = regex.exec(result)) !== null) { - const citationNumber = match[1]; - const citationJson = match[2]; - citations.push({ citationNumber, citationJson }); - } - citationCondition(t, citations); -}; - -const checkMDList = (text: string) => { - const lines = text.split('\n'); - const listItemRegex = /^( {2})*(-|\u2010-\u2015|\*|\+)? .+$/; - let prevIndent = null; - - for (const line of lines) { - if (line.trim() === '') continue; - if (!listItemRegex.test(line)) { - return false; - } - - const currentIndent = line.match(/^( *)/)?.[0].length!; - if (Number.isNaN(currentIndent) || currentIndent % 2 !== 0) { - return false; - } - - if (prevIndent !== null && currentIndent > 0) { - const indentDiff = currentIndent - prevIndent; - // allow 1 level of indentation difference - if (indentDiff > 2) { - return false; - } - } - - if (line.trim().startsWith('-')) { - prevIndent = currentIndent; - } - } - - return true; -}; - -const checkUrl = (url: string) => { - try { - new URL(url); - return true; - } catch { - return false; - } -}; - -const checkStreamObjects = (result: string) => { - try { - const streamObjects = JSON.parse(result); - z.array(StreamObjectSchema).parse(streamObjects); - return true; - } catch { - return false; - } -}; - -const parseStreamObjects = (result: string): StreamObject[] => { - const streamObjects = JSON.parse(result); - return z.array(StreamObjectSchema).parse(streamObjects); -}; - -const getStreamObjectText = (result: string) => - parseStreamObjects(result) - .filter( - (chunk): chunk is Extract => - chunk.type === 'text-delta' - ) - .map(chunk => chunk.textDelta) - .join(''); - -const retry = async ( - action: string, - t: ExecutionContext, - callback: (t: ExecutionContext) => Promise -) => { - let i = 3; - while (i--) { - const ret = await t.try(async t => { - try { - await callback(t); - } catch (e) { - console.error(`Error during ${action}:`, e); - t.log(`Error during ${action}:`, e); - throw e; - } + const embeddings = await runtime.embed('route-selected', ['AFFiNE'], { + ...options, + featureKind: 'embedding', }); - if (ret.passed) { - return ret.commit(); - } else { - ret.discard({ retainLogs: true }); - t.log(ret.errors.map(e => e.message || e.name || String(e)).join('\n')); - t.log(`retrying ${action} ${3 - i}/3 ...`); - } + t.is(embeddings.length, 1); + t.true(embeddings[0].length > 0); + + const scores = await runtime.rerank( + 'route-selected', + { + query: 'collaborative editor', + candidates: [ + { id: 'relevant', text: 'AFFiNE is a collaborative editor.' }, + { id: 'irrelevant', text: 'A recipe for apple pie.' }, + ], + }, + { ...options, featureKind: 'rerank' } + ); + t.is(scores.length, 2); + t.true(scores.every(score => Number.isFinite(score))); } - t.fail(`failed to run ${action}`); -}; +); -// ==================== utils ==================== +providerTest( + 'managed image and action routes preserve public event contracts', + async t => { + const { app, runtime } = t.context; + await assertManagedRoute(runtime, 'image.generate', 'Generate image'); + await assertManagedRoute( + runtime, + 'action.mindmap.generate', + 'mindmap.generate' + ); + await app.signupV1(); + const workspace = await createWorkspace(app); -test('should validate markdown list', t => { - t.true( - checkMDList(` -- item 1 -- item 2 -`) - ); - t.true( - checkMDList(` -- item 1 - - item 1.1 -- item 2 -`) - ); - t.true( - checkMDList(` -- item 1 - - item 1.1 - - item 1.1.1 -- item 2 -`) - ); - t.true( - checkMDList(` -- item 1 - - item 1.1 - - item 1.1.1 - - item 1.1.2 -- item 2 -`) - ); - t.true( - checkMDList(` -- item 1 - - item 1.1 - - item 1.1.1 -- item 1.2 -`) - ); - t.false( - checkMDList(` -- item 1 - - item 1.1 - - item 1.1.1.1 -`) - ); - t.true( - checkMDList(` -- item 1 - - item 1.1 - - item 1.1.1.1 - item 1.1.1.1 line breaks - - item 1.1.1.2 -`), - 'should allow line breaks' - ); -}); + const imageSession = await createCopilotSession( + app, + workspace.id, + randomUUID(), + 'Generate image' + ); + const imageToken = await createCopilotMessage( + app, + imageSession, + 'A simple panda icon.' + ); + const imageEvents = sse2array( + await chatWithImages(app, imageSession, imageToken) + ); + t.false(imageEvents.some(event => event.event === 'error')); + t.true(imageEvents.some(event => event.event === 'attachment')); -// ==================== action ==================== - -const actions = [ - { - name: 'Should chat with histories', - promptName: ['Chat With AFFiNE AI'], - messages: promptMessages( - userPrompt( - ` -Hi! I’m going to send you a technical term related to real-time collaborative editing (e.g., CRDT, Operational Transformation, OT Composer, etc.). Whenever I send you a term: -1. Translate it into Chinese (send me the Chinese version). -2. Then translate that Chinese back into English (send me the retranslated English). -3. Provide a brief, English-language introduction and context for this concept. -4. In that English explanation, annotate any niche terms with their Chinese equivalents in parentheses (for example: “Conflict-Free Replicated Data Type (无冲突复制数据类型)”). -5. Finally, give the origin or “term history” (e.g., who introduced it, in which paper or year). - -If you understand, please proceed by explaining the term “CRDT.” - `.trim() - ), - assistantPrompt( - ` -1. **Chinese Translation:** -“CRDT” → **无冲突复制数据类型** - -2. **Back-Translation to English:** -无冲突复制数据类型 → **Conflict-Free Replicated Data Type** - -3. **English Introduction & Context:** -A **Conflict-Free Replicated Data Type (无冲突复制数据类型)** is an abstract data type designed for distributed systems where replicas of shared state may be modified concurrently without requiring coordination. CRDTs allow multiple users or processes to update the same data structure (for example, a shared document in a collaborative editor) at the same time. -- **Key Terms (with Chinese equivalents):** - - **Replica (副本):** Each node or client maintains its own copy of the data. - - **State-based (状态型) vs. Operation-based (操作型):** Two main CRDT classes; state-based CRDTs exchange entire state snapshots occasionally, whereas operation-based CRDTs broadcast only incremental operations. - - **Merge Function (合并函数):** A deterministic function that resolves differences between two replicas without conflicts. - -CRDTs enable **eventual consistency (最终一致性)** in real-time collaborative editors by ensuring that, after all updates propagate, every replica converges to the same state, even if operations arrive in different orders. This approach removes the need for a centralized server to resolve conflicts, making offline or peer-to-peer editing possible. - -4. **Origin / Term History:** -The term **“CRDT”** was first introduced by Marc Shapiro, Nuno Preguiça, Carlos Baquero, and Marek Zawirski in their 2011 paper titled “Conflict-free Replicated Data Types” (published in the _Stabilization, Safety, and Security of Distributed Systems (SSS)_ conference). They formalized two families of CRDTs—state-based (“Convergent Replicated Data Types” or CvRDTs) and operation-based (“Commutative Replicated Data Types” or CmRDTs)—and proved their convergence properties under asynchronous, unreliable networks. - `.trim() - ), - userPrompt( - 'Thanks! Now please just tell me the **Chinese translation** and the **back-translated English term** that you provided previously for “CRDT.” Do not reprint the full introduction—only those two lines.' + const actionSession = await createCopilotSession( + app, + workspace.id, + randomUUID(), + 'mindmap.generate' + ); + const actionToken = await createCopilotMessage( + app, + actionSession, + 'AFFiNE product architecture' + ); + const actionEvents = sse2array( + await chatWithActionStream(app, actionSession, { + actionId: 'mindmap.generate', + messageId: actionToken, + }) + ); + t.false(actionEvents.some(event => event.event === 'error')); + t.true( + actionEvents.some( + event => + event.event === 'message' || event.data?.includes('action_done') ) - ), - verifier: (t: ExecutionContext, result: string) => { - assertNotWrappedInCodeBlock(t, result); - const lower = result.toLowerCase(); - t.assert( - lower.includes('无冲突复制数据类型') && - lower.includes('conflict-free replicated data type'), - 'The response should include “无冲突复制数据类型” and “Conflict-Free Replicated Data Type”' - ); - }, - type: 'text' as const, - }, - { - name: 'Should not have citation', - promptName: ['Chat With AFFiNE AI'], - messages: singleUserPromptMessages('what is AFFiNE AI?', { - params: { - files: [ - { - blobId: 'todo_md', - fileName: 'todo.md', - fileType: 'text/markdown', - fileContent: TestAssets.TODO, - }, - ], - }, - }), - verifier: (t: ExecutionContext, result: string) => { - assertNotWrappedInCodeBlock(t, result); - assertCitation(t, result, (t, c) => { - t.assert( - c.length === 0 || - // ignore web search result - c - .map(c => JSON.parse(c.citationJson).type) - .filter(type => ['attachment', 'doc'].includes(type)).length === - 0, - `should not have citation: ${JSON.stringify(c, null, 2)}` - ); - }); - }, - type: 'text' as const, - }, - { - name: 'Should have citation', - promptName: ['Chat With AFFiNE AI'], - messages: singleUserPromptMessages('what is ssot', { - params: { - docs: [ - { - docId: 'SSOT', - docTitle: 'Single source of truth - Wikipedia', - fileType: 'text/markdown', - docContent: TestAssets.SSOT, - }, - ], - }, - }), - verifier: (t: ExecutionContext, result: string) => { - assertNotWrappedInCodeBlock(t, result); - assertCitation(t, result); - }, - type: 'text' as const, - }, - { - name: 'stream objects', - promptName: ['Chat With AFFiNE AI'], - messages: singleUserPromptMessages('what is AFFiNE AI'), - verifier: (t: ExecutionContext, result: string) => { - t.truthy(checkStreamObjects(result), 'should be valid stream objects'); - }, - type: 'object' as const, - }, - { - name: 'Gemini native text', - promptName: ['Chat With AFFiNE AI'], - messages: singleUserPromptMessages( - 'In one short sentence, explain what AFFiNE AI is and mention AFFiNE by name.' - ), - config: { model: 'gemini-3.6-flash' }, - verifier: (t: ExecutionContext, result: string) => { - assertNotWrappedInCodeBlock(t, result); - t.assert( - result.toLowerCase().includes('affine'), - 'should mention AFFiNE' - ); - }, - prefer: CopilotProviderType.Gemini, - type: 'text' as const, - }, - { - name: 'Gemini native stream objects', - promptName: ['Chat With AFFiNE AI'], - messages: singleUserPromptMessages( - 'Respond with one short sentence about AFFiNE AI and mention AFFiNE by name.' - ), - config: { model: 'gemini-3.6-flash' }, - verifier: (t: ExecutionContext, result: string) => { - t.truthy(checkStreamObjects(result), 'should be valid stream objects'); - const assembledText = getStreamObjectText(result); - t.assert( - assembledText.toLowerCase().includes('affine'), - 'should mention AFFiNE' - ); - }, - prefer: CopilotProviderType.Gemini, - type: 'object' as const, - }, - { - promptName: ['Conversation Summary'], - messages: singleUserPromptMessages('', { - params: { - messages: [ - userPrompt('what is single source of truth?'), - assistantPrompt(TestAssets.SSOT), - ], - focus: 'technical decisions', - length: 'comprehensive', - }, - }), - verifier: (t: ExecutionContext, result: string) => { - assertNotWrappedInCodeBlock(t, result); - const cleared = result.toLowerCase(); - t.assert( - cleared.includes('single source of truth') || - /single.*source/.test(cleared) || - cleared.includes('ssot'), - 'should include original keyword' - ); - }, - type: 'text' as const, - }, - { - promptName: [ - 'Summary', - 'Summary as title', - 'Explain this', - 'Write an article about this', - 'Write a twitter about this', - 'Write a poem about this', - 'Write a blog post about this', - 'Write outline', - 'Change tone to', - 'Improve writing for it', - 'Improve grammar for it', - 'Fix spelling for it', - 'Create headings', - 'Make it longer', - 'Make it shorter', - 'Section Edit', - 'Chat With AFFiNE AI', - ], - messages: singleUserPromptMessages(TestAssets.SSOT), - verifier: (t: ExecutionContext, result: string) => { - assertNotWrappedInCodeBlock(t, result); - const cleared = result.toLowerCase(); - t.assert( - cleared.includes('single source of truth') || - /single.*source/.test(cleared) || - cleared.includes('ssot'), - 'should include original keyword' - ); - }, - type: 'text' as const, - }, - { - promptName: ['Continue writing'], - messages: singleUserPromptMessages(TestAssets.AFFiNE), - verifier: (t: ExecutionContext, result: string) => { - assertNotWrappedInCodeBlock(t, result); - t.assert(result.length > 0, 'should not be empty'); - }, - type: 'text' as const, - }, - { - promptName: ['Brainstorm ideas about this', 'Brainstorm mindmap'], - messages: singleUserPromptMessages(TestAssets.AFFiNE), - verifier: (t: ExecutionContext, result: string) => { - assertNotWrappedInCodeBlock(t, result); - t.assert(checkMDList(result), 'should be a markdown list'); - }, - type: 'text' as const, - }, - { - promptName: 'Expand mind map', - messages: singleUserPromptMessages('- Single source of truth'), - verifier: (t: ExecutionContext, result: string) => { - assertNotWrappedInCodeBlock(t, result); - t.assert(checkMDList(result), 'should be a markdown list'); - }, - type: 'text' as const, - }, - { - promptName: 'Find action items from it', - messages: singleUserPromptMessages(TestAssets.TODO), - verifier: (t: ExecutionContext, result: string) => { - assertNotWrappedInCodeBlock(t, result); - t.assert(checkMDList(result), 'should be a markdown list'); - }, - type: 'text' as const, - }, - { - promptName: ['Explain this code', 'Check code error'], - messages: singleUserPromptMessages(TestAssets.Code), - verifier: (t: ExecutionContext, result: string) => { - assertNotWrappedInCodeBlock(t, result); - t.assert( - result.toLowerCase().includes('distance') || - /no.*error/.test(result.toLowerCase()), - 'explain code result should include keyword' - ); - }, - type: 'text' as const, - }, - { - promptName: 'Translate to', - messages: singleUserPromptMessages(TestAssets.SSOT, { - params: { language: 'Simplified Chinese' }, - }), - verifier: (t: ExecutionContext, result: string) => { - assertNotWrappedInCodeBlock(t, result); - const cleared = result.toLowerCase(); - t.assert( - cleared.includes('单一') || cleared.includes('SSOT'), - 'explain code result should include keyword' - ); - }, - type: 'text' as const, - }, - { - promptName: ['Generate a caption', 'Explain this image'], - messages: singleUserPromptMessages('', { - attachments: [ - 'https://cdn.affine.pro/copilot-test/Qgqy9qZT3VGIEuMIotJYoCCH.jpg', - ], - }), - verifier: (t: ExecutionContext, result: string) => { - assertNotWrappedInCodeBlock(t, result); - const content = result.toLowerCase(); - t.assert( - content.includes('classroom') || - content.includes('school') || - content.includes('sky'), - 'explain code result should include keyword' - ); - }, - type: 'text' as const, - }, - { - promptName: ['Convert to sticker', 'Remove background', 'Upscale image'], - messages: singleUserPromptMessages('', { - attachments: [ - 'https://cdn.affine.pro/copilot-test/Zkas098lkjdf-908231.jpg', - ], - }), - verifier: (t: ExecutionContext, link: string) => { - t.truthy(checkUrl(link), 'should be a valid url'); - }, - type: 'image' as const, - }, - { - promptName: ['Generate image'], - messages: singleUserPromptMessages('Panda'), - config: { quality: 'low' }, - verifier: (t: ExecutionContext, link: string) => { - t.truthy(checkUrl(link), 'should be a valid url'); - }, - type: 'image' as const, - }, -]; - -for (const { - name, - promptName, - messages, - verifier, - type, - config, - prefer, -} of actions) { - const prompts = Array.isArray(promptName) ? promptName : [promptName]; - for (const promptName of prompts) { - test( - `should be able to run action: ${promptName}${name ? ` - ${name}` : ''}`, - runIfCopilotConfigured, - async t => { - const { factory, prompt: promptService } = t.context; - const prompt = (await promptService.get(promptName))!; - t.truthy(prompt, 'should have prompt'); - const finalConfig = Object.assign({}, prompt.config, config); - const modelId = - ('model' in finalConfig ? finalConfig.model : undefined) ?? - prompt.model; - const provider = (await factory.getProviderByModel(modelId, { - prefer, - }))!; - t.truthy(provider, 'should have provider'); - await retry(`action: ${promptName}`, t, async t => { - switch (type) { - case 'text': { - const result = await getProviderRuntimeHost(provider).run.text( - { modelId }, - [ - ...promptService.finish( - prompt, - messages.reduce( - (acc, m) => Object.assign(acc, m.params), - {} - ) - ), - ...messages, - ], - finalConfig - ); - t.truthy(result, 'should return result'); - verifier?.(t, result); - break; - } - case 'object': { - const streamObjects: StreamObject[] = []; - for await (const chunk of getProviderRuntimeHost( - provider - ).run.streamObject( - { modelId }, - [ - ...promptService.finish( - prompt, - messages.reduce( - (acc, m) => Object.assign(acc, (m as any).params || {}), - {} - ) - ), - ...messages, - ], - finalConfig - )) { - streamObjects.push(chunk); - } - t.truthy(streamObjects, 'should return result'); - verifier?.(t, JSON.stringify(streamObjects)); - break; - } - case 'image': { - const finalMessage = [...messages]; - const params = {}; - if (finalMessage.length === 1) { - const latestMessage = finalMessage.pop()!; - Object.assign(params, { - content: latestMessage.content, - attachments: - 'attachments' in latestMessage - ? latestMessage.attachments - : undefined, - }); - } - const imageMessages = [ - ...promptService.finish( - prompt, - finalMessage.reduce( - (acc, m) => Object.assign(acc, m.params), - params - ) - ), - ...finalMessage, - ]; - const prepared = await getProviderRuntimeHost( - provider - ).prepare.image({ modelId }, imageMessages, finalConfig); - t.truthy(prepared, 'should prepare image request'); - const result = await llmImageDispatchPlan({ - preparedRoutes: [ - { - provider_id: prepared!.route.providerId, - protocol: prepared!.route.protocol, - model: prepared!.route.model, - config: prepared!.route.backendConfig, - request: prepared!.request, - }, - ], - }); - - t.truthy(result.response.images.length, 'should return result'); - for (const image of result.response.images) { - const link = image.data_base64 - ? `data:${image.media_type};base64,${image.data_base64}` - : image.url; - t.truthy(link); - verifier?.(t, link!); - } - break; - } - default: { - t.fail('unsupported provider type'); - break; - } - } - }); - } ); } -} +); -// ==================== action recipes ==================== - -function actionRunRecord( - input: Parameters[0] -) { - return { - id: `action-run-${randomUUID()}`, - userId: input.userId, - workspaceId: input.workspaceId, - docId: input.docId ?? null, - sessionId: input.sessionId ?? null, - userMessageId: input.userMessageId ?? null, - compatSubmissionId: input.compatSubmissionId ?? null, - assistantMessageId: null, - actionId: input.actionId, - actionVersion: input.actionVersion, - status: 'created' as const, - attempt: input.attempt ?? 1, - retryOf: input.retryOf ?? null, - inputSnapshot: (input.inputSnapshot ?? null) as Prisma.JsonValue, - result: null, - artifacts: null, - resultSummary: null, - errorCode: null, - trace: null, - createdAt: new Date(), - updatedAt: new Date(), - }; -} - -async function installActionSessionMock( - t: ExecutionContext, - { - actionId, - actionPrompt, - content, - }: { - actionId: string; - actionPrompt: Awaited>; - content: string; - } -) { - const { models, session } = t.context; - const sandbox = Sinon.createSandbox(); - const sessionId = `copilot-provider-action-${actionId}-${randomUUID()}`; - const user = await models.user.create({ - email: `copilot-provider-user-${randomUUID()}@affine.test`, - }); - const userId = user.id; - const workspace = await models.workspace.create(userId); - const workspaceId = workspace.id; - const docId = `copilot-provider-action-${actionId}-doc`; - const savedTurns: Array<{ role: string }> = []; - const userTurn = { - conversationId: sessionId, - role: 'user' as const, - content, - attachments: [], - renderTrace: [], - toolEvents: [], - metadata: { language: 'English' }, - createdAt: new Date(), - }; - const chatSession = new ChatSession( - { - userId, - sessionId, - workspaceId, - docId, - turns: [userTurn], - prompt: actionPrompt!, - }, - (prompt, turns, params, maxTokenSize, sessionId) => - t.context.prompt.renderSession( - prompt, - turns, - params, - maxTokenSize, - sessionId - ), - async state => { - savedTurns.push(...state.turns); - } - ); - - sandbox - .stub(session, 'get') - .callsFake(async id => (id === sessionId ? chatSession : null)); - sandbox.stub(session, 'appendTurn').callsFake(async input => { - savedTurns.push(input.turn); - return { ...input.turn, id: `assistant-${randomUUID()}` }; - }); - sandbox.stub(session, 'revertLatestMessage').resolves(); - sandbox - .stub(models.copilotActionRun, 'create') - .callsFake(async input => actionRunRecord(input)); - sandbox.stub(models.copilotActionRun, 'markRunning').callsFake( - async id => - ({ - id, - status: 'running', - }) as never - ); - sandbox.stub(models.copilotActionRun, 'complete').callsFake( - async (id, input) => - ({ - id, - ...input, - updatedAt: new Date(), - }) as never - ); - - return { sandbox, sessionId, userId, savedTurns }; -} - -const actionRecipeCases = [ - { - actionId: 'mindmap.generate', - content: 'apple company', - verifier: (t: ExecutionContext, result: string) => { - assertNotWrappedInCodeBlock(t, result); - t.assert(checkMDList(result), 'should be a markdown list'); - }, - }, - { - actionId: 'slides.outline', - content: 'apple company', - verifier: (t: ExecutionContext, result: string) => { - assertNotWrappedInCodeBlock(t, result); - t.assert( - result - .split('\n') - .filter(line => line.trim()) - .every(line => /^( {2})*(-|\*|\+) .+$/.test(line)), - 'should be a markdown list' - ); - t.false( - result - .split('\n') - .filter(line => line.trim()) - .every(line => { - try { - JSON.parse(line); - return true; - } catch { - return false; - } - }), - 'should not expose raw NDJSON' - ); - }, - }, -]; - -for (const { actionId, content, verifier } of actionRecipeCases) { - test.serial( - `should be able to run action recipe: ${actionId}`, - runIfCopilotConfigured, - async t => { - await retry(`action recipe: ${actionId}`, t, async t => { - const { actionStreams, prompt } = t.context; - const actionPrompt = await prompt.get(actionId); - if (!actionPrompt) { - return t.fail(`prompt ${actionId} should exist`); - } - - const { sandbox, sessionId, userId, savedTurns } = - await installActionSessionMock(t, { - actionId, - actionPrompt, - content, - }); - - let result = ''; - try { - const prepared = await actionStreams.stream(userId, sessionId, { - actionId, - actionVersion: 'v1', - modelId: actionPrompt.model, - }); - - for await (const event of prepared.stream) { - if (event.type === 'action_done' && event.status === 'succeeded') { - if (typeof event.result === 'string') { - result += event.result; - } else if (event.result && typeof event.result === 'object') { - const value = event.result as { - content?: unknown; - text?: unknown; - result?: unknown; - }; - result += - typeof value.content === 'string' - ? value.content - : typeof value.text === 'string' - ? value.text - : typeof value.result === 'string' - ? value.result - : ''; - } - } - } - } finally { - sandbox.restore(); - } - t.truthy(result, 'should return result'); - verifier(t, result); - t.true( - savedTurns.some(turn => turn.role === 'assistant'), - 'should persist assistant turn through real conversation host' - ); - }); - } - ); -} - -const TRANSCRIPT_AUDIO_CASES = [ - { - name: 'short audio', - url: 'https://cdn.affine.pro/copilot-test/MP9qDGuYgnY+ILoEAmHpp3h9Npuw2403EAYMEA.mp3', - mimeType: 'audio/mpeg', - modelId: 'gemini-3.5-flash-lite', - }, - { - name: 'middle audio', - url: 'https://cdn.affine.pro/copilot-test/2ed05eo1KvZ2tWB_BAjFo67EAPZZY-w4LylUAw.m4a', - mimeType: 'audio/m4a', - modelId: 'gemini-3.5-flash-lite', - }, - { - name: 'long audio', - url: 'https://cdn.affine.pro/copilot-test/nC9-e7P85PPI2rU29QWwf8slBNRMy92teLIIMw.opus', - mimeType: 'audio/opus', - modelId: 'gemini-3.6-flash', - }, -]; - -for (const testCase of TRANSCRIPT_AUDIO_CASES) { - test( - `should run transcript task through native action bridge: ${testCase.name}`, - runIfCopilotConfigured, - async t => { - const { models, transcript } = t.context; - const user = await models.user.create({ - email: `copilot-provider-transcript-${randomUUID()}@affine.pro`, - }); - const workspace = await models.workspace.create(user.id); - const blobId = `copilot-provider-transcript-blob-${randomUUID()}`; - const payload = TranscriptPayloadSchema.parse({ - sourceAudio: { blobId, mimeType: testCase.mimeType }, - infos: [ - { - url: testCase.url, - mimeType: testCase.mimeType, - index: 0, - }, - ], - }); - const task = await models.copilotTranscriptTask.create({ - userId: user.id, - workspaceId: workspace.id, - blobId, - strategy: 'gemini', - recipeId: 'transcript.audio.gemini', - recipeVersion: 'v1', - inputSnapshot: payload, - publicMeta: { - sourceAudio: payload.sourceAudio, - infos: payload.infos, - }, - }); - - await retry('transcript native action recipe', t, async t => { - await transcript.transcriptTask({ - taskId: task.id, - payload, - modelId: testCase.modelId, - }); - const ready = await models.copilotTranscriptTask.get(task.id); - t.is(ready?.status, 'ready'); - const parsed = TranscriptPayloadSchema.parse(ready?.protectedResult); - t.is(typeof parsed.normalizedTranscript, 'string'); - }); - } - ); -} - -// ==================== rerank ==================== - -test( - 'should be able to rerank message chunks', - runIfCopilotConfigured, +providerTest( + 'managed transcript route executes the provider-neutral job port', async t => { - const { factory } = t.context; - - await retry('rerank', t, async t => { - const query = 'Is this content relevant to programming?'; - const embeddings = [ - 'How to write JavaScript code for web development.', - 'Today is a beautiful sunny day for walking in the park.', - 'Python is a popular programming language for data science.', - 'The weather forecast predicts rain for the weekend.', - 'JavaScript frameworks like React and Angular are widely used.', - 'Cooking recipes can be found in many online blogs.', - 'Machine learning algorithms are essential for AI development.', - 'The latest smartphone models have impressive camera features.', - 'Learning to code can open up many career opportunities.', - 'The stock market is experiencing significant fluctuations.', - ]; - - const provider = (await factory.getProviderByModel('gpt-4o-mini'))!; - t.assert(provider, 'should have provider for rerank'); - - const scores = await getProviderRuntimeHost(provider).run.rerank( - { modelId: 'gpt-4o-mini' }, - { - query, - candidates: embeddings.map((text, index) => ({ - id: String(index), - text, - })), - } - ); - - t.is(scores.length, 10, 'should return scores for all chunks'); - - for (const score of scores) { - t.assert( - typeof score === 'number' && score >= 0 && score <= 1, - `score should be a number between 0 and 1, got ${score}` - ); - } - - t.log('Rerank scores:', scores); - t.is( - scores.filter(s => s > 0.5).length, - 4, - 'should have 4 related chunks' - ); + const { models, runtime, transcript } = t.context; + await assertManagedRoute( + runtime, + 'transcript.audio', + 'Transcript audio structured' + ); + const user = await models.user.create({ + email: `copilot-provider-transcript-${randomUUID()}@affine.pro`, }); + const workspace = await models.workspace.create(user.id); + const blobId = `copilot-provider-transcript-${randomUUID()}`; + const payload = TranscriptPayloadSchema.parse({ + sourceAudio: { blobId, mimeType: 'audio/mpeg' }, + infos: [ + { + url: 'https://cdn.affine.pro/copilot-test/MP9qDGuYgnY+ILoEAmHpp3h9Npuw2403EAYMEA.mp3', + mimeType: 'audio/mpeg', + index: 0, + }, + ], + }); + const task = await models.copilotTranscriptTask.create({ + userId: user.id, + workspaceId: workspace.id, + blobId, + recipeId: 'transcript.audio', + recipeVersion: 'v1', + inputSnapshot: payload, + publicMeta: { sourceAudio: payload.sourceAudio, infos: payload.infos }, + }); + + await transcript.transcriptTask({ taskId: task.id, payload }); + const ready = await models.copilotTranscriptTask.get(task.id); + t.is(ready?.status, 'ready'); + t.is( + typeof TranscriptPayloadSchema.parse(ready?.protectedResult) + .normalizedTranscript, + 'string' + ); } ); diff --git a/packages/backend/server/src/__tests__/copilot/copilot.e2e.ts b/packages/backend/server/src/__tests__/copilot/copilot.e2e.ts index ec68f29e5d..85d676d7d7 100644 --- a/packages/backend/server/src/__tests__/copilot/copilot.e2e.ts +++ b/packages/backend/server/src/__tests__/copilot/copilot.e2e.ts @@ -1,2251 +1,317 @@ +import '../../plugins/copilot'; + import { randomUUID } from 'node:crypto'; -import serverNativeModule from '@affine/server-native'; -import { ProjectRoot } from '@affine-tools/utils/path'; -import { PrismaClient } from '@prisma/client'; +import { McpAccessMode, PrismaClient } from '@prisma/client'; import type { TestFn } from 'ava'; import ava from 'ava'; -import Sinon from 'sinon'; -import { AppModule } from '../../app.module'; -import { JobQueue } from '../../base'; -import { ConfigModule } from '../../base/config'; +import { Config } from '../../base'; +import { ServerFeature, ServerService } from '../../core'; import { AuthService } from '../../core/auth'; -import { DocReader } from '../../core/doc'; -import { QuotaService } from '../../core/quota'; -import { ContextCategories, DocRole, WorkspaceRole } from '../../models'; -import { CompatSubmissionStore } from '../../plugins/copilot/compat/submission-store'; -import { CopilotContextService } from '../../plugins/copilot/context'; import { - CopilotEmbeddingJob, - MockEmbeddingClient, -} from '../../plugins/copilot/embedding'; -import { PromptService } from '../../plugins/copilot/prompt'; -import { - CopilotProviderFactory, - CopilotProviderType, - GeminiGenerativeProvider, - OpenAIProvider, -} from '../../plugins/copilot/providers'; -import { CapabilityRuntime } from '../../plugins/copilot/runtime/capability-runtime'; -import { ChatSessionService } from '../../plugins/copilot/session'; -import { CopilotStorage } from '../../plugins/copilot/storage'; -import { - installMockCopilotRuntime, - MockCopilotProvider, - Mockers, -} from '../mocks'; -import { TestingPromptService } from '../mocks/prompt-service.mock'; -import { - acceptInviteById, - createTestingApp, - createWorkspace, - inviteUser, - smallestPng, - TestingApp, - TestUser, -} from '../utils'; + ContextCategories, + DocRole, + Models, + WorkspaceMemberStatus, + WorkspaceRole, +} from '../../models'; +import { CopilotFeatureService } from '../../plugins/copilot/feature'; +import { McpCredentialService } from '../../plugins/copilot/mcp/credential'; +import { WorkspaceMcpProvider } from '../../plugins/copilot/mcp/provider'; +import { installMockCopilotRuntime, Mockers } from '../mocks'; +import { createTestingApp, createWorkspace, type TestingApp } from '../utils'; import { addContextCategory, - addContextDoc, addContextFile, - array2sse, - chatWithActionStream, chatWithImages, - chatWithStreamObject, chatWithText, - chatWithTextStream, - cleanObject, createCopilotContext, createCopilotMessage, createCopilotSession, - createDocCopilotSession, - createPinnedCopilotSession, - createWorkspaceCopilotSession, - forkCopilotSession, getCopilotSession, - getDocSessions, getHistories, - getPinnedSessions, - getTranscriptTask, - getWorkspaceSessions, - listContext, - listContextCategories, - listContextDocAndFiles, - matchFiles, - matchWorkspaceDocs, - settleTranscriptTask, sse2array, - submitTranscriptTask, - textToEventStream, - unsplashSearch, - updateCopilotSession, } from '../utils/copilot'; -const test = ava as TestFn<{ - auth: AuthService; +type Context = { app: TestingApp; - db: PrismaClient; - context: CopilotContextService; - jobs: CopilotEmbeddingJob; - prompt: TestingPromptService; - factory: CopilotProviderFactory; - storage: CopilotStorage; - u1: TestUser; -}>; -let restoreMockCopilotRuntime: (() => void) | undefined; - -const waitForStatus = async ( - loadStatus: () => Promise, - expected: string, - description: string, - attempts = 30, - intervalMs = 1000 -) => { - let status = await loadStatus(); - for (let attempt = 0; attempt < attempts; attempt++) { - if (status === expected) { - return status; - } - await new Promise(resolve => setTimeout(resolve, intervalMs)); - status = await loadStatus(); - } - throw new Error( - `${description} did not reach status "${expected}", last status: ${ - status ?? 'undefined' - }` - ); + restoreRuntime: () => void; }; +const test = ava.serial as TestFn; + test.before(async t => { - restoreMockCopilotRuntime = installMockCopilotRuntime(); - const app = await createTestingApp({ - imports: [ - ConfigModule.override({ - copilot: { - providers: { - openai: { apiKey: '1' }, - fal: {}, - gemini: { apiKey: '1' }, - }, - unsplash: { - key: process.env.UNSPLASH_ACCESS_KEY || '1', - }, - }, - }), - AppModule, - ], - tapModule: m => { - // use real JobQueue for testing - m.overrideProvider(JobQueue).useClass(JobQueue); - m.overrideProvider(DocReader).useValue({ - getFullDocContent() { - return { - title: '1', - summary: '1', - }; - }, - getWorkspaceContent() { - return {}; - }, - }); - m.overrideProvider(PromptService).useClass(TestingPromptService); - m.overrideProvider(OpenAIProvider).useClass(MockCopilotProvider); - m.overrideProvider(GeminiGenerativeProvider).useClass( - class MockGenerativeProvider extends MockCopilotProvider { - // @ts-expect-error type not typed - override type: CopilotProviderType = CopilotProviderType.Gemini; - } - ); - }, - }); - - const auth = app.get(AuthService); - const db = app.get(PrismaClient); - const context = app.get(CopilotContextService); - const prompt = app.get(PromptService) as TestingPromptService; - const storage = app.get(CopilotStorage); - const jobs = app.get(CopilotEmbeddingJob); - - t.context.app = app; - t.context.db = db; - t.context.auth = auth; - t.context.context = context; - t.context.prompt = prompt; - t.context.storage = storage; - t.context.jobs = jobs; + const restoreRuntime = installMockCopilotRuntime(); + t.context = { + app: await createTestingApp(), + restoreRuntime, + }; }); -let textPromptName = 'prompt'; -let imagePromptName = 'prompt-image'; - test.beforeEach(async t => { - Sinon.restore(); - const { app, prompt } = t.context; - await app.initTestingDB(); - prompt.reset(); - t.context.u1 = await app.signupV1(); - textPromptName = randomUUID().replaceAll('-', ''); - imagePromptName = randomUUID().replaceAll('-', ''); - - await prompt.set(textPromptName, 'test', [ - { role: 'system', content: 'hello {{word}}' }, - ]); - - await prompt.set(imagePromptName, 'test-image', [ - { role: 'system', content: 'hello {{word}}' }, - ]); + await t.context.app.initTestingDB(); }); test.after.always(async t => { - restoreMockCopilotRuntime?.(); - await t.context.app.close(); + t.context.restoreRuntime?.(); + await t.context.app?.close(); }); -// ==================== session ==================== - -test('should create session correctly', async t => { - const { app, u1 } = t.context; - - const assertCreateSession = async ( - workspaceId: string, - error: string, - asserter = async (x: any) => { - t.truthy(await x, error); - } - ) => { - await asserter( - createCopilotSession(app, workspaceId, randomUUID(), textPromptName) - ); - }; - - { - const { id } = await createWorkspace(app); - await assertCreateSession( - id, - 'should be able to create session with cloud workspace that user can access' - ); - } - - { - await assertCreateSession( - randomUUID(), - 'should be able to create session with local workspace' - ); - } - - { - const u2 = await app.createUser(); - const { id } = await createWorkspace(app); - await app.login(u2); - await assertCreateSession(id, '', async x => { - await t.throwsAsync( - x, - { instanceOf: Error }, - 'should not able to create session with cloud workspace that user cannot access' - ); - }); - - await app.switchUser(u1); - const inviteId = await inviteUser(app, id, u2.email); - await app.login(u2); - await acceptInviteById(app, id, inviteId, false); - await assertCreateSession( - id, - 'should able to create session after user have permission' - ); - } -}); - -test('should update session correctly', async t => { +test('disabled copilot hides its server feature and rejects every API transport', async t => { const { app } = t.context; + const config = app.get(Config); + const feature = app.get(CopilotFeatureService); + const server = app.get(ServerService); + await app.signupV1(); + const workspace = await createWorkspace(app); - const assertUpdateSession = async ( - sessionId: string, - error: string, - asserter = async (x: any) => { - t.truthy(await x, error); - } - ) => { - await asserter(updateCopilotSession(app, sessionId, textPromptName)); - }; - - { - const { id: workspaceId } = await createWorkspace(app); - const docId = randomUUID(); - const sessionId = await createCopilotSession( - app, - workspaceId, - docId, - textPromptName + config.copilot.enabled = false; + feature.onConfigChanged({ updates: { copilot: { enabled: false } } }); + try { + t.false(server.features.includes(ServerFeature.Copilot)); + await t.throwsAsync( + createCopilotSession( + app, + workspace.id, + randomUUID(), + 'Chat With AFFiNE AI' + ) ); - await assertUpdateSession( - sessionId, - 'should be able to update session with cloud workspace that user can access' - ); - } - - { - const sessionId = await createCopilotSession( - app, - randomUUID(), - randomUUID(), - textPromptName - ); - await assertUpdateSession( - sessionId, - 'should be able to update session with local workspace' - ); - } - - { - await app.signupV1(); - const u2 = await app.createUser(); - const { id: workspaceId } = await createWorkspace(app); - const inviteId = await inviteUser(app, workspaceId, u2.email); - await app.login(u2); - await acceptInviteById(app, workspaceId, inviteId, false); - const sessionId = await createCopilotSession( - app, - workspaceId, - randomUUID(), - textPromptName - ); - await assertUpdateSession( - sessionId, - 'should able to update session after user have permission' - ); - } - - { - const sessionId = '123456'; - await assertUpdateSession(sessionId, '', async x => { - await t.throwsAsync( - x, - { instanceOf: Error }, - 'should not able to update invalid session id' - ); - }); + await app.GET('/api/copilot/unsplash/photos').expect(403); + await app + .POST(`/api/workspaces/${workspace.id}/mcp`) + .send({ jsonrpc: '2.0', id: 1, method: 'initialize', params: {} }) + .expect(403); + } finally { + config.copilot.enabled = true; + feature.onConfigChanged({ updates: { copilot: { enabled: true } } }); } }); -test('should fetch action session by session id', async t => { +test('session, compat message, text SSE and durable history share one public contract', async t => { const { app } = t.context; - const { id: workspaceId } = await createWorkspace(app); + await app.signupV1(); + const workspace = await createWorkspace(app); + const docId = randomUUID(); const sessionId = await createCopilotSession( app, - workspaceId, + workspace.id, + docId, + 'Chat With AFFiNE AI' + ); + + t.deepEqual(await getCopilotSession(app, workspace.id, sessionId), { + id: sessionId, + docId, + parentSessionId: null, + pinned: false, + promptName: 'Chat With AFFiNE AI', + }); + + const token = await createCopilotMessage(app, sessionId, 'hello'); + const [beforeStream] = await getHistories(app, { + workspaceId: workspace.id, + docId, + }); + t.is(beforeStream.sessionId, sessionId); + t.deepEqual(beforeStream.messages, []); + + t.is( + await chatWithText(app, sessionId, token), + 'generate text to text stream' + ); + t.is( + await chatWithText(app, sessionId, token), + 'generate text to text stream' + ); + const [history] = await getHistories(app, { + workspaceId: workspace.id, + docId, + }); + t.deepEqual( + history.messages.map(message => [message.role, message.content]), + [ + ['user', 'hello'], + ['assistant', 'generate text to text stream'], + ['assistant', 'generate text to text stream'], + ] + ); + t.is(history.messages.filter(message => message.role === 'user').length, 1); + t.not(history.messages[0].id, token); +}); + +test('chat and history endpoints reject a different user', async t => { + const { app } = t.context; + const owner = await app.signupV1(); + const workspace = await createWorkspace(app); + const sessionId = await createCopilotSession( + app, + workspace.id, + randomUUID(), + 'Chat With AFFiNE AI' + ); + const token = await createCopilotMessage(app, sessionId, 'private'); + await app.signupV1(); + + await t.throwsAsync(chatWithText(app, sessionId, token)); + await t.throwsAsync(getHistories(app, { workspaceId: workspace.id })); + + await app.switchUser(owner); + t.is( + await chatWithText(app, sessionId, token), + 'generate text to text stream' + ); +}); + +test('image SSE emits persisted attachment events for action sessions', async t => { + const { app } = t.context; + await app.signupV1(); + const workspace = await createWorkspace(app); + const sessionId = await createCopilotSession( + app, + workspace.id, randomUUID(), 'Generate image' ); + const token = await createCopilotMessage(app, sessionId, 'Panda'); - const session = await getCopilotSession(app, workspaceId, sessionId); - t.truthy(session); - t.is(session.id, sessionId); - t.is(session.promptName, 'Generate image'); + const events = sse2array(await chatWithImages(app, sessionId, token)); + const attachment = events.find(event => event.event === 'attachment'); + t.truthy(attachment?.data); }); -test('should fork session correctly', async t => { - const { app, u1 } = t.context; - - const assertForkSession = async ( - workspaceId: string, - docId: string, - sessionId: string, - lastMessageId: string | undefined, - error: string, - asserter = async (x: any) => { - const forkedSessionId = await x; - t.truthy(forkedSessionId, error); - return forkedSessionId; - } - ) => - await asserter( - forkCopilotSession(app, workspaceId, docId, sessionId, lastMessageId) - ); - - // prepare session - const { id } = await createWorkspace(app); - const docId = randomUUID(); - const sessionId = await createCopilotSession(app, id, docId, textPromptName); - - let forkedSessionId: string; - // should be able to fork session - { - for (let i = 0; i < 3; i++) { - const messageId = await createCopilotMessage(app, sessionId); - await chatWithText(app, sessionId, messageId); - } - const histories = await getHistories(app, { workspaceId: id, docId }); - const latestMessageId = histories[0].messages.findLast( - m => m.role === 'assistant' - )?.id; - t.truthy(latestMessageId, 'should find last message id'); - - // should be able to fork session - forkedSessionId = await assertForkSession( - id, - docId, - sessionId, - latestMessageId!, - 'should be able to fork session with cloud workspace that user can access' - ); - } - - // should be able to fork session without latestMessageId (copy all messages) - { - forkedSessionId = await assertForkSession( - id, - docId, - sessionId, - undefined, - 'should be able to fork session without latestMessageId' - ); - } - - // should not be able to fork session with wrong latestMessageId - { - await assertForkSession( - id, - docId, - sessionId, - 'wrong-message-id', - '', - async x => { - await t.throwsAsync( - x, - { instanceOf: Error }, - 'should not able to fork session with wrong latestMessageId' - ); - } - ); - } - - { - const u2 = await app.signupV1(); - await assertForkSession(id, docId, sessionId, randomUUID(), '', async x => { - await t.throwsAsync( - x, - { instanceOf: Error }, - 'should not able to fork session with cloud workspace that user cannot access' - ); - }); - - await app.switchUser(u1); - const inviteId = await inviteUser(app, id, u2.email); - await app.switchUser(u2); - await acceptInviteById(app, id, inviteId, false); - await assertForkSession(id, docId, sessionId, randomUUID(), '', async x => { - await t.throwsAsync( - x, - { instanceOf: Error }, - 'should not able to fork a root session from other user' - ); - }); - - await app.switchUser(u1); - const histories = await getHistories(app, { workspaceId: id, docId }); - const latestMessageId = histories - .find(h => h.sessionId === forkedSessionId) - ?.messages.findLast(m => m.role === 'assistant')?.id; - t.truthy(latestMessageId, 'should find latest message id'); - - await app.switchUser(u2); - await assertForkSession( - id, - docId, - forkedSessionId, - latestMessageId!, - 'should able to fork a forked session created by other user' - ); - } -}); - -test('should be able to use test provider', async t => { +test('context API rechecks write access and filters unreadable category docs', async t => { const { app } = t.context; - - const { id } = await createWorkspace(app); - t.truthy( - await createCopilotSession(app, id, randomUUID(), textPromptName), - 'failed to create session' - ); -}); - -// ==================== message ==================== - -test('should create message correctly', async t => { - const { app } = t.context; - const pngData = await fetch(smallestPng).then(res => res.arrayBuffer()); - const cases = [ - { - title: 'should be able to create message with valid session', - invoke: (sessionId: string) => createCopilotMessage(app, sessionId), - }, - { - title: 'should be able to create message with url link', - invoke: (sessionId: string) => - createCopilotMessage(app, sessionId, undefined, [ - 'http://example.com/cat.jpg', - ]), - }, - { - title: 'should be able to create message with blob', - invoke: (sessionId: string) => - createCopilotMessage( - app, - sessionId, - undefined, - undefined, - new File([new Uint8Array(pngData)], '1.png', { type: 'image/png' }) - ), - }, - { - title: 'should be able to create message with blobs', - invoke: (sessionId: string) => - createCopilotMessage(app, sessionId, undefined, undefined, undefined, [ - new File([new Uint8Array(pngData)], '1.png', { type: 'image/png' }), - ]), - }, - ]; - - for (const testCase of cases) { - const { id } = await createWorkspace(app); - const sessionId = await createCopilotSession( - app, - id, - randomUUID(), - textPromptName - ); - const messageId = await testCase.invoke(sessionId); - t.truthy(messageId, testCase.title); - } - - { - await t.throwsAsync( - createCopilotMessage(app, randomUUID()), - { instanceOf: Error }, - 'should not able to create message with invalid session' - ); - } -}); - -// ==================== chat ==================== - -test('should be able to chat with api', async t => { - const { app, storage } = t.context; - - Sinon.stub(storage, 'handleRemoteLink').resolvesArg(2); - - const { id } = await createWorkspace(app); - { - const docId = randomUUID(); - const sessionId = await createCopilotSession( - app, - id, - docId, - textPromptName - ); - const messageId = await createCopilotMessage(app, sessionId); - const ret = await chatWithText(app, sessionId, messageId); - t.is( - ret, - 'generate text to text stream', - 'should be able to chat with text' - ); - - const ret2 = await chatWithTextStream(app, sessionId, messageId); - t.is( - ret2, - textToEventStream('generate text to text stream', messageId), - 'should be able to chat with text stream' - ); - - const [history] = await getHistories(app, { workspaceId: id, docId }); - const persistedMessageIds = history?.messages - .filter(message => message.role !== 'system') - .map(message => message.id); - t.deepEqual( - persistedMessageIds?.every(id => typeof id === 'string' && id.length > 0), - true, - 'should persist non-empty database-generated ids for chat turns' - ); - t.is( - new Set(persistedMessageIds).size, - persistedMessageIds?.length ?? 0, - 'should persist unique ids for chat turns' - ); - } - - { - const sessionId = await createCopilotSession( - app, - id, - randomUUID(), - imagePromptName - ); - const messageId = await createCopilotMessage(app, sessionId); - const ret3 = await chatWithImages(app, sessionId, messageId); - t.is( - array2sse(sse2array(ret3).filter(e => e.event !== 'event')), - textToEventStream( - ['https://example.com/gpt-image-2.jpg'], - messageId, - 'attachment' - ), - 'should be able to chat with images' - ); - } - - { - const sessionId = await createCopilotSession( - app, - id, - randomUUID(), - textPromptName - ); - const messageId = await createCopilotMessage(app, sessionId); - - const ret4 = await chatWithStreamObject(app, sessionId, messageId); - - const objects = Array.from('generate text to text stream').map(data => - JSON.stringify({ type: 'text-delta', textDelta: data }) - ); - - t.is( - ret4, - textToEventStream(objects, messageId), - 'should be able to chat with stream object' - ); - } - - Sinon.restore(); -}); - -test('should be able to chat with api by action stream', async t => { - const { app, db, prompt } = t.context; - - const { id } = await createWorkspace(app); - const beforeQuota = await app.gql( - ` - query getCopilotQuota($workspaceId: String!) { - currentUser { - copilot(workspaceId: $workspaceId) { - quota { - used - } - } - } - } - `, - { workspaceId: id } - ); - const sessionId = await createCopilotSession( - app, - id, - randomUUID(), - 'slides.outline' - ); - const messageId = await createCopilotMessage(app, sessionId, 'apple company'); - const actionPrompt = await prompt.get('slides.outline'); - t.truthy(actionPrompt); - const ret = await chatWithActionStream(app, sessionId, { - actionId: 'slides.outline', - actionVersion: 'v1', - modelId: actionPrompt?.model, - messageId, - }); - t.is( - array2sse(sse2array(ret).filter(e => e.event !== 'event')), - textToEventStream(['generate text to text stream'], messageId), - 'should be able to chat with action stream' - ); - const actionRuns = await db.aiActionRun.findMany({ - where: { sessionId }, - select: { - actionId: true, - actionVersion: true, - status: true, - assistantMessageId: true, - }, - }); - const afterQuota = await app.gql( - ` - query getCopilotQuota($workspaceId: String!) { - currentUser { - copilot(workspaceId: $workspaceId) { - quota { - used - } - } - } - } - `, - { workspaceId: id } - ); - - t.like(actionRuns[0], { - actionId: 'slides.outline', - actionVersion: 'v1', - status: 'succeeded', - }); - t.truthy(actionRuns[0]?.assistantMessageId); - t.is( - afterQuota.currentUser.copilot.quota.used, - beforeQuota.currentUser.copilot.quota.used + 1 - ); -}); - -test('should map action stream preparation errors to SSE error events', async t => { - const { app } = t.context; - - const { id } = await createWorkspace(app); - const sessionId = await createCopilotSession( - app, - id, - randomUUID(), - 'slides.outline' - ); - const messageId = await createCopilotMessage(app, sessionId, 'apple company'); - - const ret = await chatWithActionStream(app, sessionId, { - actionId: 'image.filter.unknown', - actionVersion: 'v1', - messageId, - }); - - t.true(ret.includes('error')); -}); - -test('should be able to chat with special image model', async t => { - const { app, storage } = t.context; - - Sinon.stub(storage, 'handleRemoteLink').resolvesArg(2); - - const { id } = await createWorkspace(app); - - const testWithModel = async (promptName: string, finalPrompt: string) => { - const sessionId = await createCopilotSession( - app, - id, - randomUUID(), - promptName - ); - const messageId = await createCopilotMessage(app, sessionId, 'some-tag', [ - smallestPng, - ]); - const ret3 = await chatWithImages(app, sessionId, messageId); - t.is( - ret3, - textToEventStream( - [ - 'https://example.com/gpt-image-2.jpg', - `https://example.com/generated/${encodeURIComponent(finalPrompt)}.jpg`, - ], - messageId, - 'attachment' - ), - 'should be able to chat with images' - ); - }; - - await testWithModel('Generate image', 'some-tag'); - await testWithModel( - 'Convert to sticker', - 'convert this image to sticker. you need to identify the subject matter and warp a circle of white stroke around the subject matter and with transparent background. some-tag' - ); - await testWithModel( - 'Upscale image', - 'make the image more detailed. some-tag' - ); - await testWithModel( - 'Remove background', - 'Keep the subject and remove other non-subject items. Transparent background. some-tag' - ); - - Sinon.restore(); -}); - -test('should be able to retry with api', async t => { - const { app, storage } = t.context; - - Sinon.stub(storage, 'handleRemoteLink').resolvesArg(2); - - // normal chat - { - const { id } = await createWorkspace(app); - const docId = randomUUID(); - const sessionId = await createCopilotSession( - app, - id, - docId, - textPromptName - ); - const messageId = await createCopilotMessage(app, sessionId); - // chat 2 times - await chatWithText(app, sessionId, messageId); - await chatWithText(app, sessionId, messageId); - - const histories = await getHistories(app, { workspaceId: id, docId }); - t.deepEqual( - histories.map(h => h.messages.map(m => m.content)), - [['generate text to text stream', 'generate text to text stream']], - 'should be able to list history' - ); - } - - // retry chat - { - const { id } = await createWorkspace(app); - const docId = randomUUID(); - const sessionId = await createCopilotSession( - app, - id, - docId, - textPromptName - ); - const messageId = await createCopilotMessage(app, sessionId); - await chatWithText(app, sessionId, messageId); - // retry without message id - await chatWithText(app, sessionId); - - // should only have 1 message - const histories = await getHistories(app, { workspaceId: id, docId }); - t.snapshot( - cleanObject(histories), - 'should be able to list history after retry' - ); - } - - // retry chat with new message id - { - const { id } = await createWorkspace(app); - const docId = randomUUID(); - const sessionId = await createCopilotSession( - app, - id, - docId, - textPromptName - ); - const messageId = await createCopilotMessage(app, sessionId); - await chatWithText(app, sessionId, messageId); - // retry with new message id - const newMessageId = await createCopilotMessage(app, sessionId); - await chatWithText(app, sessionId, newMessageId, '', true); - - // should only have 1 message - const histories = await getHistories(app, { workspaceId: id, docId }); - t.snapshot( - cleanObject(histories), - 'should be able to list history after retry' - ); - } - - Sinon.restore(); -}); - -test('should reject message from different session', async t => { - const { app } = t.context; - - const { id } = await createWorkspace(app); - const sessionId = await createCopilotSession( - app, - id, - randomUUID(), - textPromptName - ); - const anotherSessionId = await createCopilotSession( - app, - id, - randomUUID(), - textPromptName - ); - const anotherMessageId = await createCopilotMessage(app, anotherSessionId); - await t.throwsAsync( - chatWithText(app, sessionId, anotherMessageId), - { instanceOf: Error }, - 'should reject message from different session' - ); -}); - -test('should reject request from different user', async t => { - const { app, u1 } = t.context; - - const u2 = await app.createUser(); - const { id } = await createWorkspace(app); - const sessionId = await createCopilotSession( - app, - id, - randomUUID(), - textPromptName - ); - - // should reject message from different user - { - await app.login(u2); - await t.throwsAsync( - createCopilotMessage(app, sessionId), - { instanceOf: Error }, - 'should reject message from different user' - ); - } - - // should reject chat from different user - { - await app.switchUser(u1); - const messageId = await createCopilotMessage(app, sessionId); - { - await app.switchUser(u2); - await t.throwsAsync( - chatWithText(app, sessionId, messageId), - { instanceOf: Error }, - 'should reject chat from different user' - ); - } - } -}); - -// ==================== history ==================== - -test('should be able to list history', async t => { - const { app } = t.context; - - const { id: workspaceId } = await createWorkspace(app); - const docId = randomUUID(); - const sessionId = await createCopilotSession( - app, - workspaceId, - docId, - textPromptName - ); - - const messageId = await createCopilotMessage(app, sessionId, 'hello'); - await chatWithText(app, sessionId, messageId); - - { - const histories = await getHistories(app, { workspaceId, docId }); - t.deepEqual( - histories.map(h => h.messages.map(m => m.content)), - [['hello', 'generate text to text stream']], - 'should be able to list history' - ); - } - - { - const histories = await getHistories(app, { - workspaceId, - docId, - options: { messageOrder: 'desc' }, - }); - t.deepEqual( - histories.map(h => h.messages.map(m => m.content)), - [['generate text to text stream', 'hello']], - 'should be able to list history' - ); - } -}); - -test('should preserve persisted assistant render trace on history reload', async t => { - const { app } = t.context; - const chatRuntime = app.get(CapabilityRuntime); - Sinon.stub(chatRuntime, 'streamObject').callsFake(async function* () { - yield { type: 'reasoning', textDelta: 'Inspecting context' } as const; - yield { - type: 'tool-result', - toolCallId: 'call_1', - toolName: 'doc_read', - args: { docId: 'doc-1' }, - result: { markdown: '# AFFiNE' }, - } as const; - yield { type: 'text-delta', textDelta: 'Final ' } as const; - yield { type: 'text-delta', textDelta: 'answer' } as const; - }); - - const { id: workspaceId } = await createWorkspace(app); - const docId = randomUUID(); - const sessionId = await createCopilotSession( - app, - workspaceId, - docId, - textPromptName - ); - - const messageToken = await createCopilotMessage(app, sessionId, 'hello'); - await chatWithStreamObject(app, sessionId, messageToken); - - const histories = await app.gql( - ` - query getCopilotHistoriesWithTrace( - $workspaceId: String! - $docId: String - $options: QueryChatHistoriesInput - ) { - currentUser { - copilot(workspaceId: $workspaceId) { - histories(docId: $docId, options: $options) { - sessionId - messages { - role - content - streamObjects { - type - textDelta - toolCallId - toolName - args - result - } - } - } - } - } - } - `, - { - workspaceId, - docId, - options: { withMessages: true }, - } - ); - - const assistantMessage = - histories.currentUser.copilot.histories[0]?.messages.find( - (message: { role: string }) => message.role === 'assistant' - ); - - t.is(assistantMessage?.content, 'Final answer'); - t.deepEqual(assistantMessage?.streamObjects, [ - { - type: 'reasoning', - textDelta: 'Inspecting context', - toolCallId: null, - toolName: null, - args: null, - result: null, - }, - { - type: 'tool-result', - toolCallId: 'call_1', - toolName: 'doc_read', - args: { docId: 'doc-1' }, - result: { markdown: '# AFFiNE' }, - textDelta: null, - }, - { - type: 'text-delta', - textDelta: 'Final answer', - toolCallId: null, - toolName: null, - args: null, - result: null, - }, - ]); -}); - -test('should keep compat submission token out of durable history before stream', async t => { - const { app } = t.context; - - const { id: workspaceId } = await createWorkspace(app); - const docId = randomUUID(); - const sessionId = await createCopilotSession( - app, - workspaceId, - docId, - textPromptName - ); - - const messageToken = await createCopilotMessage(app, sessionId, 'hello'); - const histories = await getHistories(app, { workspaceId, docId }); - - t.deepEqual( - histories.flatMap(history => - history.messages.map(message => message.content) - ), - [], - 'should not persist user turn before stream starts' - ); - - await chatWithText(app, sessionId, messageToken); - const [history] = await getHistories(app, { workspaceId, docId }); - - t.truthy(history?.messages[0]?.id); - t.not( - history?.messages[0]?.id, - messageToken, - 'should return compat token instead of durable turn id' - ); -}); - -test('should accept compat submission once and keep duplicate consume idempotent', async t => { - const { app } = t.context; - - const { id: workspaceId } = await createWorkspace(app); - const docId = randomUUID(); - const sessionId = await createCopilotSession( - app, - workspaceId, - docId, - textPromptName - ); - - const beforeQuota = await app.gql( - ` - query getCopilotQuota($workspaceId: String!) { - currentUser { - copilot(workspaceId: $workspaceId) { - quota { - used - } - } - } - } - `, - { workspaceId } - ); - - const messageToken = await createCopilotMessage(app, sessionId, 'hello'); - const text = await chatWithText(app, sessionId, messageToken); - t.is(text, 'generate text to text stream'); - await chatWithText(app, sessionId, messageToken); - - const afterQuota = await app.gql( - ` - query getCopilotQuota($workspaceId: String!) { - currentUser { - copilot(workspaceId: $workspaceId) { - quota { - used - } - } - } - } - `, - { workspaceId } - ); - const [history] = await getHistories(app, { workspaceId, docId }); - - t.is( - afterQuota.currentUser.copilot.quota.used, - beforeQuota.currentUser.copilot.quota.used + 1, - 'should count accepted submission exactly once' - ); - t.true((history?.tokens ?? 0) > 0, 'should accumulate token cost'); - t.deepEqual( - history?.messages.map(message => message.content), - ['hello', 'generate text to text stream', 'generate text to text stream'] - ); - t.is( - history?.messages.filter(message => message.role === 'user').length, - 1, - 'should reuse the same durable user turn for duplicate consume' - ); - t.not( - history?.messages.find(message => message.role === 'user')?.id, - messageToken, - 'should keep compat token separate from durable user turn id' - ); -}); - -test('should allow accepted token replay after quota is exhausted', async t => { - const { app } = t.context; - const quota = app.get(QuotaService); - Sinon.stub(quota, 'getUserQuota').resolves({ - copilotActionLimit: 1, - } as never); - - const { id: workspaceId } = await createWorkspace(app); - const docId = randomUUID(); - const sessionId = await createCopilotSession( - app, - workspaceId, - docId, - textPromptName - ); - - const messageToken = await createCopilotMessage(app, sessionId, 'hello'); - t.is( - await chatWithText(app, sessionId, messageToken), - 'generate text to text stream' - ); - t.is( - await chatWithText(app, sessionId, messageToken), - 'generate text to text stream' - ); - - const [history] = await getHistories(app, { workspaceId, docId }); - t.is( - history?.messages.filter(message => message.role === 'user').length, - 1, - 'should not insert a second user turn when replaying an accepted token' - ); -}); - -test('should recover duplicate consume after accepted-cache write fails', async t => { - const { app } = t.context; - const submissions = app.get(CompatSubmissionStore); - let shouldFail = true; - Sinon.stub(submissions, 'markAccepted').callsFake(async (...args) => { - if (shouldFail) { - shouldFail = false; - throw new Error('inject accepted cache failure'); - } - return await CompatSubmissionStore.prototype.markAccepted.apply( - submissions, - args - ); - }); - - const { id: workspaceId } = await createWorkspace(app); - const docId = randomUUID(); - const sessionId = await createCopilotSession( - app, - workspaceId, - docId, - textPromptName - ); - - const messageToken = await createCopilotMessage(app, sessionId, 'hello'); - await t.throwsAsync(chatWithText(app, sessionId, messageToken), { - instanceOf: Error, - }); - - t.is( - await chatWithText(app, sessionId, messageToken), - 'generate text to text stream' - ); - - const [history] = await getHistories(app, { workspaceId, docId }); - t.is( - history?.messages.filter(message => message.role === 'user').length, - 1, - 'should reuse the durable user turn after accepted-cache failure' - ); - t.deepEqual( - history?.messages.map(message => message.content), - ['hello', 'generate text to text stream'] - ); -}); - -test('should retry token safely when durable insert failed before commit', async t => { - const { app } = t.context; - const sessions = app.get(ChatSessionService); - let shouldFail = true; - Sinon.stub(sessions, 'appendTurn').callsFake(async (...args) => { - if (shouldFail) { - shouldFail = false; - throw new Error('inject append failure'); - } - return await ChatSessionService.prototype.appendTurn.apply(sessions, args); - }); - - const { id: workspaceId } = await createWorkspace(app); - const docId = randomUUID(); - const sessionId = await createCopilotSession( - app, - workspaceId, - docId, - textPromptName - ); - - const messageToken = await createCopilotMessage(app, sessionId, 'hello'); - await t.throwsAsync(chatWithText(app, sessionId, messageToken), { - instanceOf: Error, - }); - - t.is( - await chatWithText(app, sessionId, messageToken), - 'generate text to text stream' - ); - - const [history] = await getHistories(app, { workspaceId, docId }); - t.is( - history?.messages.filter(message => message.role === 'user').length, - 1, - 'should insert the user turn exactly once after retry' - ); - t.deepEqual( - history?.messages.map(message => message.content), - ['hello', 'generate text to text stream'] - ); -}); - -test('should reject new token before durable insert when quota is exhausted', async t => { - const { app } = t.context; - const quota = app.get(QuotaService); - Sinon.stub(quota, 'getUserQuota').resolves({ - copilotActionLimit: 1, - } as never); - - const { id: workspaceId } = await createWorkspace(app); - const docId = randomUUID(); - const sessionId = await createCopilotSession( - app, - workspaceId, - docId, - textPromptName - ); - - const firstMessageToken = await createCopilotMessage(app, sessionId, 'hello'); - await chatWithText(app, sessionId, firstMessageToken); - - const secondMessageToken = await createCopilotMessage( - app, - sessionId, - 'new action' - ); - await t.throwsAsync(chatWithText(app, sessionId, secondMessageToken), { - instanceOf: Error, - }); - - const [history] = await getHistories(app, { workspaceId, docId }); - t.deepEqual( - history?.messages.map(message => message.content), - ['hello', 'generate text to text stream'] - ); -}); - -test('should preload prompt messages when withPrompt is enabled', async t => { - const { app, prompt } = t.context; - - const promptName = randomUUID().replaceAll('-', ''); - await prompt.set(promptName, 'test', [ - { role: 'system', content: 'system prompt' }, - { role: 'user', content: 'preloaded question' }, - ]); - - const { id: workspaceId } = await createWorkspace(app); - const docId = randomUUID(); - const sessionId = await createCopilotSession( - app, - workspaceId, - docId, - promptName - ); - const messageId = await createCopilotMessage(app, sessionId, 'hello'); - await chatWithText(app, sessionId, messageId); - - const withoutPrompt = await getHistories(app, { - workspaceId, - docId, - options: { withPrompt: false }, - }); - const withPrompt = await getHistories(app, { - workspaceId, - docId, - options: { withPrompt: true }, - }); - const chatsWithPrompt = await app.gql( - ` - query getCopilotChatsWithPrompt( - $workspaceId: String! - $docId: String! - $pagination: PaginationInput! - $options: QueryChatHistoriesInput - ) { - currentUser { - copilot(workspaceId: $workspaceId) { - chats(pagination: $pagination, docId: $docId, options: $options) { - totalCount - edges { - node { - sessionId - messages { - content - } - } - } - } - } - } - } - `, - { - workspaceId, - docId, - pagination: { first: 10, offset: 0 }, - options: { withMessages: true, withPrompt: true }, - } - ); - - t.deepEqual( - withoutPrompt[0]?.messages.map(message => message.content), - ['hello', 'generate text to text stream'] - ); - t.deepEqual( - withPrompt[0]?.messages.map(message => message.content), - ['preloaded question', 'hello', 'generate text to text stream'] - ); - t.deepEqual( - chatsWithPrompt.currentUser.copilot.chats.edges[0]?.node.messages.map( - (message: { content: string }) => message.content - ), - ['preloaded question', 'hello', 'generate text to text stream'] - ); -}); - -test('should keep action sessions visible in session and chat metadata queries', async t => { - const { app } = t.context; - - const { id: workspaceId } = await createWorkspace(app); - const sessionId = await createCopilotSession( - app, - workspaceId, - randomUUID(), - 'Generate image' - ); - - const sessionsResult = await app.gql( - ` - query getCopilotSessions($workspaceId: String!) { - currentUser { - copilot(workspaceId: $workspaceId) { - sessions { - id - promptName - } - } - } - } - `, - { workspaceId } - ); - const chatsResult = await app.gql( - ` - query getCopilotChats($workspaceId: String!, $pagination: PaginationInput!) { - currentUser { - copilot(workspaceId: $workspaceId) { - chats(pagination: $pagination) { - edges { - node { - sessionId - promptName - action - messages { - content - } - } - } - } - } - } - } - `, - { - workspaceId, - pagination: { first: 10, offset: 0 }, - } - ); - - t.true( - sessionsResult.currentUser.copilot.sessions.some( - (session: { id: string; promptName: string }) => - session.id === sessionId && session.promptName === 'Generate image' - ), - 'should expose action session in sessions()' - ); - t.true( - chatsResult.currentUser.copilot.chats.edges.some( - (edge: { - node: { - sessionId: string; - promptName: string; - messages: { content: string }[]; - }; - }) => - edge.node.sessionId === sessionId && - edge.node.promptName === 'Generate image' && - edge.node.messages.length === 0 - ), - 'should expose action session metadata in chats(withMessages: false)' - ); -}); - -test('should reject request that user have not permission', async t => { - const { app, u1 } = t.context; - - const u2 = await app.createUser(); - const { id: workspaceId } = await createWorkspace(app); - - // should reject request that user have not permission - { - await app.login(u2); - await t.throwsAsync( - getHistories(app, { workspaceId }), - { instanceOf: Error }, - 'should reject request that user have not permission' - ); - } - - // should able to list history after user have permission - { - await app.switchUser(u1); - const inviteId = await inviteUser(app, workspaceId, u2.email); - await app.switchUser(u2); - await acceptInviteById(app, workspaceId, inviteId, false); - - t.deepEqual( - await getHistories(app, { workspaceId }), - [], - 'should able to list history after user have permission' - ); - } - - { - const docId = randomUUID(); - const sessionId = await createCopilotSession( - app, - workspaceId, - docId, - textPromptName - ); - - const messageId = await createCopilotMessage(app, sessionId); - await chatWithText(app, sessionId, messageId); - - const histories = await getHistories(app, { workspaceId, docId }); - t.deepEqual( - histories.map(h => h.messages.map(m => m.content)), - [['generate text to text stream']], - 'should able to list history' - ); - - await app.switchUser(u1); - t.deepEqual( - await getHistories(app, { workspaceId }), - [], - 'should not list history created by another user' - ); - } -}); - -test('should be able to search image from unsplash', async t => { - const { app } = t.context; - - const resp = await unsplashSearch(app); - t.not(resp.status, 404, 'route should be exists'); -}); - -test('should be able to manage context', async t => { - const { app, context, jobs } = t.context; - const waitForMatches = async ( - loader: () => Promise, - expectedLength = 1 - ) => { - let matches = await loader(); - for (let attempt = 0; attempt < 30; attempt++) { - if ((matches?.length ?? 0) >= expectedLength) { - return matches; - } - await new Promise(resolve => setTimeout(resolve, 1000)); - matches = await loader(); - } - return matches; - }; - - const { id: workspaceId } = await createWorkspace(app); - const sessionId = await createCopilotSession( - app, - workspaceId, - randomUUID(), - textPromptName - ); - - // use mocked embedding client - Sinon.stub(context, 'embeddingClient').get(() => new MockEmbeddingClient()); - Sinon.stub(jobs, 'embeddingClient').get(() => new MockEmbeddingClient()); - - { - await t.throwsAsync( - createCopilotContext(app, workspaceId, randomUUID()), - { instanceOf: Error }, - 'should throw error if create context with invalid session id' - ); - - const context = await createCopilotContext(app, workspaceId, sessionId); - - const list = await listContext(app, workspaceId, sessionId); - t.deepEqual( - list.map(f => ({ id: f.id })), - [{ id: context }], - 'should list context' - ); - } - - const fs = await import('node:fs'); - const buffer = fs.readFileSync( - ProjectRoot.join('packages/common/native/fixtures/sample.pdf').toFileUrl() - ); - - // match files - { - const contextId = await createCopilotContext(app, workspaceId, sessionId); - - const { id: fileId } = await addContextFile( - app, - contextId, - 'sample.pdf', - buffer - ); - - await waitForStatus( - async () => - (await listContextDocAndFiles(app, workspaceId, sessionId, contextId)) - ?.files?.[0]?.status, - 'finished', - 'context file embedding', - 60 - ); - - const { files } = - (await listContextDocAndFiles(app, workspaceId, sessionId, contextId)) || - {}; - t.deepEqual(cleanObject(files, ['id', 'error', 'createdAt']), [ - { - blobId: 'Ip3vuwzubwJnOlzeKQ0Gc-daDcMc7EOYnIqypOyn4bs', - chunkSize: 1, - name: 'sample.pdf', - status: 'finished', - }, - ]); - - const result = await waitForMatches( - () => matchFiles(app, contextId, 'test', 1), - 1 - ); - if (!result) { - t.fail('should return context matches'); - return; - } - t.is(result.length, 1, 'should match context'); - t.is(result[0].fileId, fileId, 'should match file id'); - } - - // match docs - { - const sessionId = await createCopilotSession( - app, - workspaceId, - randomUUID(), - textPromptName - ); - const contextId = await createCopilotContext(app, workspaceId, sessionId); - - const docId = 'docId1'; - await t.context.db.snapshot.create({ - data: { - workspaceId: workspaceId, - id: docId, - blob: Buffer.from([1, 1]), - state: Buffer.from([1, 1]), - updatedAt: new Date(), - createdAt: new Date(), - }, - }); - - await addContextDoc(app, contextId, docId); - - await waitForStatus( - async () => - (await listContextDocAndFiles(app, workspaceId, sessionId, contextId)) - ?.docs?.[0]?.status ?? undefined, - 'finished', - 'context doc embedding', - 60 - ); - - const { docs } = - (await listContextDocAndFiles(app, workspaceId, sessionId, contextId)) || - {}; - t.deepEqual(cleanObject(docs, ['error', 'createdAt']), [ - { - id: docId, - status: 'finished', - }, - ]); - - const result = await waitForMatches( - () => matchWorkspaceDocs(app, contextId, 'test', 1), - 1 - ); - if (!result) { - t.fail('should return workspace doc matches'); - return; - } - t.is(result.length, 1, 'should match context'); - t.is(result[0].docId, docId, 'should match doc id'); - } -}); - -test('should reject context reads from another user', async t => { - const { app, context, jobs, u1 } = t.context; - - const u2 = await app.signupV1(); - await app.switchUser(u1); - - const { id: workspaceId } = await createWorkspace(app); - const sessionId = await createCopilotSession( - app, - workspaceId, - randomUUID(), - textPromptName - ); - - Sinon.stub(context, 'embeddingClient').get(() => new MockEmbeddingClient()); - Sinon.stub(jobs, 'embeddingClient').get(() => new MockEmbeddingClient()); - - const contextId = await createCopilotContext(app, workspaceId, sessionId); - await addContextFile(app, contextId, 'sample.txt', Buffer.from('test file')); - - await app.switchUser(u2); - - await t.throwsAsync( - app.gql(` - query { - currentUser { - copilot { - contexts(contextId: "${contextId}") { - id - } - } - } - } - `) - ); - await t.throwsAsync(matchFiles(app, contextId, 'test', 1)); -}); - -test('should skip unauthorized docs when adding context category', async t => { - const { app, context, jobs, u1 } = t.context; - + const models = app.get(Models); + const owner = await app.signupV1(); + const workspace = await createWorkspace(app); const member = await app.signupV1(); - await app.switchUser(u1); + await models.workspaceUser.set( + workspace.id, + member.id, + WorkspaceRole.Collaborator, + { status: WorkspaceMemberStatus.Accepted } + ); - const { id: workspaceId } = await createWorkspace(app); - await app.create(Mockers.WorkspaceUser, { - workspaceId, - userId: member.id, - type: WorkspaceRole.Collaborator, - }); + const sessionId = await createCopilotSession( + app, + workspace.id, + randomUUID(), + 'Chat With AFFiNE AI' + ); + const contextId = await createCopilotContext(app, workspace.id, sessionId); + await models.workspaceUser.set( + workspace.id, + member.id, + WorkspaceRole.External + ); + await t.throwsAsync( + addContextFile(app, contextId, 'sample.txt', Buffer.from('test')) + ); - const readableSnapshot = await app.create(Mockers.DocSnapshot, { - workspaceId, - user: u1, + await models.workspaceUser.set( + workspace.id, + member.id, + WorkspaceRole.Collaborator, + { status: WorkspaceMemberStatus.Accepted } + ); + const readable = await app.create(Mockers.DocSnapshot, { + workspaceId: workspace.id, + user: owner, }); - const hiddenSnapshot = await app.create(Mockers.DocSnapshot, { - workspaceId, - user: u1, - }); - - await app.create(Mockers.DocMeta, { - workspaceId, - docId: readableSnapshot.id, - title: 'readable-doc', + const hidden = await app.create(Mockers.DocSnapshot, { + workspaceId: workspace.id, + user: owner, }); await app.create(Mockers.DocMeta, { - workspaceId, - docId: hiddenSnapshot.id, - title: 'hidden-doc', + workspaceId: workspace.id, + docId: readable.id, + title: 'readable', + }); + await app.create(Mockers.DocMeta, { + workspaceId: workspace.id, + docId: hidden.id, + title: 'hidden', defaultRole: DocRole.None, }); - - Sinon.stub(context, 'embeddingClient').get(() => new MockEmbeddingClient()); - Sinon.stub(jobs, 'embeddingClient').get(() => new MockEmbeddingClient()); - - await app.switchUser(member); - const sessionId = await createCopilotSession( - app, - workspaceId, - randomUUID(), - textPromptName - ); - const contextId = await createCopilotContext(app, workspaceId, sessionId); const category = await addContextCategory( app, contextId, ContextCategories.Collection, - 'fav', - [readableSnapshot.id, hiddenSnapshot.id] + 'favorites', + [readable.id, hidden.id] ); - t.deepEqual( category.docs.map(doc => doc.id), - [readableSnapshot.id] - ); - - const ret = await listContextCategories( - app, - workspaceId, - sessionId, - contextId - ); - t.deepEqual( - ret?.collections?.[0]?.docs.map(doc => doc.id), - [readableSnapshot.id] + [readable.id] ); }); -test('should be able to transcript', async t => { - const { app, db } = t.context; - - const { id: workspaceId } = await createWorkspace(app); - const transcriptOutput = [ - { a: 'A', s: 30, e: 45, t: 'Hello, everyone.' }, - { - a: 'B', - s: 46, - e: 70, - t: 'Hi, thank you for joining the meeting today.', - }, - ]; - const summaryOutput = { - title: 'Weekly Sync', - durationMinutes: 12, - attendees: ['A', 'B'], - keyPoints: ['Reviewed launch status'], - actionItems: [ - { - description: 'Send recap', - owner: 'A', - deadline: 'Friday', - }, - ], - decisions: ['Ship on Monday'], - openQuestions: ['Need final QA sign-off'], - blockers: ['Waiting on analytics'], - }; - const formatTime = (seconds: number) => { - const hours = Math.floor(seconds / 3600); - const minutes = Math.floor((seconds % 3600) / 60); - const secs = Math.floor(seconds % 60); - return [hours, minutes, secs] - .map(value => value.toString().padStart(2, '0')) - .join(':'); - }; - const buildTranscriptActionResult = ( - route: { - provider_id?: string; - request?: { - messages?: Array<{ content?: string | Array<{ text?: string }> }>; - }; - }, - model: string, - metadataFallback: { - sourceAudio?: unknown; - quality?: unknown; - infos?: unknown; - sliceManifest?: Array<{ startSec?: number }> | null; - } = {} - ) => { - const getContentText = (content?: string | Array<{ text?: string }>) => - typeof content === 'string' - ? content - : content?.map(item => item.text ?? '').join(''); - const metadataContent = route.request?.messages - ?.map(message => getContentText(message.content)) - .find(content => content?.startsWith('{')); - const metadata = { - ...metadataFallback, - ...(metadataContent ? JSON.parse(metadataContent) : {}), - }; - const sliceManifest: Array<{ startSec?: number }> = metadata.sliceManifest - ?.length - ? metadata.sliceManifest - : [{ startSec: 0 }]; - const normalizedSegments = sliceManifest.flatMap( - (slice: { startSec?: number }) => - transcriptOutput.map(segment => { - const startSec = (slice.startSec ?? 0) + segment.s; - const endSec = (slice.startSec ?? 0) + segment.e; - const speaker = segment.a; - return { - speaker, - start: formatTime(startSec), - end: formatTime(endSec), - startSec, - endSec, - text: segment.t, - }; - }) - ); - return { - sourceAudio: metadata.sourceAudio ?? null, - quality: metadata.quality ?? null, - infos: metadata.infos ?? null, - sliceManifest: metadata.sliceManifest ?? null, - normalizedSegments, - normalizedTranscript: normalizedSegments - .map(segment => `${segment.start} ${segment.speaker}: ${segment.text}`) - .join('\n'), - summaryJson: summaryOutput, - providerMeta: { - provider: 'gemini', - model, - }, - }; - }; - const originalActionPreparedStream = (serverNativeModule as any) - .runNativeActionRecipePreparedStream; - (serverNativeModule as any).runNativeActionRecipePreparedStream = ( - input: { - recipeId: string; - recipeVersion?: string; - input?: { - sourceAudio?: unknown; - quality?: unknown; - infos?: unknown; - sliceManifest?: Array<{ startSec?: number }> | null; - preparedRoutes?: { - transcribe?: Array<{ - provider_id?: string; - request?: { - messages?: Array<{ - content?: string | Array<{ text?: string }>; - }>; - }; - }>; - }; - }; - }, - callback: (error: Error | null, eventJson: string) => void - ) => { - if (!input.recipeId.startsWith('transcript.audio.')) { - return originalActionPreparedStream(input, callback); - } - - const route = input.input?.preparedRoutes?.transcribe?.[0] ?? {}; - const result = buildTranscriptActionResult( - route, - 'gemini-3.5-flash-lite', - input.input ?? {} - ); - const actionVersion = input.recipeVersion ?? 'v1'; - const events = [ - { - type: 'action_start', - actionId: input.recipeId, - actionVersion, - status: 'running', - }, - { - type: 'step_start', - actionId: input.recipeId, - actionVersion, - stepId: 'transcribe', - status: 'running', - }, - { - type: 'step_end', - actionId: input.recipeId, - actionVersion, - stepId: 'transcribe', - status: 'running', - }, - { - type: 'action_done', - actionId: input.recipeId, - actionVersion, - status: 'succeeded', - result, - trace: { - actionId: input.recipeId, - actionVersion, - 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, '__AFFINE_LLM_STREAM_END__'); - return { abort() {} }; - }; - t.teardown(() => { - (serverNativeModule as any).runNativeActionRecipePreparedStream = - originalActionPreparedStream; +test('MCP credentials remain endpoint-bound through rotate, revoke and expiry', async t => { + const { app } = t.context; + const auth = app.get(AuthService); + const credentials = app.get(McpCredentialService); + const db = app.get(PrismaClient); + const models = app.get(Models); + const provider = app.get(WorkspaceMcpProvider); + const user = await auth.signUp(`mcp-${randomUUID()}@affine.pro`, '123456'); + const target = await models.workspace.create(user.id); + const other = await models.workspace.create(user.id); + const issued = await credentials.create({ + userId: user.id, + workspaceId: target.id, + name: 'Claude Desktop', + accessMode: McpAccessMode.READ_ONLY, + expirationDays: 90, }); - { - const job = await submitTranscriptTask( - app, - workspaceId, - '1', - '1.mp3', - [Buffer.from([1, 1])], - { - sourceAudio: { - mimeType: 'audio/ogg', - durationMs: 120000, - sampleRate: 48000, - channels: 2, - }, - quality: { - degraded: true, - overflowCount: 4, - }, - sliceManifest: [ - { - index: 0, - fileName: '1-0.opus', - mimeType: 'audio/opus', - startSec: 12, - durationSec: 58, - byteSize: 2, - }, - ], - } - ); - t.truthy(job.id, 'should have job id'); + t.like(await credentials.authenticate(issued.token, target.id), { + userId: user.id, + workspaceId: target.id, + accessMode: McpAccessMode.READ_ONLY, + }); + await t.throwsAsync(credentials.authenticate(issued.token, other.id)); + const response = await app + .POST(`/api/workspaces/${target.id}/mcp`) + .set('Authorization', `Bearer ${issued.token}`) + .send({ jsonrpc: '2.0', id: 1, method: 'initialize', params: {} }) + .expect(200); + t.like(response.body, { jsonrpc: '2.0', id: 1 }); + t.deepEqual( + (await provider.for(user.id, target.id, McpAccessMode.READ_ONLY)).tools.map( + tool => tool.name + ), + ['read_document', 'semantic_search', 'keyword_search'] + ); - await waitForStatus( - async () => { - const status = (await getTranscriptTask(app, workspaceId, job.id)) - ?.status; - if (status === 'failed') { - const task = await db.aiTranscriptTask.findUnique({ - where: { id: job.id }, - select: { errorCode: true }, - }); - throw new Error( - `audio transcription job failed: ${ - task?.errorCode ?? 'unknown error' - }` - ); - } - return status; - }, - 'finished', - 'audio transcription job' - ); + const rotated = await credentials.rotate( + issued.credential.id, + user.id, + target.id, + 30 + ); + t.like((await credentials.list(user.id, target.id))[0], { + status: 'ROTATING', + }); + await credentials.authenticate(issued.token, target.id); + await credentials.revoke(rotated.credential.id, user.id, target.id); + await t.throwsAsync(credentials.authenticate(issued.token, target.id)); + await t.throwsAsync(credentials.authenticate(rotated.token, target.id)); - const result = await settleTranscriptTask(app, workspaceId, job.id); - t.is(result.summaryJson?.title, 'Weekly Sync'); - t.is(result.summaryJson?.actionItems[0]?.description, 'Send recap'); - t.is(result.sourceAudio?.blobId, '1'); - t.is(result.sourceAudio?.mimeType, 'audio/ogg'); - t.is(result.quality?.degraded, true); - t.is(result.quality?.overflowCount, 4); - t.is(result.normalizedSegments?.[0]?.start, '00:00:42'); - t.is(result.normalizedSegments?.[0]?.text, 'Hello, everyone.'); - t.true( - result.summaryJson?.keyPoints.includes('Reviewed launch status') ?? false - ); - } - - { - const job = await submitTranscriptTask( - app, - workspaceId, - '2', - '2.mp3', - [Buffer.from([1, 1]), Buffer.from([1, 2])], - { - sliceManifest: [ - { - index: 0, - fileName: '2-0.opus', - mimeType: 'audio/opus', - startSec: 0, - durationSec: 600, - byteSize: 2, - }, - { - index: 1, - fileName: '2-1.opus', - mimeType: 'audio/opus', - startSec: 605, - durationSec: 120, - byteSize: 2, - }, - ], - } - ); - t.truthy(job.id, 'should have job id'); - - await waitForStatus( - async () => { - const status = (await getTranscriptTask(app, workspaceId, job.id)) - ?.status; - if (status === 'failed') { - const task = await db.aiTranscriptTask.findUnique({ - where: { id: job.id }, - select: { errorCode: true }, - }); - throw new Error( - `audio transcription job failed: ${ - task?.errorCode ?? 'unknown error' - }` - ); - } - return status; - }, - 'finished', - 'audio transcription job' - ); - - const result = await settleTranscriptTask(app, workspaceId, job.id); - t.deepEqual( - result.normalizedSegments?.map(segment => segment.start), - ['00:00:30', '00:00:46', '00:10:35', '00:10:51'] - ); - t.is( - result.normalizedTranscript?.split('\n')[2], - '00:10:35 A: Hello, everyone.' - ); - } -}); - -test('should create different session types and validate prompt constraints', async t => { - const { app } = t.context; - const { id: workspaceId } = await createWorkspace(app); - - const validateSession = async ( - description: string, - workspaceId: string, - createPromise: Promise - ) => { - const sessionId = await createPromise; - - t.truthy(sessionId, description); - t.snapshot( - cleanObject( - [await getCopilotSession(app, workspaceId, sessionId)], - ['id', 'workspaceId', 'promptName'] - ), - `should create session with ${description}` - ); - return sessionId; - }; - - await validateSession( - 'should create workspace session with text prompt', - workspaceId, - createWorkspaceCopilotSession(app, workspaceId, textPromptName) - ); - await validateSession( - 'should create pinned session with text prompt', - workspaceId, - createPinnedCopilotSession(app, workspaceId, 'pinned-doc', textPromptName) - ); - await validateSession( - 'should create doc session with text prompt', - workspaceId, - createDocCopilotSession(app, workspaceId, 'normal-doc', textPromptName) - ); -}); - -test('should list histories for different session types correctly', async t => { - const { app } = t.context; - const { id: workspaceId } = await createWorkspace(app); - const pinnedDocId = 'pinned-doc'; - const docId = 'normal-doc'; - - // create sessions and add messages - const [workspaceSessionId, pinnedSessionId, docSessionId] = await Promise.all( - [ - createWorkspaceCopilotSession(app, workspaceId, textPromptName), - createPinnedCopilotSession(app, workspaceId, pinnedDocId, textPromptName), - createDocCopilotSession(app, workspaceId, docId, textPromptName), - ] - ); - - await Promise.all([ - createCopilotMessage(app, workspaceSessionId, 'workspace message'), - createCopilotMessage(app, pinnedSessionId, 'pinned message'), - createCopilotMessage(app, docSessionId, 'doc message'), - ]); - - const testHistoryQuery = async ( - queryFn: () => Promise, - opts: { - sessionIds?: string[]; - sessionId?: string; - pinned?: boolean; - isEmpty?: boolean; - }, - description: string - ) => { - const s = await queryFn(); - - if (opts.isEmpty) { - t.is(s.length, 0, `should return ${description}`); - return; - } - - if (opts.sessionIds) { - t.is(s.length, opts.sessionIds.length, `should return ${description}`); - const ids = s.map(h => h.sessionId).sort((a, b) => a.localeCompare(b)); - const expectedIds = opts.sessionIds.sort((a, b) => a.localeCompare(b)); - t.deepEqual(ids, expectedIds, `should return correct ${description}`); - } else if (opts.sessionId) { - t.is(s.length, 1, `should return ${description}`); - t.is( - s[0].sessionId, - opts.sessionId, - `should return correct ${description}` - ); - if (opts.pinned !== undefined) { - t.is(s[0].pinned, opts.pinned, `pinned status for ${description}`); - } - } - }; - - // test for getHistories - await testHistoryQuery( - () => getHistories(app, { workspaceId, docId: null }), - { sessionId: workspaceSessionId }, - 'workspace session history' - ); - await testHistoryQuery( - () => getHistories(app, { workspaceId, docId: pinnedDocId }), - { sessionId: pinnedSessionId }, - 'pinned session history' - ); - await testHistoryQuery( - () => getHistories(app, { workspaceId, docId }), - { sessionId: docSessionId }, - 'doc session history' - ); - - // test for getWorkspaceSessions - await testHistoryQuery( - () => getWorkspaceSessions(app, { workspaceId }), - { sessionId: workspaceSessionId, pinned: false }, - 'workspace-level sessions' - ); - - // test for getDocSessions - await testHistoryQuery( - () => - getDocSessions(app, { workspaceId, docId, options: { pinned: false } }), - { sessionId: docSessionId, pinned: false }, - 'doc sessions' - ); - - await testHistoryQuery( - () => getDocSessions(app, { workspaceId, docId: pinnedDocId }), - { sessionId: pinnedSessionId, pinned: true }, - 'pinned doc sessions' - ); - - // test for getPinnedSessions - await testHistoryQuery( - () => getPinnedSessions(app, { workspaceId }), - { sessionId: pinnedSessionId, pinned: true }, - 'pinned sessions' - ); - - await testHistoryQuery( - () => getPinnedSessions(app, { workspaceId, docId: pinnedDocId }), - { sessionId: pinnedSessionId, pinned: true }, - 'pinned session for specific doc' - ); - - await testHistoryQuery( - () => getPinnedSessions(app, { workspaceId, docId }), - { isEmpty: true }, - 'no pinned sessions for non-pinned doc' - ); + const disabled = await credentials.create({ + userId: user.id, + workspaceId: target.id, + name: 'Disabled user', + accessMode: McpAccessMode.READ_ONLY, + expirationDays: 30, + }); + await models.user.update(user.id, { disabled: true }); + await t.throwsAsync(credentials.authenticate(disabled.token, target.id)); + await models.user.update(user.id, { disabled: false }); + await db.mcpCredential.update({ + where: { id: disabled.credential.id }, + data: { expiresAt: new Date(0) }, + }); + await t.throwsAsync(credentials.authenticate(disabled.token, target.id)); }); diff --git a/packages/backend/server/src/__tests__/copilot/copilot.spec.ts b/packages/backend/server/src/__tests__/copilot/copilot.spec.ts deleted file mode 100644 index 0af1b515ce..0000000000 --- a/packages/backend/server/src/__tests__/copilot/copilot.spec.ts +++ /dev/null @@ -1,2708 +0,0 @@ -import { createHash, randomUUID } from 'node:crypto'; -import { Readable } from 'node:stream'; - -import { ProjectRoot } from '@affine-tools/utils/path'; -import { McpAccessMode, PrismaClient } from '@prisma/client'; -import type { TestFn } from 'ava'; -import ava from 'ava'; -import { nanoid } from 'nanoid'; -import Sinon from 'sinon'; -import { z } from 'zod'; - -import { - EventBus, - JobQueue, - RequestMutex, - SpaceAccessDenied, -} from '../../base'; -import { ConfigModule } from '../../base/config'; -import { AuthService } from '../../core/auth'; -import { QuotaModule } from '../../core/quota'; -import { QuotaStateService } from '../../core/quota/state'; -import { StorageModule, WorkspaceBlobStorage } from '../../core/storage'; -import { - ContextCategories, - CopilotSessionModel, - Models, - WorkspaceMemberStatus, - WorkspaceModel, - WorkspaceRole, -} from '../../models'; -import { addDocToRootDoc, type LlmToolCallbackRequest } from '../../native'; -import { CopilotModule } from '../../plugins/copilot'; -import { CopilotContextService } from '../../plugins/copilot/context'; -import { CopilotContextResolver } from '../../plugins/copilot/context/resolver'; -import { - chatMessageFromTurn, - turnFromChatMessage, -} from '../../plugins/copilot/core'; -import { CopilotCronJobs } from '../../plugins/copilot/cron'; -import { - CopilotEmbeddingClientService, - CopilotEmbeddingJob, - MockEmbeddingClient, -} from '../../plugins/copilot/embedding'; -import { McpCredentialService } from '../../plugins/copilot/mcp/credential'; -import { WorkspaceMcpProvider } from '../../plugins/copilot/mcp/provider'; -import { PromptService } from '../../plugins/copilot/prompt'; -import { - CopilotProviderFactory, - CopilotProviderType, - ModelInputType, - ModelOutputType, - OpenAIProvider, -} from '../../plugins/copilot/providers'; -import { TextStreamParser } from '../../plugins/copilot/providers/utils'; -import { CopilotResolver } from '../../plugins/copilot/resolver'; -import { ActionRuntimeBridge } from '../../plugins/copilot/runtime/action-runtime-bridge'; -import { CapabilityRuntime } from '../../plugins/copilot/runtime/capability-runtime'; -import { - parsePromptRenderContract, - parsePromptSessionContract, -} from '../../plugins/copilot/runtime/contracts'; -import { projectActionEventToChatEvent } from '../../plugins/copilot/runtime/hosts/action-stream-host'; -import { CapabilityPolicyHost } from '../../plugins/copilot/runtime/hosts/capability-policy-host'; -import { ConversationHost } from '../../plugins/copilot/runtime/hosts/conversation-host'; -import { ImageResultHost } from '../../plugins/copilot/runtime/hosts/image-result-host'; -import { ModelSelectionPolicy } from '../../plugins/copilot/runtime/model-selection-policy'; -import { PromptRuntime } from '../../plugins/copilot/runtime/prompt-runtime'; -import { getProviderRuntimeHost } from '../../plugins/copilot/runtime/provider-runtime-context'; -import { executeToolCall } from '../../plugins/copilot/runtime/tool/bridge'; -import { TurnOrchestrator } from '../../plugins/copilot/runtime/turn-orchestrator'; -import { ChatSessionService } from '../../plugins/copilot/session'; -import { CopilotStorage } from '../../plugins/copilot/storage'; -import type { - CopilotToolExecuteOptions, - CopilotToolSet, -} from '../../plugins/copilot/tools'; -import { CopilotTranscriptionService } from '../../plugins/copilot/transcript'; -import { CopilotWorkspaceService } from '../../plugins/copilot/workspace'; -import { PaymentModule } from '../../plugins/payment'; -import { SubscriptionService } from '../../plugins/payment/service'; -import { SubscriptionStatus } from '../../plugins/payment/types'; -import { installMockCopilotRuntime, MockCopilotProvider } from '../mocks'; -import { TestingPromptService } from '../mocks/prompt-service.mock'; -import { createTestingApp, TestingApp } from '../utils'; -import { singleUserPromptMessages, systemPrompt } from './prompt-test-helper'; - -type Context = { - auth: AuthService; - module: TestingApp; - db: PrismaClient; - event: EventBus; - models: Models; - workspace: WorkspaceModel; - workspaceStorage: WorkspaceBlobStorage; - copilotSession: CopilotSessionModel; - context: CopilotContextService; - prompt: TestingPromptService; - transcript: CopilotTranscriptionService; - workspaceEmbedding: CopilotWorkspaceService; - factory: CopilotProviderFactory; - session: ChatSessionService; - promptRuntime: PromptRuntime; - chatRuntime: CapabilityRuntime; - conversationHost: ConversationHost; - embeddingClients: CopilotEmbeddingClientService; - jobs: CopilotEmbeddingJob; - imageResults: ImageResultHost; - orchestrator: TurnOrchestrator; - storage: CopilotStorage; - actionBridge: ActionRuntimeBridge; - cronJobs: CopilotCronJobs; - subscription: SubscriptionService; - quotaState: QuotaStateService; - mcpCredentials: McpCredentialService; - mcpProvider: WorkspaceMcpProvider; -}; - -const buildTurn = ( - sessionId: string, - message: Parameters[0] -) => turnFromChatMessage(message, sessionId); - -const cleanSnapshotObject = (obj: unknown, omittedKeys: string[] = []) => - JSON.parse( - JSON.stringify(obj, (k, v) => - ['id', 'createdAt', ...omittedKeys].includes(k) || - v === null || - (typeof v === 'object' && !Object.keys(v).length) - ? undefined - : v - ) - ); - -const cleanFinalMessages = (messages: unknown) => - cleanSnapshotObject(messages, ['attachments']); - -const test = ava as TestFn; -let userId: string; -let restoreMockCopilotNativeRuntime: (() => void) | undefined; - -test.before(async t => { - restoreMockCopilotNativeRuntime = installMockCopilotRuntime(); - const module = await createTestingApp({ - imports: [ - ConfigModule.override({ - copilot: { - providers: { - openai: { - apiKey: process.env.COPILOT_OPENAI_API_KEY ?? '1', - }, - fal: { - apiKey: process.env.COPILOT_FAL_API_KEY ?? '1', - }, - anthropic: { - apiKey: process.env.COPILOT_ANTHROPIC_API_KEY ?? '1', - }, - }, - exa: { - key: process.env.COPILOT_EXA_API_KEY ?? '1', - }, - }, - }), - PaymentModule, - QuotaModule, - StorageModule, - CopilotModule, - ], - tapModule: builder => { - builder.overrideProvider(RequestMutex).useValue({ - acquire: async () => ({ - async [Symbol.asyncDispose]() {}, - }), - }); - builder.overrideProvider(PromptService).useClass(TestingPromptService); - builder.overrideProvider(OpenAIProvider).useClass(MockCopilotProvider); - builder.overrideProvider(SubscriptionService).useClass( - class { - select() { - return { getSubscription: async () => undefined }; - } - } - ); - }, - }); - - const auth = module.get(AuthService); - const db = module.get(PrismaClient); - const event = module.get(EventBus); - const models = module.get(Models); - const workspace = module.get(WorkspaceModel); - const workspaceStorage = module.get(WorkspaceBlobStorage); - const copilotSession = module.get(CopilotSessionModel); - const prompt = module.get(PromptService) as TestingPromptService; - const factory = module.get(CopilotProviderFactory); - - const session = module.get(ChatSessionService); - const promptRuntime = module.get(PromptRuntime); - const chatRuntime = module.get(CapabilityRuntime); - const conversationHost = module.get(ConversationHost); - const imageResults = module.get(ImageResultHost); - const orchestrator = module.get(TurnOrchestrator); - const actionBridge = module.get(ActionRuntimeBridge); - const storage = module.get(CopilotStorage); - - const context = module.get(CopilotContextService); - const embeddingClients = module.get(CopilotEmbeddingClientService); - const jobs = module.get(CopilotEmbeddingJob); - const transcript = module.get(CopilotTranscriptionService); - const workspaceEmbedding = module.get(CopilotWorkspaceService); - const cronJobs = module.get(CopilotCronJobs); - const subscription = module.get(SubscriptionService); - const quotaState = module.get(QuotaStateService); - - t.context.module = module; - t.context.auth = auth; - t.context.db = db; - t.context.event = event; - t.context.models = models; - t.context.workspace = workspace; - t.context.workspaceStorage = workspaceStorage; - t.context.copilotSession = copilotSession; - t.context.prompt = prompt; - t.context.factory = factory; - t.context.session = session; - t.context.promptRuntime = promptRuntime; - t.context.chatRuntime = chatRuntime; - t.context.conversationHost = conversationHost; - t.context.imageResults = imageResults; - t.context.orchestrator = orchestrator; - t.context.actionBridge = actionBridge; - t.context.storage = storage; - t.context.context = context; - t.context.embeddingClients = embeddingClients; - t.context.jobs = jobs; - t.context.transcript = transcript; - t.context.workspaceEmbedding = workspaceEmbedding; - t.context.cronJobs = cronJobs; - t.context.subscription = subscription; - t.context.quotaState = quotaState; - t.context.mcpCredentials = module.get(McpCredentialService); - t.context.mcpProvider = module.get(WorkspaceMcpProvider); - - await module.initTestingDB(); -}); - -let promptName = 'prompt'; - -test.beforeEach(async t => { - Sinon.restore(); - const { auth, prompt } = t.context; - prompt.reset(); - const user = await auth.signUp(`test-${randomUUID()}@affine.pro`, '123456'); - userId = user.id; - promptName = randomUUID().replaceAll('-', ''); -}); - -test.after.always(async t => { - restoreMockCopilotNativeRuntime?.(); - await t.context.module?.close(); -}); - -test('document cleanup reconciles missing and restored copilot state before ack', async t => { - const { db, jobs, models, module, workspace } = t.context; - const queue = module.get(JobQueue); - const deleteEmbedding = Sinon.spy( - models.copilotContext, - 'purgeWorkspaceEmbedding' - ); - const scheduleEmbedding = Sinon.stub( - jobs, - 'addDocEmbeddingQueueFromEvent' - ).resolves(); - - for (const [docId, restored, cleanupVersion] of [ - ['missing-doc', false, 'missing-version'], - ['restored-doc', true, 'restored-version'], - ] as const) { - const ws = await workspace.create(userId); - const root = addDocToRootDoc(Buffer.from([0, 0]), docId, docId); - const missingRoot = addDocToRootDoc( - Buffer.from([0, 0]), - 'live-doc', - 'Live' - ); - await db.snapshot.create({ - data: { - workspaceId: ws.id, - id: ws.id, - blob: restored ? root : missingRoot, - state: Buffer.from([0, 0]), - updatedAt: new Date(), - createdAt: new Date(), - }, - }); - if (restored) { - await db.snapshot.create({ - data: { - workspaceId: ws.id, - id: docId, - blob: addDocToRootDoc(Buffer.from([0, 0]), 'content', 'Content'), - state: Buffer.from([0, 0]), - updatedAt: new Date(), - createdAt: new Date(), - }, - }); - } - - await jobs.reconcileDocumentCleanup({ - workspaceId: ws.id, - docId, - cleanupVersion, - }); - - if (restored) { - t.true(scheduleEmbedding.calledOnceWith({ workspaceId: ws.id, docId })); - t.false(deleteEmbedding.called); - } else { - t.true(deleteEmbedding.calledOnceWith(ws.id, docId)); - t.false(scheduleEmbedding.called); - } - const { payload } = await module.queue.waitFor( - 'backendRuntime.ackDocumentCleanupEffect' - ); - t.deepEqual(payload, { - workspaceId: ws.id, - docId, - cleanupVersion, - effect: 'copilot', - }); - deleteEmbedding.resetHistory(); - scheduleEmbedding.resetHistory(); - (queue.add as Sinon.SinonStub).resetHistory(); - } -}); - -test('MCP credentials stay bound to their endpoint, workspace, and profile', async t => { - const { db, mcpCredentials, mcpProvider, models, workspace } = t.context; - const ws = await workspace.create(userId); - const other = await workspace.create(userId); - const issued = await mcpCredentials.create({ - userId, - workspaceId: ws.id, - name: 'Claude Desktop', - accessMode: McpAccessMode.READ_ONLY, - expirationDays: 90, - }); - - const authenticated = await mcpCredentials.authenticate(issued.token, ws.id); - t.is(authenticated.userId, userId); - await t.throwsAsync(mcpCredentials.authenticate(issued.token, other.id)); - - const response = await t.context.module - .POST(`/api/workspaces/${ws.id}/mcp`) - .set('Authorization', `Bearer ${issued.token}`) - .send({ jsonrpc: '2.0', id: 1, method: 'initialize', params: {} }) - .expect(200); - t.like(response.body, { - jsonrpc: '2.0', - id: 1, - }); - - const server = await mcpProvider.for(userId, ws.id, McpAccessMode.READ_ONLY); - t.deepEqual( - server.tools.map(tool => tool.name), - ['read_document', 'semantic_search', 'keyword_search'] - ); - - const rotated = await mcpCredentials.rotate( - issued.credential.id, - userId, - ws.id, - 30 - ); - const listed = await mcpCredentials.list(userId, ws.id); - t.is(listed.length, 1); - t.is(listed[0].status, 'ROTATING'); - await mcpCredentials.authenticate(issued.token, ws.id); - await mcpCredentials.revoke(rotated.credential.id, userId, ws.id); - await t.throwsAsync(mcpCredentials.authenticate(issued.token, ws.id)); - await t.throwsAsync(mcpCredentials.authenticate(rotated.token, ws.id)); - - const disabled = await mcpCredentials.create({ - userId, - workspaceId: ws.id, - name: 'Disabled user', - accessMode: McpAccessMode.READ_ONLY, - expirationDays: 30, - }); - await models.user.update(userId, { disabled: true }); - await t.throwsAsync(mcpCredentials.authenticate(disabled.token, ws.id)); - - await models.user.update(userId, { disabled: false }); - await db.mcpCredential.update({ - where: { id: disabled.credential.id }, - data: { expiresAt: new Date(0) }, - }); - await t.throwsAsync(mcpCredentials.authenticate(disabled.token, ws.id)); -}); - -test('should reject context file uploads after workspace write access is revoked', async t => { - const { auth, context, models, prompt, session, storage, workspace } = - t.context; - const contextResolver = await t.context.module.resolve( - CopilotContextResolver - ); - - const owner = await auth.signUp(`test-${randomUUID()}@affine.pro`, '123456'); - const member = await auth.signUp(`test-${randomUUID()}@affine.pro`, '123456'); - const ws = await workspace.create(owner.id); - - await models.workspaceUser.set(ws.id, member.id, WorkspaceRole.Collaborator, { - status: WorkspaceMemberStatus.Accepted, - }); - await prompt.set(promptName, 'test', [ - { role: 'system', content: 'hello {{word}}' }, - ]); - - const sessionId = await session.create({ - userId: member.id, - workspaceId: ws.id, - docId: randomUUID(), - promptName, - pinned: false, - }); - const contextSession = await context.create(sessionId); - await models.workspaceUser.set(ws.id, member.id, WorkspaceRole.External); - - Sinon.stub(context, 'canEmbedding').get(() => true); - Sinon.stub(context, 'embeddingClient').get(() => new MockEmbeddingClient()); - const put = Sinon.stub(storage, 'put').resolves(); - const buffer = Buffer.from('test pdf'); - - await t.throwsAsync( - contextResolver.addContextFile( - { id: member.id } as any, - { - req: { - headers: { - 'content-length': String(buffer.length), - }, - }, - } as any, - { contextId: contextSession.id }, - { - filename: 'sample.pdf', - mimetype: 'application/pdf', - createReadStream: () => Readable.from(buffer), - } as any - ), - { - instanceOf: SpaceAccessDenied, - } - ); - - t.false(put.called); -}); - -test('should prioritize user-added context file embedding jobs', async t => { - const { context, jobs, prompt, session, storage, workspace } = t.context; - const contextResolver = await t.context.module.resolve( - CopilotContextResolver - ); - - const ws = await workspace.create(userId); - await prompt.set(promptName, 'test', [ - { role: 'system', content: 'hello {{word}}' }, - ]); - - const sessionId = await session.create({ - userId, - workspaceId: ws.id, - docId: randomUUID(), - promptName, - pinned: false, - }); - const contextSession = await context.create(sessionId); - - Sinon.stub(context, 'canEmbedding').get(() => true); - Sinon.stub(context, 'embeddingClient').get(() => new MockEmbeddingClient()); - const put = Sinon.stub(storage, 'put').resolves(); - const queue = Sinon.stub(jobs, 'addFileEmbeddingQueue').resolves(); - const buffer = Buffer.from('test pdf'); - - await contextResolver.addContextFile( - { id: userId } as any, - { - req: { - headers: { - 'content-length': String(buffer.length), - }, - }, - } as any, - { contextId: contextSession.id }, - { - filename: 'sample.pdf', - mimetype: 'application/pdf', - createReadStream: () => Readable.from(buffer), - } as any - ); - - t.true(put.calledOnce); - t.true(queue.calledOnce); - t.deepEqual(queue.firstCall.args[0], { - userId, - workspaceId: ws.id, - contextId: contextSession.id, - blobId: createHash('sha256').update(buffer).digest('base64url'), - fileId: queue.firstCall.args[0].fileId, - fileName: 'sample.pdf', - }); - t.deepEqual(queue.firstCall.args[1], { priority: 0 }); -}); - -test('should resolve context sessions with the shared embedding client', async t => { - const { context, embeddingClients, prompt, session, workspace } = t.context; - - const ws = await workspace.create(userId); - await prompt.set(promptName, 'test', [ - { role: 'system', content: 'hello {{word}}' }, - ]); - - const sessionId = await session.create({ - userId, - workspaceId: ws.id, - docId: randomUUID(), - promptName, - pinned: false, - }); - const client = new MockEmbeddingClient(); - - Sinon.stub(embeddingClients, 'refresh').resolves(undefined); - Sinon.stub(embeddingClients, 'getClient').returns(client); - await context.onConfigChanged(); - - const contextSession = await context.create(sessionId); - t.is(context.embeddingClient, client); - await t.notThrowsAsync(context.get(contextSession.id)); -}); - -test('should be able to render prompt', async t => { - const { prompt } = t.context; - - const msg = { - role: 'system' as const, - content: 'translate {{src_language}} to {{dest_language}}: {{content}}', - params: { src_language: ['eng'], dest_language: ['chs', 'jpn', 'kor'] }, - }; - const params = { - src_language: 'eng', - dest_language: 'chs', - content: 'hello world', - }; - - await prompt.set(promptName, 'test', [msg]); - const testPrompt = await prompt.get(promptName); - t.assert(testPrompt, 'should have prompt'); - t.is( - prompt.finish(testPrompt!, params).pop()?.content, - 'translate eng to chs: hello world', - 'should render the prompt' - ); - t.deepEqual( - testPrompt?.paramKeys, - Object.keys(params), - 'should have param keys' - ); - t.deepEqual(testPrompt?.params, msg.params, 'should have params'); - // will use first option if a params not provided - t.deepEqual(prompt.finish(testPrompt!, { src_language: 'abc' }), [ - { - content: 'translate eng to chs: ', - params: { dest_language: 'chs', src_language: 'eng' }, - role: 'system', - }, - ]); -}); - -test('should be able to render listed prompt', async t => { - const { prompt } = t.context; - - const msg = { - role: 'system' as const, - content: 'links:\n{{#links}}- {{.}}\n{{/links}}', - }; - const params = { - links: ['https://affine.pro', 'https://github.com/toeverything/affine'], - }; - - await prompt.set(promptName, 'test', [msg]); - const testPrompt = await prompt.get(promptName); - - t.is( - prompt.finish(testPrompt!, params).pop()?.content, - 'links:\n- https://affine.pro\n- https://github.com/toeverything/affine\n', - 'should render the prompt' - ); -}); - -test('PromptContract should preserve render/session payloads and reject legacy aliases', t => { - const render = parsePromptRenderContract({ - messages: [ - { - role: 'system', - content: 'Return JSON only.', - responseFormat: { - type: 'json_schema', - responseSchemaJson: { - type: 'object', - properties: { summary: { type: 'string' } }, - required: ['summary'], - }, - schemaHash: 'schema-hash', - }, - }, - ], - templateParams: {}, - renderParams: { tone: 'brief' }, - }); - - t.deepEqual( - { messages: render.messages, warnings: [] }, - { - messages: render.messages, - warnings: [], - } - ); - - const session = parsePromptSessionContract({ - prompt: { - model: 'gpt-5-mini', - promptTokens: 12, - templateParams: {}, - messages: [systemPrompt('Return JSON only.')], - }, - turns: singleUserPromptMessages('hello'), - renderParams: { tone: 'brief' }, - maxTokenSize: 1024, - }); - - t.is(session.prompt.model, 'gpt-5-mini'); - - const error = t.throws(() => - parsePromptRenderContract({ - messages: [ - { - role: 'system', - content: 'Return JSON only.', - responseFormat: { - type: 'json_schema', - schemaJson: { type: 'object' }, - schemaHash: 'schema-hash', - }, - }, - ], - templateParams: {}, - renderParams: {}, - }) - ); - - t.truthy(error); -}); - -test('capability runtime should require explicit structured schema contract', async t => { - const runtime = new CapabilityRuntime({} as never, {} as never); - - const error = await t.throwsAsync(() => - runtime.generateStructuredValue( - { modelId: 'gpt-5-mini' }, - singleUserPromptMessages('Summarize AFFiNE.'), - { - responseSchemaJson: { - type: 'object', - properties: { summary: { type: 'string' } }, - required: ['summary'], - additionalProperties: false, - }, - } - ) - ); - - t.true(error instanceof Error); - t.regex(error.message, /Structured schema contract is required/); -}); - -// ==================== session ==================== - -test('should be able to manage chat session', async t => { - const { prompt, session } = t.context; - - await prompt.set(promptName, 'model', [ - { role: 'system', content: 'hello {{word}}' }, - ]); - - const params = { word: 'world' }; - const commonParams = { docId: 'test', workspaceId: 'test', pinned: false }; - - const sessionId = await session.create({ - userId, - promptName, - ...commonParams, - }); - t.truthy(sessionId, 'should create session'); - - const s = (await session.get(sessionId))!; - t.is(s.config.sessionId, sessionId, 'should get session'); - t.is(s.config.promptName, promptName, 'should have prompt name'); - t.is(s.model, 'model', 'should have model'); - - s.pushTurn( - buildTurn(sessionId, { - role: 'user', - content: 'hello', - createdAt: new Date(), - }) - ); - - const finalMessages = cleanFinalMessages(s.finish(params)); - t.snapshot(finalMessages, 'should generate the final message'); - await s.save(); - - const s1 = (await session.get(sessionId))!; - t.deepEqual( - cleanFinalMessages(s1.finish(params)), - finalMessages, - 'should same as before message' - ); - t.snapshot( - cleanFinalMessages(s1.finish(params)), - 'should generate different message with another params' - ); - - // should get main session after fork if re-create a chat session for same docId and workspaceId - { - const newSessionId = await session.create({ - userId, - promptName, - ...commonParams, - }); - t.is(newSessionId, sessionId, 'should get same session id'); - } - - // should create a fresh session when reuseLatestChat is explicitly disabled - { - const newSessionId = await session.create({ - userId, - promptName, - ...commonParams, - reuseLatestChat: false, - }); - t.not( - newSessionId, - sessionId, - 'should create new session id when reuseLatestChat is false' - ); - } -}); - -test('should be able to update chat session prompt', async t => { - const { prompt, session } = t.context; - - // Set up a prompt to be used in the session - await prompt.set(promptName, 'model', [ - { role: 'system', content: 'hello {{word}}' }, - ]); - - // Create a session - const sessionId = await session.create({ - promptName, - docId: 'test', - workspaceId: 'test', - userId, - pinned: false, - }); - t.truthy(sessionId, 'should create session'); - - // Update the session - const updatedSessionId = await session.update({ - sessionId, - promptName: 'Chat With AFFiNE AI', - userId, - }); - t.is(updatedSessionId, sessionId, 'should update session with same id'); - - // Verify the session was updated - const updatedSession = await session.get(sessionId); - t.truthy(updatedSession, 'should retrieve updated session'); - t.is( - updatedSession?.config.promptName, - 'Chat With AFFiNE AI', - 'should have updated prompt name' - ); -}); - -test('should be able to fork chat session', async t => { - const { auth, prompt, session } = t.context; - - await prompt.set(promptName, 'model', [ - { role: 'system', content: 'hello {{word}}' }, - ]); - - const params = { word: 'world' }; - const commonParams = { docId: 'test', workspaceId: 'test', pinned: false }; - // create session - const sessionId = await session.create({ - userId, - promptName, - ...commonParams, - }); - const s = (await session.get(sessionId))!; - s.pushTurn( - buildTurn(sessionId, { - role: 'user', - content: 'hello', - createdAt: new Date(), - }) - ); - s.pushTurn( - buildTurn(sessionId, { - role: 'assistant', - content: 'world', - createdAt: new Date(), - }) - ); - s.pushTurn( - buildTurn(sessionId, { - role: 'user', - content: 'aaa', - createdAt: new Date(), - }) - ); - s.pushTurn( - buildTurn(sessionId, { - role: 'assistant', - content: 'bbb', - createdAt: new Date(), - }) - ); - await s.save(); - - // fork session - const latestMessageId = (await session.getState(sessionId))?.turns.find( - turn => turn.role === 'assistant' - )?.id; - t.truthy(latestMessageId); - const forkedSessionId1 = await session.fork({ - userId, - sessionId, - latestMessageId: latestMessageId!, - ...commonParams, - }); - t.not(sessionId, forkedSessionId1, 'should fork a new session'); - - const newUser = await auth.signUp('darksky.1@affine.pro', '123456'); - const forkedSessionId2 = await session.fork({ - userId: newUser.id, - sessionId, - latestMessageId: latestMessageId!, - ...commonParams, - }); - t.not( - forkedSessionId1, - forkedSessionId2, - 'should fork new session with same params' - ); - - // fork session without latestMessageId - const forkedSessionId3 = await session.fork({ - userId, - sessionId, - ...commonParams, - }); - - // fork session with wrong latestMessageId - await t.throwsAsync( - session.fork({ - userId, - sessionId, - latestMessageId: 'wrong-message-id', - ...commonParams, - }), - { - instanceOf: Error, - }, - 'should not able to fork new session with wrong latestMessageId' - ); - - // check forked session messages - { - const s2 = (await session.get(forkedSessionId1))!; - - const finalMessages = s2.finish(params); - t.snapshot( - cleanSnapshotObject(finalMessages), - 'should generate the final message' - ); - } - - // check second times forked session - { - const s2 = (await session.get(forkedSessionId2))!; - - // should overwrite user id - t.is(s2.config.userId, newUser.id, 'should have same user id'); - - const finalMessages = s2.finish(params); - t.snapshot( - cleanSnapshotObject(finalMessages), - 'should generate the final message' - ); - } - - // check third times forked session - { - const s3 = (await session.get(forkedSessionId3))!; - const finalMessages = s3.finish(params); - t.snapshot( - cleanSnapshotObject(finalMessages), - 'should generate the final message' - ); - } - - // check original session messages - { - const s4 = (await session.get(sessionId))!; - const finalMessages = s4.finish(params); - t.snapshot( - cleanSnapshotObject(finalMessages), - 'should generate the final message' - ); - } - - // should get main session after fork if re-create a chat session for same docId and workspaceId - { - const newSessionId = await session.create({ - userId, - promptName, - ...commonParams, - }); - t.is(newSessionId, sessionId, 'should get same session id'); - } -}); - -test('should schedule title generation as a background job', async t => { - const { prompt, session, module, workspace } = t.context; - const jobs = module.get(JobQueue); - - const ws = await workspace.create(userId); - await prompt.set(promptName, 'model', [ - { role: 'system', content: 'hello {{word}}' }, - ]); - - const sessionId = await session.create({ - userId, - promptName, - docId: 'test', - workspaceId: ws.id, - pinned: false, - }); - const chatSession = await session.get(sessionId); - t.truthy(chatSession); - - const addJob = jobs.add as Sinon.SinonStub; - addJob.resetHistory(); - addJob.resolves(); - - chatSession!.pushTurn( - buildTurn(sessionId, { - role: 'user', - content: 'hello', - createdAt: new Date(), - }) - ); - await chatSession!.save(); - - t.true(addJob.calledOnce); - t.deepEqual(addJob.firstCall.args, [ - 'copilot.session.generateTitle', - { sessionId }, - { priority: 100 }, - ]); -}); - -test('should merge latest user turn content and attachments into prompt', async t => { - const { prompt, session } = t.context; - - await prompt.set(promptName, 'model', [ - { role: 'system', content: 'hello {{word}}' }, - ]); - - for (const testCase of [ - { - title: 'text message', - message: { content: 'hello' }, - project: (messages: { content: string }[]) => - messages.map(({ content }) => content), - expected: ['hello world', 'hello'], - }, - { - title: 'attachment message', - message: { attachments: ['https://affine.pro/example.jpg'] as string[] }, - project: (messages: { attachments?: unknown }[]) => - messages.map(({ attachments }) => attachments), - expected: [undefined, ['https://affine.pro/example.jpg']], - }, - { - title: 'empty message', - message: {}, - project: (messages: { content: string }[]) => - messages.map(({ content }) => content), - expected: ['hello world'], - }, - ]) { - const sessionId = await session.create({ - docId: 'test', - workspaceId: 'test', - userId, - promptName, - pinned: false, - }); - const s = (await session.get(sessionId))!; - s.pushTurn( - buildTurn(sessionId, { - role: 'user', - content: testCase.message.content ?? '', - attachments: testCase.message.attachments, - createdAt: new Date(), - }) - ); - t.deepEqual( - testCase.project(s.finish({ word: 'world' })), - testCase.expected, - testCase.title - ); - } -}); - -test('should preserve file handle attachments when merging user content into prompt', async t => { - const { prompt, session } = t.context; - - await prompt.set(promptName, 'model', [ - { role: 'user', content: '{{content}}' }, - ]); - - const sessionId = await session.create({ - docId: 'test', - workspaceId: 'test', - userId, - promptName, - pinned: false, - }); - const s = (await session.get(sessionId))!; - - s.pushTurn( - buildTurn(sessionId, { - role: 'user', - content: 'Summarize this file', - attachments: [ - { - kind: 'file_handle', - fileHandle: 'file_123', - mimeType: 'application/pdf', - }, - ], - createdAt: new Date(), - }) - ); - const finalMessages = s.finish({}); - - t.deepEqual(finalMessages, [ - { - role: 'user', - content: 'Summarize this file', - attachments: [ - { - kind: 'file_handle', - fileHandle: 'file_123', - mimeType: 'application/pdf', - }, - ], - params: { - content: 'Summarize this file', - }, - }, - ]); -}); - -test('should preserve assistant render trace when converting between chat message and turn', t => { - const sessionId = randomUUID(); - const createdAt = new Date('2025-01-01T00:00:00.000Z'); - const message = { - id: 'message-1', - role: 'assistant' as const, - content: 'Final answer', - attachments: [ - { - kind: 'file_handle' as const, - fileHandle: 'file_123', - mimeType: 'application/pdf', - }, - ], - params: { - schemaVersion: 'v1', - }, - streamObjects: [ - { type: 'reasoning' as const, textDelta: 'Plan' }, - { - type: 'tool-call' as const, - toolCallId: 'call_1', - toolName: 'doc_read', - args: { docId: 'doc-1' }, - rawArgumentsText: '{"docId":"doc-1"}', - thought: 'Need the current doc', - }, - { type: 'text-delta' as const, textDelta: 'Final answer' }, - { - type: 'tool-result' as const, - toolCallId: 'call_2', - toolName: 'doc_keyword_search', - args: { query: 'affine' }, - result: { hits: ['doc-2'] }, - }, - ], - createdAt, - }; - - const turn = turnFromChatMessage(message, sessionId); - - t.deepEqual(turn.renderTrace, message.streamObjects); - t.deepEqual( - turn.toolEvents.map(event => event.type), - ['tool_call', 'tool_result'] - ); - t.deepEqual(chatMessageFromTurn(turn), message); -}); - -test('should save message correctly', async t => { - const { prompt, session } = t.context; - - await prompt.set(promptName, 'model', [ - { role: 'system', content: 'hello {{word}}' }, - ]); - - const sessionId = await session.create({ - docId: 'test', - workspaceId: 'test', - userId, - promptName, - pinned: false, - }); - const s = (await session.get(sessionId))!; - - s.pushTurn( - buildTurn(sessionId, { - role: 'user', - content: 'hello', - createdAt: new Date(), - }) - ); - t.is(s.stashTurns.length, 1, 'should get stash turns'); - await s.save(); - t.is(s.stashTurns.length, 0, 'should empty stash turns after save'); -}); - -test('should revert message correctly', async t => { - const { prompt, session } = t.context; - - // init session - let sessionId: string; - { - await prompt.set(promptName, 'model', [ - { role: 'system', content: 'hello {{word}}' }, - ]); - - sessionId = await session.create({ - docId: 'test', - workspaceId: 'test', - userId, - promptName, - pinned: false, - }); - const s = (await session.get(sessionId))!; - - s.pushTurn( - buildTurn(sessionId, { - role: 'user', - content: '1', - createdAt: new Date(), - }) - ); - await s.save(); - } - - // check ChatSession behavior - { - const s = (await session.get(sessionId))!; - s.pushTurn( - buildTurn(sessionId, { - role: 'assistant', - content: '2', - createdAt: new Date(), - }) - ); - s.pushTurn( - buildTurn(sessionId, { - role: 'user', - content: '3', - createdAt: new Date(), - }) - ); - s.pushTurn( - buildTurn(sessionId, { - role: 'assistant', - content: '4', - createdAt: new Date(), - }) - ); - await s.save(); - const beforeRevert = s.finish({ word: 'world' }); - t.snapshot( - cleanSnapshotObject(beforeRevert), - 'should have three messages before revert' - ); - - { - s.revertLatestMessage(false); - const afterRevert = s.finish({ word: 'world' }); - t.snapshot( - cleanSnapshotObject(afterRevert), - 'should remove assistant message after revert' - ); - } - - { - s.revertLatestMessage(true); - const afterRevert = s.finish({ word: 'world' }); - t.snapshot( - cleanSnapshotObject(afterRevert), - 'should remove assistant message after revert' - ); - } - } - - // check database behavior - { - let s = (await session.get(sessionId))!; - - const beforeRevert = s.finish({ word: 'world' }); - t.snapshot( - cleanSnapshotObject(beforeRevert), - 'should have three messages before revert' - ); - - { - await session.revertLatestMessage(sessionId, false); - s = (await session.get(sessionId))!; - const afterRevert = s.finish({ word: 'world' }); - t.snapshot( - cleanSnapshotObject(afterRevert), - 'should remove assistant message after revert' - ); - } - - { - await session.revertLatestMessage(sessionId, true); - s = (await session.get(sessionId))!; - const afterRevert = s.finish({ word: 'world' }); - t.snapshot( - cleanSnapshotObject(afterRevert), - 'should remove assistant message after revert' - ); - } - } -}); - -test('should handle params correctly in chat session', async t => { - const { prompt, session } = t.context; - - await prompt.set(promptName, 'model', [ - { role: 'system', content: 'hello {{word}}' }, - ]); - - const sessionId = await session.create({ - docId: 'test', - workspaceId: 'test', - userId, - promptName, - pinned: false, - }); - - const s = (await session.get(sessionId))!; - - // Case 1: When params is provided directly - { - const directParams = { word: 'direct' }; - const messages = s.finish(directParams); - t.is(messages[0].content, 'hello direct', 'should use provided params'); - } - - // Case 2: When no params provided but last message has params - { - s.pushTurn( - buildTurn(sessionId, { - role: 'user', - content: 'test message', - params: { word: 'fromMessage' }, - createdAt: new Date(), - }) - ); - const messages = s.finish({}); - t.is( - messages[0].content, - 'hello fromMessage', - 'should use params from last message' - ); - } - - // Case 3: When neither params provided nor last message has params - { - s.pushTurn( - buildTurn(sessionId, { - role: 'user', - content: 'test message without params', - createdAt: new Date(), - }) - ); - const messages = s.finish({}); - t.is(messages[0].content, 'hello ', 'should use empty params'); - } -}); - -// ==================== provider ==================== - -test('should be able to get provider', async t => { - const { factory } = t.context; - - { - const p = await factory.getProvider({ outputType: ModelOutputType.Text }); - t.is( - p?.type.toString(), - 'openai', - 'should get provider support text-to-text' - ); - } - - { - const p = await factory.getProvider({ - outputType: ModelOutputType.Image, - inputTypes: [ModelInputType.Image], - modelId: 'lora/image-to-image', - }); - t.is( - p?.type.toString(), - 'fal', - 'should get provider supporting image output' - ); - } - - { - const p = await factory.getProvider( - { - outputType: ModelOutputType.Image, - inputTypes: [ModelInputType.Image], - }, - { prefer: CopilotProviderType.FAL } - ); - t.is( - p?.type.toString(), - 'fal', - 'should get provider supporting text output with image input' - ); - } - - // if a model is not defined and not available in online api - // it should return null - { - const p = await factory.getProvider({ - outputType: ModelOutputType.Text, - inputTypes: [ModelInputType.Text], - modelId: 'gpt-4-not-exist', - }); - t.falsy(p, 'should not get provider'); - } -}); - -test('should resolve provider by prefixed model id', async t => { - const { factory } = t.context; - - const resolved = await factory.resolveProvider({ - modelId: 'openai-default/test', - outputType: ModelOutputType.Text, - }); - t.truthy(resolved, 'should resolve prefixed model id'); - if (!resolved) { - throw new Error('should resolve prefixed model id'); - } - - t.is(resolved.provider.type, CopilotProviderType.OpenAI); - - const result = await getProviderRuntimeHost(resolved.provider).run.text( - { modelId: resolved.modelId }, - [{ role: 'user', content: 'hello' }], - undefined, - resolved.execution - ); - t.is(result, 'generate text to text'); -}); - -test('should fallback to null when prefixed provider id does not exist', async t => { - const { factory } = t.context; - - const provider = await factory.getProviderByModel('unknown/test'); - t.is(provider, null); -}); - -// ==================== action runtime ==================== - -const wrapAsyncIter = async (iter: AsyncIterable) => { - const result: T[] = []; - for await (const r of iter) { - result.push(r); - } - return result; -}; - -test('tool bridge should validate zod tool args before execution', async t => { - const execute = Sinon.stub().resolves({ message: 'executed' }); - const tools: CopilotToolSet = { - safeTool: { - inputSchema: z.object({ - name: z.string(), - }), - execute, - }, - }; - const options: CopilotToolExecuteOptions = {}; - const request: LlmToolCallbackRequest = { - callId: 'call-1', - name: 'safeTool', - args: { name: 123 }, - rawArgumentsText: '{"name":123}', - }; - - const response = await executeToolCall(tools, request, options); - - t.true(response.isError); - t.true(response.output instanceof Object); - t.regex((response.output as { message: string }).message, /Expected string/); - t.false(execute.called); -}); - -test('tool bridge should execute with parsed zod args and preserve callback response metadata', async t => { - const execute = Sinon.stub().resolves({ message: 'executed' }); - const tools: CopilotToolSet = { - safeTool: { - inputSchema: z.object({ - name: z.string().trim(), - }), - execute, - }, - }; - const options: CopilotToolExecuteOptions = {}; - const request: LlmToolCallbackRequest = { - callId: 'call-2', - name: 'safeTool', - args: { name: ' AFFiNE ' }, - rawArgumentsText: '{"name":" AFFiNE "}', - }; - - const response = await executeToolCall(tools, request, options); - - t.false(response.isError ?? false); - t.deepEqual(response, { - callId: 'call-2', - name: 'safeTool', - args: { name: ' AFFiNE ' }, - rawArgumentsText: '{"name":" AFFiNE "}', - argumentParseError: undefined, - output: { message: 'executed' }, - }); - t.deepEqual(execute.firstCall.args, [{ name: 'AFFiNE' }, options]); -}); - -test('tool bridge should reject malformed or unknown tool calls without executing tools', async t => { - const execute = Sinon.stub().resolves({ message: 'executed' }); - const tools: CopilotToolSet = { - safeTool: { - inputSchema: z.record(z.unknown()), - execute, - }, - }; - - const invalidJsonResponse = await executeToolCall( - tools, - { - callId: 'call-3', - name: 'safeTool', - args: {}, - rawArgumentsText: '{"unterminated"', - argumentParseError: 'Unexpected end of JSON input', - }, - {} - ); - const missingToolResponse = await executeToolCall( - tools, - { - callId: 'call-4', - name: 'missingTool', - args: {}, - rawArgumentsText: '{}', - }, - {} - ); - - t.deepEqual(invalidJsonResponse, { - callId: 'call-3', - name: 'safeTool', - args: {}, - rawArgumentsText: '{"unterminated"', - argumentParseError: 'Unexpected end of JSON input', - isError: true, - output: { - message: 'Invalid tool arguments JSON', - rawArguments: '{"unterminated"', - error: 'Unexpected end of JSON input', - }, - }); - t.deepEqual(missingToolResponse, { - callId: 'call-4', - name: 'missingTool', - args: {}, - rawArgumentsText: '{}', - argumentParseError: undefined, - isError: true, - output: { message: 'Tool not found: missingTool' }, - }); - t.false(execute.called); -}); - -test('tool bridge should not mutate global object prototype for adversarial args', async t => { - const execute = Sinon.stub().resolves({ message: 'executed' }); - const tools: CopilotToolSet = { - safeTool: { - inputSchema: z.record(z.unknown()), - execute, - }, - }; - const request: LlmToolCallbackRequest = { - callId: 'call-5', - name: 'safeTool', - args: JSON.parse('{"__proto__":{"polluted":"yes"}}'), - rawArgumentsText: '{"__proto__":{"polluted":"yes"}}', - }; - - const response = await executeToolCall(tools, request, {}); - - t.true(Object.prototype.hasOwnProperty.call(request.args, '__proto__')); - t.false(response.isError ?? false); - t.is((Object.prototype as Record).polluted, undefined); - t.deepEqual(execute.firstCall.args[0], {}); -}); - -test('action stream should expose successful text action result as message', t => { - t.deepEqual( - projectActionEventToChatEvent('message-1', { - type: 'action_done', - actionId: 'slides.outline', - actionVersion: 'v1', - status: 'succeeded', - runId: 'run-1', - result: '- Launch deck', - }), - { - type: 'message', - id: 'message-1', - data: '- Launch deck', - } - ); -}); - -test('turn orchestrator should persist generated image links through image result host', async t => { - const { conversationHost, imageResults, orchestrator, chatRuntime, module } = - t.context; - const capabilityPolicy = module.get(CapabilityPolicyHost); - const session = { - latestUserTurn: { attachments: ['https://example.com/source.png'] }, - config: { sessionId: 'session-1' }, - finish: Sinon.stub().returns([ - { - role: 'system', - content: 'generate image', - params: { quality: 'hd', seed: '7' }, - }, - ]), - } as any; - - Sinon.stub(conversationHost, 'prepareTurn').resolves({ - messageId: 'message-1', - params: {}, - session, - latestTurn: undefined, - } as any); - Sinon.stub(capabilityPolicy, 'selectChat').resolves({ - model: 'test-image-model', - providerOptions: { format: 'png' }, - } as any); - Sinon.stub(chatRuntime, 'streamImageArtifacts').callsFake(async function* () { - yield { url: 'https://remote.example/1.png', media_type: 'image/png' }; - yield { url: 'https://remote.example/2.png', media_type: 'image/png' }; - }); - const persistNativeArtifact = Sinon.stub( - imageResults, - 'persistNativeArtifact' - ).callsFake( - async (_userId, _workspaceId, artifact) => `stored:${artifact.url}` - ); - const persistAssistantTurn = Sinon.stub( - conversationHost, - 'persistAssistantTurn' - ).resolves(); - - const prepared = await orchestrator.streamImages('user-1', 'session-1', { - modelId: 'chat-model', - }); - const result = await wrapAsyncIter(prepared.stream); - - t.deepEqual(result, [ - 'stored:https://remote.example/1.png', - 'stored:https://remote.example/2.png', - ]); - t.deepEqual( - (chatRuntime.streamImageArtifacts as Sinon.SinonStub).firstCall.args[0], - { - modelId: undefined, - inputTypes: [ModelInputType.Image], - } - ); - t.deepEqual( - (chatRuntime.streamImageArtifacts as Sinon.SinonStub).firstCall.args[2], - { - format: 'png', - quality: 'hd', - seed: 7, - signal: undefined, - } - ); - t.deepEqual( - persistNativeArtifact.getCalls().map(call => call.args), - [ - [ - 'user-1', - 'session-1', - { url: 'https://remote.example/1.png', media_type: 'image/png' }, - ], - [ - 'user-1', - 'session-1', - { url: 'https://remote.example/2.png', media_type: 'image/png' }, - ], - ] - ); - t.true(persistAssistantTurn.calledOnce); - t.deepEqual(persistAssistantTurn.firstCall.args[1].attachments, result); -}); - -test('TextStreamParser should format different types of chunks correctly', t => { - // Define interfaces for fixtures - interface BaseFixture { - chunk: any; - description: string; - } - - interface ContentFixture extends BaseFixture { - expected: string; - } - - interface ErrorFixture extends BaseFixture { - errorMessage: string; - } - - type ChunkFixture = ContentFixture | ErrorFixture; - - // Define test fixtures for different chunk types - const fixtures: Record = { - textDelta: { - chunk: { - type: 'text-delta' as const, - text: 'Hello world', - }, - expected: 'Hello world', - description: 'should format text-delta correctly', - }, - reasoning: { - chunk: { - type: 'reasoning-delta' as const, - text: 'I need to think about this', - }, - expected: '\n> [!]\n> I need to think about this', - description: 'should format reasoning as callout', - }, - webSearch: { - chunk: { - type: 'tool-call' as const, - toolName: 'web_search_exa' as const, - toolCallId: 'test-id-1', - input: { query: 'test query', mode: 'AUTO' as const }, - }, - expected: '\n> [!]\n> \n> Searching the web "test query"\n> ', - description: 'should format web search tool call correctly', - }, - webCrawl: { - chunk: { - type: 'tool-call' as const, - toolName: 'web_crawl_exa' as const, - toolCallId: 'test-id-2', - input: { url: 'https://example.com' }, - }, - expected: '\n> [!]\n> \n> Crawling the web "https://example.com"\n> ', - description: 'should format web crawl tool call correctly', - }, - toolResult: { - chunk: { - type: 'tool-result' as const, - toolName: 'web_search_exa' as const, - toolCallId: 'test-id-1', - input: { query: 'test query', mode: 'AUTO' as const }, - output: [ - { - title: 'Test Title', - url: 'https://test.com', - content: 'Test content', - favicon: undefined, - publishedDate: undefined, - author: undefined, - }, - { - title: null, - url: 'https://example.com', - content: 'Example content', - favicon: undefined, - publishedDate: undefined, - author: undefined, - }, - ], - } as any, - expected: - '\n> [!]\n> \n> \n> \n> [Test Title](https://test.com)\n> \n> \n> \n> [https://example.com](https://example.com)\n> \n> \n> ', - description: 'should format tool result correctly', - }, - error: { - chunk: { - type: 'error' as const, - error: { type: 'testError', message: 'Test error message' }, - }, - errorMessage: 'Test error message', - description: 'should throw error for error chunks', - }, - }; - - // Test each chunk type individually - Object.entries(fixtures).forEach(([_name, fixture]) => { - const parser = new TextStreamParser(); - if ('errorMessage' in fixture) { - t.throws( - () => parser.parse(fixture.chunk), - { message: fixture.errorMessage }, - fixture.description - ); - } else { - const result = parser.parse(fixture.chunk); - t.is(result, fixture.expected, fixture.description); - } - }); -}); - -test('TextStreamParser should process a sequence of message chunks', t => { - const parser = new TextStreamParser(); - - // Define test fixtures for mixed chunks sequence - const mixedChunksFixture = { - chunks: [ - // Reasoning chunks - { - id: nanoid(), - type: 'reasoning-delta' as const, - text: 'The user is asking about', - }, - { - id: nanoid(), - type: 'reasoning-delta' as const, - text: ' recent advances in quantum computing', - }, - { - id: nanoid(), - type: 'reasoning-delta' as const, - text: ' and how it might impact', - }, - { - id: nanoid(), - type: 'reasoning-delta' as const, - text: ' cryptography and data security.', - }, - { - id: nanoid(), - type: 'reasoning-delta' as const, - text: ' I should provide information on quantum supremacy achievements', - }, - - // Text delta - { - id: nanoid(), - type: 'text-delta' as const, - text: 'Let me search for the latest breakthroughs in quantum computing and their ', - }, - - // Tool call - { - type: 'tool-call' as const, - toolCallId: 'toolu_01ABCxyz123456789', - toolName: 'web_search_exa' as const, - input: { - query: 'latest quantum computing breakthroughs cryptography impact', - }, - }, - - // Tool result - { - type: 'tool-result' as const, - toolCallId: 'toolu_01ABCxyz123456789', - toolName: 'web_search_exa' as const, - input: { - query: 'latest quantum computing breakthroughs cryptography impact', - }, - output: [ - { - title: 'IBM Unveils 1000-Qubit Quantum Processor', - url: 'https://example.com/tech/quantum-computing-milestone', - }, - ], - }, - - // More text deltas - { - id: nanoid(), - type: 'text-delta' as const, - text: 'implications for security.', - }, - { - id: nanoid(), - type: 'text-delta' as const, - text: '\n\nQuantum computing has made ', - }, - { - id: nanoid(), - type: 'text-delta' as const, - text: 'remarkable progress in the past year. ', - }, - { - id: nanoid(), - type: 'text-delta' as const, - text: 'The development of more stable qubits has accelerated research significantly.', - }, - ], - expected: - '\n> [!]\n> The user is asking about recent advances in quantum computing and how it might impact cryptography and data security. I should provide information on quantum supremacy achievements\n\nLet me search for the latest breakthroughs in quantum computing and their \n> [!]\n> \n> Searching the web "latest quantum computing breakthroughs cryptography impact"\n> \n> \n> \n> [IBM Unveils 1000-Qubit Quantum Processor](https://example.com/tech/quantum-computing-milestone)\n> \n> \n> \n\nimplications for security.\n\nQuantum computing has made remarkable progress in the past year. The development of more stable qubits has accelerated research significantly.', - description: - 'should format the entire stream correctly with proper sequence', - }; - - // Process all chunks sequentially - let result = ''; - for (const chunk of mixedChunksFixture.chunks) { - result += parser.parse(chunk); - } - - // Check final processed output - t.is(result, mixedChunksFixture.expected, mixedChunksFixture.description); -}); - -// ==================== context ==================== -test('should be able to manage context', async t => { - const { - context, - event, - jobs, - prompt, - session, - storage, - workspace, - workspaceStorage, - } = t.context; - - const ws = await workspace.create(userId); - - await prompt.set(promptName, 'model', [ - { role: 'system', content: 'hello {{word}}' }, - ]); - const chatSession = await session.create({ - docId: 'test', - workspaceId: ws.id, - userId, - promptName, - pinned: false, - }); - - // use mocked embedding client - Sinon.stub(context, 'embeddingClient').get(() => new MockEmbeddingClient()); - Sinon.stub(jobs, 'embeddingClient').get(() => new MockEmbeddingClient()); - - { - await t.throwsAsync( - context.create(randomUUID()), - { instanceOf: Error }, - 'should throw error if create context with invalid session id' - ); - - const session = context.create(chatSession); - await t.notThrowsAsync(session, 'should create context with chat session'); - - await t.notThrowsAsync( - context.get((await session).id), - 'should get context after create' - ); - - await t.throwsAsync( - context.get(randomUUID()), - { instanceOf: Error }, - 'should throw error if get context with invalid id' - ); - } - - const fs = await import('node:fs'); - const buffer = fs.readFileSync( - ProjectRoot.join('packages/common/native/fixtures/sample.pdf').toFileUrl() - ); - - { - const session = await context.create(chatSession); - - // file record - { - await storage.put(userId, session.workspaceId, 'blob', buffer); - const file = await session.addFile( - 'blob', - 'sample.pdf', - 'application/pdf' - ); - - const handler = Sinon.spy(event, 'emit'); - - await jobs.embedPendingFile({ - userId, - workspaceId: session.workspaceId, - contextId: session.id, - blobId: file.blobId, - fileId: file.id, - fileName: file.name, - }); - - t.deepEqual(handler.lastCall.args, [ - 'workspace.file.embed.finished', - { - contextId: session.id, - workspaceId: session.workspaceId, - fileId: file.id, - chunkSize: 1, - }, - ]); - - const list = session.files; - t.deepEqual( - list.map(f => f.id), - [file.id], - 'should list file id' - ); - - const result = await session.matchFiles('test', 1, undefined, 1); - t.is(result.length, 1, 'should match context'); - t.is(result[0].fileId, file.id, 'should match file id'); - } - - // blob record - { - const blobId = 'test-blob'; - await workspaceStorage.put(session.workspaceId, blobId, buffer); - - await jobs.embedPendingBlob({ workspaceId: session.workspaceId, blobId }); - - const result = await t.context.context.matchWorkspaceBlobs( - session.workspaceId, - 'test', - 1, - undefined, - 1 - ); - t.is(result.length, 1, 'should match blob embedding'); - t.is(result[0].blobId, blobId, 'should match blob id'); - } - - // doc record - - const addDoc = async () => { - const docId = randomUUID(); - await t.context.db.snapshot.create({ - data: { - workspaceId: session.workspaceId, - id: docId, - blob: Buffer.from([1, 1]), - state: Buffer.from([1, 1]), - updatedAt: new Date(), - createdAt: new Date(), - }, - }); - return docId; - }; - - { - const docId = await addDoc(); - await session.addDocRecord(docId); - const docs = session.docs.map(d => d.id); - t.deepEqual(docs, [docId], 'should list doc id'); - - await session.removeDocRecord(docId); - t.deepEqual(session.docs, [], 'should remove doc id'); - } - - // tag record - { - const tagId = randomUUID(); - - const docId1 = await addDoc(); - const docId2 = await addDoc(); - - { - await session.addCategoryRecord(ContextCategories.Tag, tagId, [docId1]); - const tags = session.tags.map(t => t.id); - t.deepEqual(tags, [tagId], 'should list tag id'); - - const docs = session.tags.flatMap(l => l.docs.map(d => d.id)); - t.deepEqual(docs, [docId1], 'should list doc ids'); - } - - { - await session.addCategoryRecord(ContextCategories.Tag, tagId, [docId2]); - - const docs = session.tags.flatMap(l => l.docs.map(d => d.id)); - t.deepEqual(docs, [docId1, docId2], 'should list doc ids'); - } - - await session.removeCategoryRecord(ContextCategories.Tag, tagId); - t.deepEqual(session.tags, [], 'should remove tag id'); - } - - // collection record - { - const collectionId = randomUUID(); - - const docId1 = await addDoc(); - const docId2 = await addDoc(); - { - await session.addCategoryRecord( - ContextCategories.Collection, - collectionId, - [docId1] - ); - const collection = session.collections.map(l => l.id); - t.deepEqual(collection, [collectionId], 'should list collection id'); - - const docs = session.collections.flatMap(l => l.docs.map(d => d.id)); - t.deepEqual(docs, [docId1], 'should list doc ids'); - } - - { - await session.addCategoryRecord( - ContextCategories.Collection, - collectionId, - [docId2] - ); - - const docs = session.collections.flatMap(l => l.docs.map(d => d.id)); - t.deepEqual(docs, [docId1, docId2], 'should list doc ids'); - } - - await session.removeCategoryRecord( - ContextCategories.Collection, - collectionId - ); - t.deepEqual(session.collections, [], 'should remove collection id'); - } - } -}); - -// ==================== workspace embedding ==================== -test('should be able to manage workspace embedding', async t => { - const { db, jobs, workspace, workspaceEmbedding, context, prompt, session } = - t.context; - - // use mocked embedding client - Sinon.stub(context, 'embeddingClient').get(() => new MockEmbeddingClient()); - Sinon.stub(jobs, 'embeddingClient').get(() => new MockEmbeddingClient()); - - const ws = await workspace.create(userId); - - // should create workspace embedding - { - const { blobId, file } = await workspaceEmbedding.addFile(userId, ws.id, { - filename: 'test.txt', - mimetype: 'text/plain', - encoding: 'utf-8', - createReadStream: () => { - return new Readable({ - read() { - this.push(Buffer.from('content')); - this.push(null); - }, - }); - }, - }); - await workspaceEmbedding.queueFileEmbedding({ - userId, - workspaceId: ws.id, - blobId, - fileId: file.fileId, - fileName: file.fileName, - }); - await jobs.embedPendingFile({ - userId, - workspaceId: ws.id, - contextId: undefined, - blobId, - fileId: file.fileId, - fileName: file.fileName, - }); - - let ret = 0; - while (!ret) { - await new Promise(resolve => setTimeout(resolve, 1000)); - ret = await db.aiWorkspaceFileEmbedding.count({ - where: { workspaceId: ws.id, fileId: file.fileId }, - }); - } - } - - // should create workspace embedding with file - { - await prompt.set(promptName, 'model', [ - { role: 'system', content: 'hello {{word}}' }, - ]); - const sessionId = await session.create({ - docId: 'test', - workspaceId: ws.id, - userId, - promptName, - pinned: false, - }); - const contextSession = await context.create(sessionId); - - const ret = await contextSession.matchFiles('test', 1, undefined, 1); - t.is(ret.length, 1, 'should match workspace context'); - t.is(ret[0].content, 'content', 'should match content'); - - await workspace.update(ws.id, { enableDocEmbedding: false }); - - const ret2 = await contextSession.matchFiles('test', 1, undefined, 1); - t.is(ret2.length, 0, 'should not match workspace context'); - } -}); - -test('should handle generateSessionTitle correctly under various conditions', async t => { - const { prompt, session, promptRuntime, workspace, copilotSession } = - t.context; - - await prompt.set(promptName, 'model', [ - { role: 'user', content: '{{content}}' }, - ]); - const createSession = async ( - options: { - userMessage?: string; - assistantMessage?: string; - existingTitle?: string; - } = {} - ) => { - const ws = await workspace.create(userId); - const sessionId = await session.create({ - docId: 'test-doc', - workspaceId: ws.id, - userId, - promptName, - pinned: false, - }); - - if (options.existingTitle) { - await copilotSession.update({ - userId, - sessionId, - title: options.existingTitle, - }); - } - - const chatSession = await session.get(sessionId); - if (chatSession) { - if (options.userMessage) { - chatSession.pushTurn( - buildTurn(sessionId, { - role: 'user', - content: options.userMessage, - createdAt: new Date(), - }) - ); - } - if (options.assistantMessage) { - chatSession.pushTurn( - buildTurn(sessionId, { - role: 'assistant', - content: options.assistantMessage, - createdAt: new Date(), - }) - ); - } - await chatSession.save(); - } - - return sessionId; - }; - - const testCases = [ - { - name: 'should generate title when conditions are met', - setup: () => - createSession({ - userMessage: 'What is machine learning?', - assistantMessage: - 'Machine learning is a subset of artificial intelligence.', - }), - mockFn: () => 'What is Machine Learning?', - expectSnapshot: true, - }, - { - name: 'should not generate title when session already has title', - setup: () => - createSession({ - userMessage: 'Test message', - assistantMessage: 'Test response', - existingTitle: 'Existing Title', - }), - mockFn: () => 'New Title', - expectSnapshot: true, - expectNotCalled: true, - }, - { - name: 'should not generate title when no user messages exist', - setup: () => - createSession({ assistantMessage: 'Hello! How can I help you?' }), - mockFn: () => 'New Title', - expectSnapshot: true, - expectNotCalled: true, - }, - { - name: 'should not generate title when no assistant messages exist', - setup: () => createSession({ userMessage: 'What is AI?' }), - mockFn: () => 'New Title', - expectSnapshot: true, - expectNotCalled: true, - }, - { - name: 'should handle errors gracefully', - setup: () => - createSession({ - userMessage: 'Test question', - assistantMessage: 'Test answer', - }), - mockFn: () => { - throw new Error('Mock error for testing'); - }, - expectError: 'Mock error for testing', - }, - ]; - - for (const testCase of testCases) { - const sessionId = await testCase.setup(); - let chatWithPromptCalled = false; - - const mockStub = Sinon.stub(promptRuntime, 'runText').callsFake( - async () => { - chatWithPromptCalled = true; - return testCase.mockFn(); - } - ); - - if (testCase.expectError) { - await t.throwsAsync( - () => session.generateSessionTitle({ sessionId }), - { message: testCase.expectError }, - testCase.name - ); - } else { - await session.generateSessionTitle({ sessionId }); - - if (testCase.expectSnapshot) { - const sessionState = await session.getState(sessionId); - t.snapshot( - { - chatWithPromptCalled: testCase.expectNotCalled - ? chatWithPromptCalled - : undefined, - title: sessionState?.conversation.title, - exists: !!sessionState, - }, - testCase.name - ); - } - } - - mockStub.restore(); - } - - { - const sessionId = await createSession({ - userMessage: 'Explain quantum computing briefly', - assistantMessage: 'Quantum computing uses quantum mechanics principles.', - }); - - let capturedArgs: any[] = []; - Sinon.stub(promptRuntime, 'runText').callsFake(async (...args) => { - capturedArgs = args; - return 'Quantum Computing Explained'; - }); - - await session.generateSessionTitle({ sessionId }); - - t.snapshot( - { - promptName: capturedArgs[0], - content: capturedArgs[1]?.content, - }, - 'should use correct prompt for title generation' - ); - } -}); - -test('should handle copilot cron jobs correctly', async t => { - const { cronJobs, copilotSession } = t.context; - - // mock calls - const mockCleanupResult = { removed: 2, cleaned: 3 }; - const mockSessions = [ - { id: 'session1', _count: { messages: 1 } }, - { id: 'session2', _count: { messages: 2 } }, - ]; - const cleanupStub = Sinon.stub( - copilotSession, - 'cleanupEmptySessions' - ).resolves(mockCleanupResult); - const toBeGenerateStub = Sinon.stub( - copilotSession, - 'toBeGenerateTitle' - ).resolves(mockSessions); - const jobAddStub = cronJobs['jobs'].add as Sinon.SinonStub; - jobAddStub.resetHistory(); - jobAddStub.resolves(); - - // daily cleanup job scheduling - { - await cronJobs.dailyCleanupJob(); - t.snapshot( - jobAddStub.getCalls().map(call => ({ - args: call.args, - })), - 'daily job scheduling calls' - ); - - jobAddStub.reset(); - cleanupStub.reset(); - toBeGenerateStub.reset(); - } - - // cleanup empty sessions - { - // mock - cleanupStub.resolves(mockCleanupResult); - toBeGenerateStub.resolves(mockSessions); - - await cronJobs.cleanupEmptySessions(); - t.snapshot( - cleanupStub.getCalls().map(call => ({ - args: call.args.map(arg => (arg instanceof Date ? 'Date' : arg)), // Replace Date with string for stable snapshot - })), - 'cleanup empty sessions calls' - ); - } - - // generate missing titles - await cronJobs.generateMissingTitles(); - t.snapshot( - { - modelCalls: toBeGenerateStub.getCalls().map(call => ({ - args: call.args, - })), - jobCalls: jobAddStub.getCalls().map(call => ({ - args: call.args, - })), - }, - 'title generation calls' - ); - - cleanupStub.restore(); - toBeGenerateStub.restore(); - jobAddStub.resetHistory(); -}); - -test('model selection policy should resolve requested optional models consistently', async t => { - const { module } = t.context; - const modelSelection = module.get(ModelSelectionPolicy); - - t.deepEqual( - modelSelection.resolveRequestedModel({ - defaultModel: 'gpt-5.6-luna', - optionalModels: ['gpt-5.6-luna', 'gpt-5.6-terra', 'claude-sonnet-4-6'], - requestedModelId: 'gpt-5.6-terra', - }), - { - selectedModel: 'gpt-5.6-terra', - matchedOptionalModel: true, - } - ); - - t.deepEqual( - modelSelection.resolveRequestedModel({ - defaultModel: 'gpt-5.6-luna', - optionalModels: ['gpt-5.6-luna', 'gpt-5.6-terra', 'claude-sonnet-4-6'], - requestedModelId: 'openai-default/gpt-5.6-terra', - }), - { - selectedModel: 'openai-default/gpt-5.6-terra', - matchedOptionalModel: true, - } - ); - - t.deepEqual( - modelSelection.resolveRequestedModel({ - defaultModel: 'gpt-5.6-luna', - optionalModels: ['gpt-5.6-luna', 'gpt-5.6-terra', 'claude-sonnet-4-6'], - requestedModelId: 'not-in-optional', - }), - { - selectedModel: 'gpt-5.6-luna', - matchedOptionalModel: false, - } - ); - - t.is( - modelSelection.resolveRequestedModel({ - defaultModel: 'gpt-5.6-luna', - optionalModels: ['gpt-5.6-luna', 'gpt-5.6-terra', 'claude-sonnet-4-6'], - requestedModelId: 'not-in-optional', - }).selectedModel, - 'gpt-5.6-luna' - ); -}); - -test('capability policy host should gate pro model requests by subscription status', async t => { - const { quotaState, subscription, module } = t.context; - const capabilityPolicy = module.get(CapabilityPolicyHost); - - const mockStatus = (status?: SubscriptionStatus) => { - Sinon.restore(); - Sinon.stub(subscription, 'select').callsFake(() => ({ - // @ts-expect-error mock - getSubscription: async () => (status ? { status } : null), - })); - Sinon.stub(quotaState, 'reconcileUserQuotaState').resolves({ - plan: status === SubscriptionStatus.Active ? 'pro' : 'free', - flags: {}, - } as Awaited>); - }; - - // payment disabled -> allow requested if in optional; pro not blocked - { - const model1 = await capabilityPolicy.resolveChatModel({ - userId, - defaultModel: 'gpt-5.6-luna', - optionalModels: ['gpt-5.6-luna', 'gpt-5.6-terra', 'claude-sonnet-4-6'], - proModels: ['gpt-5.6-terra', 'claude-sonnet-4-6'], - requestedModelId: 'gpt-5.6-terra', - paymentEnabled: false, - }); - t.snapshot(model1, 'should honor requested pro model'); - - const model1WithPrefix = await capabilityPolicy.resolveChatModel({ - userId, - defaultModel: 'gpt-5.6-luna', - optionalModels: ['gpt-5.6-luna', 'gpt-5.6-terra', 'claude-sonnet-4-6'], - proModels: ['gpt-5.6-terra', 'claude-sonnet-4-6'], - requestedModelId: 'openai-default/gpt-5.6-terra', - paymentEnabled: false, - }); - t.is( - model1WithPrefix, - 'openai-default/gpt-5.6-terra', - 'should honor requested prefixed pro model' - ); - - const model2 = await capabilityPolicy.resolveChatModel({ - userId, - defaultModel: 'gpt-5.6-luna', - optionalModels: ['gpt-5.6-luna', 'gpt-5.6-terra', 'claude-sonnet-4-6'], - proModels: ['gpt-5.6-terra', 'claude-sonnet-4-6'], - requestedModelId: 'not-in-optional', - paymentEnabled: false, - }); - t.snapshot(model2, 'should fallback to default model'); - } - - // payment enabled + trialing: requesting pro should fallback to default - { - mockStatus(SubscriptionStatus.Trialing); - const model3 = await capabilityPolicy.resolveChatModel({ - userId, - defaultModel: 'gpt-5.6-luna', - optionalModels: ['gpt-5.6-luna', 'gpt-5.6-terra', 'claude-sonnet-4-6'], - proModels: ['gpt-5.6-terra', 'claude-sonnet-4-6'], - requestedModelId: 'gpt-5.6-terra', - paymentEnabled: true, - }); - t.snapshot( - model3, - 'should fallback to default model when requesting pro model during trialing' - ); - - const model3WithPrefix = await capabilityPolicy.resolveChatModel({ - userId, - defaultModel: 'gpt-5.6-luna', - optionalModels: ['gpt-5.6-luna', 'gpt-5.6-terra', 'claude-sonnet-4-6'], - proModels: ['gpt-5.6-terra', 'claude-sonnet-4-6'], - requestedModelId: 'openai-default/gpt-5.6-terra', - paymentEnabled: true, - }); - t.is( - model3WithPrefix, - 'gpt-5.6-luna', - 'should fallback to default model when requesting prefixed pro model during trialing' - ); - - const model4 = await capabilityPolicy.resolveChatModel({ - userId, - defaultModel: 'gpt-5.6-luna', - optionalModels: ['gpt-5.6-luna', 'gpt-5.6-terra', 'claude-sonnet-4-6'], - proModels: ['gpt-5.6-terra', 'claude-sonnet-4-6'], - requestedModelId: 'gpt-5.6-luna', - paymentEnabled: true, - }); - t.snapshot(model4, 'should honor requested non-pro model during trialing'); - - const model5 = await capabilityPolicy.resolveChatModel({ - userId, - defaultModel: 'gpt-5.6-luna', - optionalModels: ['gpt-5.6-luna', 'gpt-5.6-terra', 'claude-sonnet-4-6'], - proModels: ['gpt-5.6-terra', 'claude-sonnet-4-6'], - paymentEnabled: true, - }); - t.snapshot( - model5, - 'should pick default model when no requested model during trialing' - ); - } - - // payment enabled + active: without requested -> default model; requested pro should be honored - { - mockStatus(SubscriptionStatus.Active); - const model6 = await capabilityPolicy.resolveChatModel({ - userId, - defaultModel: 'gpt-5.6-luna', - optionalModels: ['gpt-5.6-luna', 'gpt-5.6-terra', 'claude-sonnet-4-6'], - proModels: ['gpt-5.6-terra', 'claude-sonnet-4-6'], - paymentEnabled: true, - }); - t.snapshot( - model6, - 'should pick default model when no requested model during active' - ); - - const model7 = await capabilityPolicy.resolveChatModel({ - userId, - defaultModel: 'gpt-5.6-luna', - optionalModels: ['gpt-5.6-luna', 'gpt-5.6-terra', 'claude-sonnet-4-6'], - proModels: ['gpt-5.6-terra', 'claude-sonnet-4-6'], - requestedModelId: 'claude-sonnet-4-6', - paymentEnabled: true, - }); - t.snapshot(model7, 'should honor requested pro model during active'); - - const model7WithPrefix = await capabilityPolicy.resolveChatModel({ - userId, - defaultModel: 'gpt-5.6-luna', - optionalModels: ['gpt-5.6-luna', 'gpt-5.6-terra', 'claude-sonnet-4-6'], - proModels: ['gpt-5.6-terra', 'claude-sonnet-4-6'], - requestedModelId: 'openai-default/claude-sonnet-4-6', - paymentEnabled: true, - }); - t.is( - model7WithPrefix, - 'openai-default/claude-sonnet-4-6', - 'should honor requested prefixed pro model during active' - ); - - const model8 = await capabilityPolicy.resolveChatModel({ - userId, - defaultModel: 'gpt-5.6-luna', - optionalModels: ['gpt-5.6-luna', 'gpt-5.6-terra', 'claude-sonnet-4-6'], - proModels: ['gpt-5.6-terra', 'claude-sonnet-4-6'], - requestedModelId: 'not-in-optional', - paymentEnabled: true, - }); - t.snapshot( - model8, - 'should fallback to default model when requesting non-optional model during active' - ); - } -}); - -test('prompt runtime should resolve prefixed optional models consistently', async t => { - const { prompt, promptRuntime, chatRuntime } = t.context; - - const promptName = randomUUID().replaceAll('-', ''); - await prompt.set( - promptName, - 'gpt-5.6-luna', - [{ role: 'user', content: '{{content}}' }], - { proModels: ['gpt-5.6-terra'] }, - { optionalModels: ['gpt-5.6-terra'] } - ); - - const textStub = Sinon.stub(chatRuntime, 'text').resolves('ok'); - - await promptRuntime.runText( - promptName, - { content: 'hello' }, - { modelId: 'openai-default/gpt-5.6-terra' } - ); - t.is( - textStub.firstCall.args[0].modelId, - 'openai-default/gpt-5.6-terra', - 'should preserve accepted provider-prefixed optional model' - ); - - await promptRuntime.runText( - promptName, - { content: 'hello' }, - { modelId: 'openai-default/not-in-optional' } - ); - t.is( - textStub.secondCall.args[0].modelId, - 'gpt-5.6-luna', - 'should fallback to default model for non-optional prefixed model' - ); -}); - -test('resolver models should use resolved provider metadata for display names', async t => { - const { prompt, factory, module } = t.context; - const resolver = module.get(CopilotResolver); - - const promptName = randomUUID().replaceAll('-', ''); - await prompt.set( - promptName, - 'gpt-5.6-luna', - [{ role: 'system', content: 'test' }], - { proModels: ['gpt-5.6-terra'] }, - { optionalModels: ['gpt-5.6-luna', 'gpt-5.6-terra'] } - ); - - const resolveProvider = Sinon.stub(factory, 'resolveProvider').callsFake( - async cond => - ({ - providerId: 'openai-default', - rawModelId: cond.modelId, - modelId: cond.modelId, - profile: { - id: 'openai-default', - type: CopilotProviderType.OpenAI, - enabled: true, - priority: 10, - config: {}, - middleware: {}, - }, - provider: { - resolveModel: (modelId: string) => ({ - id: modelId, - name: `Resolved ${modelId}`, - }), - }, - }) as any - ); - - const models = await resolver.models(promptName); - - t.deepEqual(models.optionalModels, [ - { id: 'gpt-5.6-luna', name: 'Resolved gpt-5.6-luna' }, - { id: 'gpt-5.6-terra', name: 'Resolved gpt-5.6-terra' }, - ]); - t.deepEqual(models.proModels, [ - { id: 'gpt-5.6-terra', name: 'Resolved gpt-5.6-terra' }, - ]); - t.true( - resolveProvider.alwaysCalledWithMatch({ - outputType: ModelOutputType.Text, - }) - ); -}); diff --git a/packages/backend/server/src/__tests__/copilot/execution-metrics.spec.ts b/packages/backend/server/src/__tests__/copilot/execution-metrics.spec.ts deleted file mode 100644 index 5e646ca2ef..0000000000 --- a/packages/backend/server/src/__tests__/copilot/execution-metrics.spec.ts +++ /dev/null @@ -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', - } - ); -}); diff --git a/packages/backend/server/src/__tests__/copilot/host-services.spec.ts b/packages/backend/server/src/__tests__/copilot/host-services.spec.ts deleted file mode 100644 index 03aede827e..0000000000 --- a/packages/backend/server/src/__tests__/copilot/host-services.spec.ts +++ /dev/null @@ -1,1790 +0,0 @@ -import test from 'ava'; -import Sinon from 'sinon'; - -import { type Models } from '../../models'; -import { CopilotAccessPolicy } from '../../plugins/copilot/access'; -import type { ByokFeatureKind } from '../../plugins/copilot/byok/types'; -import { HistoryAttachmentUrlProjector } from '../../plugins/copilot/compat/history-attachment-url-projector'; -import { CompatHistoryProjector } from '../../plugins/copilot/compat/history-projector'; -import { HistoryPromptPreloadProjector } from '../../plugins/copilot/compat/history-prompt-preload-projector'; -import { HistoryVisibilityPolicy } from '../../plugins/copilot/compat/history-visibility-policy'; -import { ConversationPolicy } from '../../plugins/copilot/conversation/policy'; -import type { Turn } from '../../plugins/copilot/core'; -import { CopilotEmbeddingClientService } from '../../plugins/copilot/embedding/client'; -import { CopilotProviderType } from '../../plugins/copilot/providers/types'; -import { - projectActionResultToAssistantTurn, - summarizeActionResult, -} from '../../plugins/copilot/runtime/action-output-projector'; -import { ActionRuntimeBridge } from '../../plugins/copilot/runtime/action-runtime-bridge'; -import { - ActionStreamHost, - projectActionEventToChatEvent, -} from '../../plugins/copilot/runtime/hosts/action-stream-host'; -import { - admittedAttachmentToPromptAttachment, - AttachmentAdmissionHost, -} from '../../plugins/copilot/runtime/hosts/attachment-admission'; -import { - planAdmittedAttachmentMaterialization, - planHostUrlAttachmentMaterialization, -} from '../../plugins/copilot/runtime/hosts/attachment-materialization-planner'; -import { - AttachmentMaterializer, - resolveAttachmentFetchUrl, -} from '../../plugins/copilot/runtime/hosts/attachment-materializer'; -import { ConversationHost } from '../../plugins/copilot/runtime/hosts/conversation-host'; -import { ImageResultHost } from '../../plugins/copilot/runtime/hosts/image-result-host'; -import { ResponsePostprocessor } from '../../plugins/copilot/runtime/hosts/response-postprocessor'; -import { TurnPersistence } from '../../plugins/copilot/runtime/hosts/turn-persistence'; -import { ToolRuntime } from '../../plugins/copilot/runtime/tool-runtime'; - -function stubTurnPersistence( - persistProjectedResult: Sinon.SinonStub = Sinon.stub().resolves(null) -) { - return { - persistProjectedResult, - } as unknown as TurnPersistence; -} - -function stubConversationSession(latestUserTurn?: unknown) { - return { - config: { - sessionId: 'session-1', - userId: 'user-1', - workspaceId: 'workspace-1', - }, - model: 'gpt-4o-mini', - stashTurns: latestUserTurn ? [latestUserTurn] : [], - latestUserTurn, - revertLatestMessage: Sinon.stub(), - }; -} - -test('ConversationPolicy should treat zero quota limit as exhausted', async t => { - const policy = new ConversationPolicy( - { - userFeature: { has: Sinon.stub().resolves(false) }, - copilotSession: { countUserMessages: Sinon.stub().resolves(0) }, - } as any, - { - getUserQuota: Sinon.stub().resolves({ copilotActionLimit: 0 }), - } as any - ); - - t.false(await policy.hasQuota('user-1')); - await t.throwsAsync(policy.checkQuota('user-1')); -}); - -type TurnRouteAccessCase = { - name: string; - profiles: Array<{ id: string }>; - featureKind?: 'embedding' | 'rerank' | 'workspace_indexing'; - byokLeaseId?: string; - quotaBackedRoutesAllowed?: boolean; - expectedQuotaCalls: number; - expectedError?: string; - expectedQuotaBackedRoutesAllowed?: boolean; -}; - -const turnRouteAccessCases: TurnRouteAccessCase[] = [ - { - name: 'checks quota when BYOK does not cover the route', - profiles: [], - expectedQuotaCalls: 1, - expectedError: 'quota exceeded', - }, - { - name: 'skips quota when BYOK covers the route', - profiles: [{ id: 'profile-1' }], - byokLeaseId: 'lease-1', - expectedQuotaCalls: 0, - expectedQuotaBackedRoutesAllowed: undefined, - }, - { - name: 'preserves explicit quota-backed route disable override', - profiles: [], - quotaBackedRoutesAllowed: false, - expectedQuotaCalls: 0, - expectedQuotaBackedRoutesAllowed: false, - }, - { - name: 'does not check user quota for unmetered service features', - profiles: [], - featureKind: 'rerank', - expectedQuotaCalls: 0, - expectedQuotaBackedRoutesAllowed: true, - }, -]; - -for (const matrixCase of turnRouteAccessCases) { - test(`CopilotAccessPolicy resolve turn route access: ${matrixCase.name}`, async t => { - const checkQuota = Sinon.stub().rejects(new Error('quota exceeded')); - const getProfiles = Sinon.stub().resolves(matrixCase.profiles); - const access = new CopilotAccessPolicy( - { checkQuota } as any, - { getProfiles } as any - ); - - const promise = access.resolveTurnRouteAccess({ - userId: 'user-1', - workspaceId: 'workspace-1', - byokLeaseId: matrixCase.byokLeaseId, - featureKind: matrixCase.featureKind, - quotaBackedRoutesAllowed: matrixCase.quotaBackedRoutesAllowed, - }); - - if (matrixCase.expectedError) { - await t.throwsAsync(promise, { message: matrixCase.expectedError }); - } else { - const routeAccess = await promise; - t.is( - routeAccess.quotaBackedRoutesAllowed, - matrixCase.expectedQuotaBackedRoutesAllowed - ); - } - t.is(checkQuota.callCount, matrixCase.expectedQuotaCalls); - if (matrixCase.expectedQuotaCalls) { - Sinon.assert.calledWithExactly(checkQuota, 'user-1'); - } - if (matrixCase.byokLeaseId) { - Sinon.assert.calledWithMatch(getProfiles, { - byokLeaseId: matrixCase.byokLeaseId, - }); - } - }); -} - -type ByokCoverageCase = { - featureKind?: ByokFeatureKind; - expected: { local: boolean; server: boolean }; -}; - -const byokCoverageCases: ByokCoverageCase[] = [ - { featureKind: 'chat', expected: { local: true, server: true } }, - { featureKind: 'action', expected: { local: true, server: true } }, - { featureKind: 'image', expected: { local: true, server: true } }, - { featureKind: 'transcript', expected: { local: false, server: true } }, - { featureKind: 'embedding', expected: { local: false, server: true } }, - { - featureKind: 'workspace_indexing', - expected: { local: false, server: true }, - }, - { featureKind: 'rerank', expected: { local: false, server: true } }, - { expected: { local: true, server: true } }, -]; - -for (const matrixCase of byokCoverageCases) { - test(`CopilotAccessPolicy should resolve BYOK coverage for ${matrixCase.featureKind ?? 'default'}`, async t => { - const getProfiles = Sinon.stub().resolves([]); - const access = new CopilotAccessPolicy( - { hasQuota: Sinon.stub().resolves(true) } as any, - { getProfiles } as any - ); - - await access.getByokProfiles({ - userId: 'user-1', - workspaceId: 'workspace-1', - featureKind: matrixCase.featureKind, - }); - - t.like(getProfiles.firstCall.args[0], { - userId: 'user-1', - workspaceId: 'workspace-1', - }); - t.is(getProfiles.firstCall.args[0].featureKind, matrixCase.featureKind); - t.deepEqual(getProfiles.firstCall.args[1], matrixCase.expected); - }); -} - -test('CopilotAccessPolicy assertQuotaOrByok should honor quota-backed route disable', async t => { - const checkQuota = Sinon.stub().resolves(undefined); - const access = new CopilotAccessPolicy( - { checkQuota } as any, - { getProfiles: Sinon.stub().resolves([]) } as any - ); - - await t.throwsAsync( - access.assertQuotaOrByok({ - userId: 'user-1', - workspaceId: 'workspace-1', - featureKind: 'transcript', - quotaBackedRoutesAllowed: false, - }) - ); - Sinon.assert.notCalled(checkQuota); -}); - -test('ConversationHost should delegate empty no-message stream access', async t => { - const session = stubConversationSession(); - const resolveTurnRouteAccess = Sinon.stub().rejects( - new Error('quota exceeded') - ); - const host = new ConversationHost( - { - get: Sinon.stub().resolves(session), - revertLatestMessage: Sinon.stub().resolves(undefined), - } as any, - {} as any, - {} as any, - { resolveTurnRouteAccess } as any - ); - - await t.throwsAsync(host.prepareTurn('user-1', 'session-1', {}), { - message: 'quota exceeded', - }); - Sinon.assert.calledOnceWithMatch(resolveTurnRouteAccess, { - userId: 'user-1', - workspaceId: 'workspace-1', - }); -}); - -test('ConversationHost should return access decision for empty no-message stream', async t => { - const session = stubConversationSession(); - const resolveTurnRouteAccess = Sinon.stub().resolves({ - byokProfiles: [{ id: 'profile-1' }], - quotaBackedRoutesAllowed: undefined, - }); - const host = new ConversationHost( - { - get: Sinon.stub().resolves(session), - revertLatestMessage: Sinon.stub().resolves(undefined), - } as any, - {} as any, - {} as any, - { resolveTurnRouteAccess } as any - ); - - const prepared = await host.prepareTurn('user-1', 'session-1', {}); - - t.is(prepared.latestTurn, undefined); - t.is(prepared.quotaBackedRoutesAllowed, undefined); - Sinon.assert.calledOnce(resolveTurnRouteAccess); -}); - -test('ConversationHost should replay accepted tokens without rechecking quota', async t => { - const acceptedTurn: Turn = { - id: 'turn-1', - conversationId: 'session-1', - role: 'user', - content: 'hello', - attachments: [], - metadata: {}, - renderTrace: [], - toolEvents: [], - createdAt: new Date(), - }; - const session = { - ...stubConversationSession(acceptedTurn), - findTurn: Sinon.stub().withArgs('turn-1').returns(acceptedTurn), - }; - const resolveTurnRouteAccess = Sinon.stub().rejects( - new Error('quota exceeded') - ); - const host = new ConversationHost( - { - get: Sinon.stub().resolves(session), - revertLatestMessage: Sinon.stub().resolves(undefined), - } as any, - { - getAccepted: Sinon.stub().resolves({ - sessionId: 'session-1', - turnId: 'turn-1', - }), - } as any, - {} as any, - { resolveTurnRouteAccess } as any - ); - - const prepared = await host.prepareTurn('user-1', 'session-1', { - messageId: 'message-1', - }); - - t.is(prepared.latestTurn, acceptedTurn); - t.true(prepared.quotaBackedRoutesAllowed); - Sinon.assert.notCalled(resolveTurnRouteAccess); -}); - -test('ConversationHost should replay durable tokens without rechecking quota', async t => { - const durableTurn: Turn = { - id: 'turn-1', - conversationId: 'session-1', - role: 'user', - content: 'hello', - attachments: [], - metadata: {}, - renderTrace: [], - toolEvents: [], - createdAt: new Date(), - }; - const session = { - ...stubConversationSession(durableTurn), - findTurn: Sinon.stub().withArgs('turn-1').returns(durableTurn), - pushPersistedTurn: Sinon.stub(), - }; - const resolveTurnRouteAccess = Sinon.stub().rejects( - new Error('quota exceeded') - ); - const markAccepted = Sinon.stub().resolves(undefined); - const host = new ConversationHost( - { - get: Sinon.stub().resolves(session), - findTurnByCompatSubmissionId: Sinon.stub().resolves(durableTurn), - revertLatestMessage: Sinon.stub().resolves(undefined), - } as any, - { - getAccepted: Sinon.stub().resolves(undefined), - markAccepted, - } as any, - { - acquire: Sinon.stub().resolves({ - [Symbol.asyncDispose]: Sinon.stub().resolves(undefined), - }), - } as any, - { resolveTurnRouteAccess } as any - ); - - const prepared = await host.prepareTurn('user-1', 'session-1', { - messageId: 'message-1', - }); - - t.is(prepared.latestTurn, durableTurn); - t.true(prepared.quotaBackedRoutesAllowed); - Sinon.assert.calledOnceWithMatch(markAccepted, 'message-1', { - sessionId: 'session-1', - turnId: 'turn-1', - }); - Sinon.assert.notCalled(resolveTurnRouteAccess); -}); - -test('ToolRuntime should pass route context into prompt-backed tools', async t => { - const promptRuntime = { - runText: Sinon.stub().resolves('done'), - }; - const runtime = new ToolRuntime( - {} as any, - {} as any, - {} as any, - {} as any, - {} as any, - {} as any, - promptRuntime as any, - {} as any - ); - - const tools = await runtime.getTools( - { - tools: ['codeArtifact'], - user: 'user-1', - session: 'session-1', - workspace: 'workspace-1', - byokLeaseId: 'lease-1', - featureKind: 'chat', - quotaBackedRoutesAllowed: false, - }, - 'gpt-4o-mini' - ); - - const result = await tools.code_artifact.execute?.( - { title: 'Demo', userPrompt: 'build a page' }, - {} - ); - - t.like(result as object, { title: 'Demo' }); - Sinon.assert.calledOnceWithMatch( - promptRuntime.runText, - 'Code Artifact', - { content: 'build a page' }, - { - providerOptions: { - user: 'user-1', - session: 'session-1', - workspace: 'workspace-1', - byokLeaseId: 'lease-1', - featureKind: 'chat', - quotaBackedRoutesAllowed: false, - }, - } - ); -}); - -test('ResponsePostprocessor should build text, object and image assistant turns', t => { - const postprocessor = new ResponsePostprocessor(); - - const textTurn = postprocessor.buildTextAssistantTurn('session-1', 'hello'); - const objectTurn = postprocessor.buildObjectAssistantTurn('session-1', [ - { type: 'text-delta', textDelta: 'hel' }, - { type: 'text-delta', textDelta: 'lo' }, - ]); - const imageTurn = postprocessor.buildImageAssistantTurn('session-1', [ - 'https://example.com/image.png', - ]); - - t.like(textTurn, { - conversationId: 'session-1', - role: 'assistant', - content: 'hello', - attachments: [], - }); - t.like(objectTurn, { - conversationId: 'session-1', - role: 'assistant', - content: 'hello', - }); - t.like(imageTurn, { - conversationId: 'session-1', - role: 'assistant', - content: '', - attachments: ['https://example.com/image.png'], - }); -}); - -test('TurnPersistence should delegate assistant turn persistence through ConversationHost', async t => { - const persistAssistantTurn = Sinon.stub().resolves(); - const persistence = new TurnPersistence( - { persistAssistantTurn } as any, - new ResponsePostprocessor() - ); - const session = { - config: { sessionId: 'session-1' }, - } as any; - - await persistence.persistObjectResult( - session, - [{ type: 'text-delta', textDelta: 'done' }], - true - ); - - t.is(persistAssistantTurn.callCount, 1); - const [persistedSession, persistedTurn, persistedAborted] = - persistAssistantTurn.firstCall.args; - t.is(persistedSession, session); - t.is(persistedAborted, true); - t.like(persistedTurn, { - conversationId: 'session-1', - role: 'assistant', - content: 'done', - }); -}); - -test('TurnPersistence should persist text and image assistant turns through ConversationHost', async t => { - const persistAssistantTurn = Sinon.stub().resolves(); - const persistence = new TurnPersistence( - { persistAssistantTurn } as any, - new ResponsePostprocessor() - ); - const session = { - config: { sessionId: 'session-1' }, - } as any; - - await persistence.persistTextResult(session, 'plain text', false); - await persistence.persistImageResult( - session, - ['https://example.com/generated.png'], - false - ); - - t.is(persistAssistantTurn.callCount, 2); - t.like(persistAssistantTurn.firstCall.args[1], { - conversationId: 'session-1', - role: 'assistant', - content: 'plain text', - attachments: [], - }); - t.like(persistAssistantTurn.secondCall.args[1], { - conversationId: 'session-1', - role: 'assistant', - content: '', - attachments: ['https://example.com/generated.png'], - }); -}); - -test('ImageResultHost should persist native base64 artifact with native MIME', async t => { - const storage = { - put: Sinon.stub().resolves('data:image/webp;base64,aW1n'), - handleRemoteLink: Sinon.stub(), - }; - const host = new ImageResultHost(storage as any); - - const persisted = await host.persistNativeArtifact('user-1', 'workspace-1', { - data_base64: 'aW1n', - media_type: 'image/webp', - }); - - t.is(persisted, 'data:image/webp;base64,aW1n'); - Sinon.assert.calledOnceWithMatch( - storage.put, - 'user-1', - 'workspace-1', - Sinon.match.string, - Buffer.from('aW1n', 'base64'), - 'image/webp' - ); -}); - -test('action result projection should map final result to assistant turn', t => { - const session = { - config: { sessionId: 'session-1' }, - stashTurns: [{ id: 'assistant-1' }], - }; - - const turn = projectActionResultToAssistantTurn({ - session: session as any, - actionId: 'mindmap.generate', - wasAborted: false, - result: { - content: 'done', - attachments: ['https://example.com/a.png'], - params: { mode: 'mindmap' }, - }, - }); - - t.like(turn, { - conversationId: 'session-1', - role: 'assistant', - content: 'done', - attachments: ['https://example.com/a.png'], - metadata: { mode: 'mindmap' }, - }); - t.deepEqual(turn?.renderTrace, []); -}); - -test('action result projection should summarize primitive text result', t => { - const turn = projectActionResultToAssistantTurn({ - session: { - config: { sessionId: 'session-1' }, - stashTurns: [], - } as any, - actionId: 'mindmap.generate', - wasAborted: false, - result: 'plain text', - }); - - t.like(turn, { - conversationId: 'session-1', - role: 'assistant', - content: 'plain text', - attachments: [], - }); - t.is(summarizeActionResult('plain text'), 'plain text'); -}); - -test('ActionRuntimeBridge should persist projected assistant message id', async t => { - class TestActionRuntimeBridge extends ActionRuntimeBridge { - protected override runNativeStream() { - return (async function* () { - yield { - type: 'action_done' as const, - actionId: 'mindmap.generate', - actionVersion: 'v1', - status: 'succeeded' as const, - result: 'done', - }; - })(); - } - } - const completedRuns: unknown[] = []; - const persistProjectedResult = Sinon.stub().resolves('assistant-after-save'); - const actionRun = { - create: async () => ({ id: 'run-1' }), - markRunning: async (id: string) => ({ id, status: 'running' }), - complete: async (id: string, input: unknown) => { - completedRuns.push({ id, input }); - return { id, ...(input as Record) }; - }, - }; - const bridge = new TestActionRuntimeBridge( - { - copilotActionRun: actionRun, - } as unknown as Models, - stubTurnPersistence(persistProjectedResult), - undefined - ); - - for await (const event of bridge.runStream({ - userId: 'user-1', - workspaceId: 'workspace-1', - session: { - config: { sessionId: 'session-1' }, - stashTurns: [], - } as any, - actionId: 'mindmap.generate', - actionVersion: 'v1', - })) { - void event; - } - - t.is(persistProjectedResult.callCount, 1); - t.like((completedRuns[0] as { input: Record }).input, { - assistantMessageId: 'assistant-after-save', - }); -}); - -test('action result projection should map image result url to assistant attachments', t => { - const turn = projectActionResultToAssistantTurn({ - session: { - config: { sessionId: 'session-1' }, - stashTurns: [], - } as any, - actionId: 'image.filter.pixel', - wasAborted: false, - result: { url: 'https://example.com/final.png' }, - }); - - t.deepEqual(turn?.attachments, ['https://example.com/final.png']); -}); - -test('CopilotEmbeddingClientService should keep dispatch client across global config refreshes', async t => { - const taskPolicy = { - resolveEmbeddingModelId: () => 'text-embedding-3-large', - }; - const runtime = { - embeddingConfigured: Sinon.stub() - .onFirstCall() - .resolves(true) - .onSecondCall() - .resolves(false), - }; - const service = new CopilotEmbeddingClientService( - taskPolicy as any, - runtime as any - ); - - const first = await service.refresh(); - t.truthy(first); - t.truthy(service.getClient()); - - const second = await service.refresh(); - t.truthy(second); - t.is(service.getClient(), second); - Sinon.assert.calledTwice(runtime.embeddingConfigured); - Sinon.assert.alwaysCalledWithExactly( - runtime.embeddingConfigured, - 'text-embedding-3-large' - ); -}); - -test('CopilotEmbeddingClientService should keep workspace-routed embedding client without global provider', async t => { - const taskPolicy = { - resolveEmbeddingModelId: () => 'gemini-embedding-001', - resolveRerankModelId: () => 'gpt-4o-mini', - }; - const runtime = { - embeddingConfigured: Sinon.stub().resolves(false), - }; - const service = new CopilotEmbeddingClientService( - taskPolicy as any, - runtime as any - ); - - const client = await service.refresh(); - - t.truthy(client); - t.is(service.getClient(), client); - Sinon.assert.calledOnceWithExactly( - runtime.embeddingConfigured, - 'gemini-embedding-001' - ); -}); - -test('CopilotEmbeddingClientService should pass workspace context into embedding routes', async t => { - const signal = new AbortController().signal; - const taskPolicy = { - resolveEmbeddingModelId: () => 'gemini-embedding-001', - resolveRerankModelId: () => 'gpt-4o-mini', - }; - const runtime = { - embeddingConfigured: Sinon.stub().resolves(true), - embed: Sinon.stub().resolves([[0.1]]), - rerank: Sinon.stub().resolves([0.8]), - }; - const service = new CopilotEmbeddingClientService( - taskPolicy as any, - runtime as any - ); - const client = await service.refresh(); - - t.truthy(client); - await client?.getEmbeddings(['hello'], { - workspaceId: 'workspace-1', - userId: 'user-1', - featureKind: 'workspace_indexing', - signal, - }); - - Sinon.assert.calledOnceWithMatch( - runtime.embed, - 'gemini-embedding-001', - ['hello'], - { - dimensions: Sinon.match.number, - workspace: 'workspace-1', - user: 'user-1', - featureKind: 'workspace_indexing', - signal, - } - ); - - await client?.reRank( - 'hello', - [{ chunk: 0, content: 'hello', distance: 0.2 }], - 1, - { - workspaceId: 'workspace-1', - userId: 'user-1', - featureKind: 'workspace_indexing', - signal, - } - ); - - Sinon.assert.calledOnceWithMatch( - runtime.rerank, - 'gpt-4o-mini', - { - query: 'hello', - candidates: [{ id: '0', text: 'hello' }], - }, - { - workspace: 'workspace-1', - user: 'user-1', - featureKind: 'rerank', - signal, - } - ); -}); - -test('CompatHistoryProjector should compose visibility, prompt preload and attachment url projection', t => { - const projector = new CompatHistoryProjector( - new HistoryVisibilityPolicy(), - new HistoryPromptPreloadProjector({ - finish: () => [ - { - role: 'assistant', - content: 'preload', - createdAt: new Date('2026-01-01T00:00:00.000Z'), - }, - ], - } as any), - new HistoryAttachmentUrlProjector() - ); - const createdAt = new Date('2026-01-01T00:00:00.000Z'); - const updatedAt = new Date('2026-01-01T00:10:00.000Z'); - - const visible = projector.projectHistory( - { - conversation: { - id: 'session-1', - userId: 'user-1', - workspaceId: 'workspace-1', - docId: null, - parentId: null, - pinned: false, - title: 'History', - createdAt, - updatedAt, - } as any, - turns: [ - { - conversationId: 'session-1', - role: 'user', - content: 'show the file', - attachments: [{ kind: 'url', url: 'https://example.com/file.pdf' }], - renderTrace: [], - toolEvents: [], - metadata: {}, - createdAt: updatedAt, - }, - ], - prompt: { - name: 'builtin', - action: 'summary', - model: 'gpt-5-mini', - optionalModels: [], - params: {}, - source: 'built_in', - } as any, - tokenCost: 42, - }, - { - requestUserId: 'user-1', - action: true, - withMessages: true, - withPrompt: true, - } - ); - - t.truthy(visible); - t.is(visible?.messages.length, 2); - t.is(visible?.messages[0]?.content, 'preload'); - t.deepEqual(visible?.messages[1]?.attachments, [ - 'https://example.com/file.pdf', - ]); - - const hidden = projector.projectHistory( - { - conversation: { - id: 'session-2', - userId: 'another-user', - workspaceId: 'workspace-1', - docId: null, - parentId: null, - pinned: false, - title: 'Hidden', - createdAt, - updatedAt, - } as any, - turns: [], - prompt: { - name: 'builtin', - action: 'summary', - model: 'gpt-5-mini', - optionalModels: [], - params: {}, - source: 'built_in', - } as any, - tokenCost: 0, - }, - { - requestUserId: 'user-1', - action: false, - withMessages: true, - } - ); - - t.is(hidden, undefined); -}); - -test('AttachmentAdmissionHost should reject remote attachments through host fetch admission', async t => { - const materializer = { - fetchRemoteAttachment: Sinon.stub().rejects(new Error('SSRF blocked')), - }; - const host = new AttachmentAdmissionHost( - materializer as unknown as AttachmentMaterializer - ); - - await t.throwsAsync( - host.admitPromptAttachment('http://127.0.0.1/internal.png', { - userId: 'user-1', - workspaceId: 'workspace-1', - sessionId: 'session-1', - }), - { message: /SSRF blocked/ } - ); - Sinon.assert.calledOnceWithExactly( - materializer.fetchRemoteAttachment, - 'http://127.0.0.1/internal.png', - Sinon.match({ - maxBytes: 64 * 1024 * 1024, - }) - ); -}); - -test('AttachmentAdmissionHost should prefer trusted host MIME over data URL prefix', async t => { - const host = new AttachmentAdmissionHost({ - fetchRemoteAttachment: Sinon.stub(), - } as unknown as AttachmentMaterializer); - const data = Buffer.from('audio-bytes', 'utf8').toString('base64'); - - const admitted = await host.admitPromptAttachment( - { - attachment: `data:image/png;base64,${data}`, - mimeType: 'audio/webm', - }, - { - userId: 'user-1', - workspaceId: 'workspace-1', - sessionId: 'session-1', - } - ); - - t.like(admitted, { - kind: 'bytes', - mimeType: 'audio/webm', - size: Buffer.byteLength('audio-bytes'), - }); -}); - -test('AttachmentAdmissionHost should keep declared Gemini audio MIME after remote prefetch', async t => { - const materializer = { - fetchRemoteAttachment: Sinon.stub().resolves({ - data: Buffer.from('audio-bytes', 'utf8').toString('base64'), - mimeType: 'image/png', - }), - }; - const host = new AttachmentAdmissionHost( - materializer as unknown as AttachmentMaterializer - ); - - const admitted = await host.admitPromptAttachment( - { - kind: 'url', - url: 'https://example.com/recording', - mimeType: 'audio/mpeg', - providerHint: { provider: CopilotProviderType.Gemini, kind: 'audio' }, - }, - { - userId: 'user-1', - workspaceId: 'workspace-1', - } - ); - const promptAttachment = admittedAttachmentToPromptAttachment(admitted); - - t.is(admitted.mimeType, 'audio/mpeg'); - t.deepEqual(promptAttachment, { - kind: 'bytes', - data: Buffer.from('audio-bytes', 'utf8').toString('base64'), - encoding: 'base64', - mimeType: 'audio/mpeg', - fileName: undefined, - providerHint: { provider: 'gemini', kind: 'audio' }, - }); -}); - -test('AttachmentMaterializer should resolve gs attachments through storage HTTPS fetch URL', t => { - t.is( - resolveAttachmentFetchUrl('gs://bucket/audio.opus').toString(), - 'https://storage.googleapis.com/bucket/audio.opus' - ); - t.is( - resolveAttachmentFetchUrl( - 'gs://bucket/folder/audio.opus?alt=media' - ).toString(), - 'https://storage.googleapis.com/bucket/folder/audio.opus?alt=media' - ); -}); - -test('ActionRuntimeBridge should persist action run status around native stream', async t => { - class TestActionRuntimeBridge extends ActionRuntimeBridge { - protected override runNativeStream() { - return (async function* () { - yield { - type: 'action_start' as const, - actionId: 'mindmap.generate', - actionVersion: 'v1', - status: 'running' as const, - }; - yield { - type: 'action_done' as const, - actionId: 'mindmap.generate', - actionVersion: 'v1', - status: 'succeeded' as const, - result: { nodes: [{ text: 'Root' }] }, - }; - })(); - } - } - const createdRuns: unknown[] = []; - const completedRuns: unknown[] = []; - const actionRun = { - create: async (input: unknown) => { - createdRuns.push(input); - return { id: 'run-1' }; - }, - markRunning: async (id: string) => ({ id, status: 'running' }), - complete: async (id: string, input: unknown) => { - completedRuns.push({ id, input }); - return { id, ...(input as Record) }; - }, - }; - const bridge = new TestActionRuntimeBridge( - { - copilotActionRun: actionRun, - } as unknown as Models, - stubTurnPersistence(), - undefined - ); - - const events = []; - for await (const event of bridge.runStream({ - userId: 'user-1', - workspaceId: 'workspace-1', - session: undefined, - actionId: 'mindmap.generate', - actionVersion: 'v1', - inputSnapshot: { prompt: 'make map' }, - nativeInput: { - input: { - mockOutput: { - generate: { - nodes: [{ text: 'Root' }], - }, - }, - }, - }, - })) { - events.push(event); - } - - t.is(events[0]?.runId, 'run-1'); - t.is(events.at(-1)?.type, 'action_done'); - t.like(createdRuns[0], { - userId: 'user-1', - workspaceId: 'workspace-1', - actionId: 'mindmap.generate', - actionVersion: 'v1', - inputSnapshot: { prompt: 'make map' }, - }); - t.like(completedRuns[0] as { id: string; input: Record }, { - id: 'run-1', - input: { - status: 'succeeded', - result: { nodes: [{ text: 'Root' }] }, - artifacts: [], - resultSummary: '{"nodes":[{"text":"Root"}]}', - errorCode: null, - trace: undefined, - assistantMessageId: null, - }, - }); -}); - -test('ActionRuntimeBridge should derive retry attempt from previous action run', async t => { - const createdRuns: unknown[] = []; - class TestActionRuntimeBridge extends ActionRuntimeBridge { - protected override runNativeStream() { - return (async function* () { - yield { - type: 'action_done' as const, - actionId: 'mindmap.generate', - actionVersion: 'v1', - status: 'succeeded' as const, - result: { content: 'retry attempt derived' }, - }; - })(); - } - } - const actionRun = { - get: async (id: string) => ({ - id, - userId: 'user-1', - workspaceId: 'workspace-1', - sessionId: null, - actionId: 'mindmap.generate', - actionVersion: 'v1', - attempt: 2, - }), - create: async (input: unknown) => { - createdRuns.push(input); - return { id: 'run-3' }; - }, - markRunning: async (id: string) => ({ id, status: 'running' }), - complete: async (id: string, input: unknown) => ({ id, input }), - }; - const bridge = new TestActionRuntimeBridge( - { copilotActionRun: actionRun } as unknown as Models, - stubTurnPersistence(), - undefined - ); - - for await (const event of bridge.runStream({ - userId: 'user-1', - workspaceId: 'workspace-1', - actionId: 'mindmap.generate', - actionVersion: 'v1', - retryOf: 'run-2', - })) { - void event; - } - - t.like(createdRuns[0] as Record, { - attempt: 3, - retryOf: 'run-2', - }); -}); - -test('ActionRuntimeBridge should reject retry source from different action owner', async t => { - class TestActionRuntimeBridge extends ActionRuntimeBridge { - protected override runNativeStream(): never { - throw new Error('owner mismatch should reject before native stream'); - } - } - const actionRun = { - get: async (id: string) => ({ - id, - userId: 'other-user', - workspaceId: 'workspace-1', - sessionId: null, - actionId: 'mindmap.generate', - actionVersion: 'v1', - attempt: 1, - }), - create: async () => { - throw new Error('create should not be called'); - }, - }; - const bridge = new TestActionRuntimeBridge( - { copilotActionRun: actionRun } as unknown as Models, - stubTurnPersistence(), - undefined - ); - - await t.throwsAsync( - async () => { - for await (const event of bridge.runStream({ - userId: 'user-1', - workspaceId: 'workspace-1', - actionId: 'mindmap.generate', - actionVersion: 'v1', - retryOf: 'run-1', - })) { - void event; - } - }, - { message: /does not match current action/ } - ); -}); - -test('ActionRuntimeBridge should validate retry source before accepting explicit attempt', async t => { - class TestActionRuntimeBridge extends ActionRuntimeBridge { - protected override runNativeStream(): never { - throw new Error('explicit retry attempt should reject before stream'); - } - } - const actionRun = { - get: async (id: string) => ({ - id, - userId: 'other-user', - workspaceId: 'workspace-1', - sessionId: null, - actionId: 'mindmap.generate', - actionVersion: 'v1', - attempt: 1, - }), - create: async () => { - throw new Error('create should not be called'); - }, - }; - const bridge = new TestActionRuntimeBridge( - { copilotActionRun: actionRun } as unknown as Models, - stubTurnPersistence(), - undefined - ); - - await t.throwsAsync( - async () => { - for await (const event of bridge.runStream({ - userId: 'user-1', - workspaceId: 'workspace-1', - actionId: 'mindmap.generate', - actionVersion: 'v1', - retryOf: 'run-1', - attempt: 3, - })) { - void event; - } - }, - { message: /does not match current action/ } - ); -}); - -test('ActionRuntimeBridge should reject retry source bound to another session', async t => { - class TestActionRuntimeBridge extends ActionRuntimeBridge { - protected override runNativeStream(): never { - throw new Error('session mismatch should reject before native stream'); - } - } - const actionRun = { - get: async (id: string) => ({ - id, - userId: 'user-1', - workspaceId: 'workspace-1', - sessionId: 'previous-session', - actionId: 'mindmap.generate', - actionVersion: 'v1', - attempt: 1, - }), - create: async () => { - throw new Error('create should not be called'); - }, - }; - const bridge = new TestActionRuntimeBridge( - { copilotActionRun: actionRun } as unknown as Models, - stubTurnPersistence(), - undefined - ); - - await t.throwsAsync( - async () => { - for await (const event of bridge.runStream({ - userId: 'user-1', - workspaceId: 'workspace-1', - actionId: 'mindmap.generate', - actionVersion: 'v1', - retryOf: 'run-1', - })) { - void event; - } - }, - { message: /does not match current action/ } - ); -}); - -test('ActionRuntimeBridge should persist attachments and lightweight trace', async t => { - const completedRuns: unknown[] = []; - class TestActionRuntimeBridge extends ActionRuntimeBridge { - protected override runNativeStream() { - return (async function* () { - yield { - type: 'attachment' as const, - actionId: 'image.filter.pixel', - actionVersion: 'v1', - attachment: { url: 'https://example.com/pixel.png' }, - }; - yield { - type: 'action_done' as const, - actionId: 'image.filter.pixel', - actionVersion: 'v1', - status: 'succeeded' as const, - result: { - content: 'done', - artifacts: [{ url: 'https://example.com/final.png' }], - }, - }; - })(); - } - } - const actionRun = { - create: async () => ({ id: 'run-1' }), - markRunning: async (id: string) => ({ id, status: 'running' }), - complete: async (id: string, input: unknown) => { - completedRuns.push({ id, input }); - return { id, ...(input as Record) }; - }, - }; - const bridge = new TestActionRuntimeBridge( - { copilotActionRun: actionRun } as unknown as Models, - stubTurnPersistence(), - undefined - ); - - for await (const event of bridge.runStream({ - userId: 'user-1', - workspaceId: 'workspace-1', - actionId: 'image.filter.pixel', - actionVersion: 'v1', - })) { - void event; - } - - t.like((completedRuns[0] as { input: Record }).input, { - status: 'succeeded', - }); - t.deepEqual( - (completedRuns[0] as { input: { artifacts: unknown } }).input.artifacts, - [ - { url: 'https://example.com/pixel.png' }, - { url: 'https://example.com/final.png' }, - ] - ); - t.is( - (completedRuns[0] as { input: { trace: unknown } }).input.trace, - undefined - ); -}); - -test('ActionRuntimeBridge should inject prepared structured routes into native input', async t => { - const capturedInputs: unknown[] = []; - class TestActionRuntimeBridge extends ActionRuntimeBridge { - protected override runNativeStream(input: unknown) { - capturedInputs.push(input); - return (async function* () { - yield { - type: 'action_done' as const, - actionId: 'mindmap.generate', - actionVersion: 'v1', - status: 'succeeded' as const, - result: { content: 'ok' }, - }; - })(); - } - } - const actionRun = { - create: async () => ({ id: 'run-1' }), - markRunning: async (id: string) => ({ id, status: 'running' }), - complete: async (id: string, input: unknown) => ({ id, input }), - }; - const plans = { - buildStructuredPlan: async (model: { modelId?: string }) => { - t.deepEqual(model, { modelId: 'model-1' }); - return { - nativeDispatch: { - structured: { - routes: [{ provider: 'openai', modelId: 'model-1' }], - }, - }, - }; - }, - }; - const bridge = new TestActionRuntimeBridge( - { copilotActionRun: actionRun } as unknown as Models, - stubTurnPersistence(), - plans as any - ); - - for await (const event of bridge.runStream({ - userId: 'user-1', - workspaceId: 'workspace-1', - actionId: 'mindmap.generate', - actionVersion: 'v1', - prepareStructuredRoutes: { - stepId: 'generate', - modelId: 'model-1', - messages: [{ role: 'user', content: 'make a map' }], - }, - })) { - void event; - } - - const nativeInput = capturedInputs[0] as { - input: { preparedRoutes: Record }; - }; - t.deepEqual(nativeInput.input.preparedRoutes.generate, [ - { provider: 'openai', modelId: 'model-1' }, - ]); -}); - -test('ActionRuntimeBridge should inject prepared image routes and persist attachment events', async t => { - const capturedInputs: unknown[] = []; - const completedRuns: unknown[] = []; - class TestActionRuntimeBridge extends ActionRuntimeBridge { - protected override runNativeStream(input: unknown) { - capturedInputs.push(input); - return (async function* () { - yield { - type: 'attachment' as const, - actionId: 'image.filter.sketch', - actionVersion: 'v1', - attachment: { url: 'data:image/png;base64,aW1hZ2U=' }, - }; - yield { - type: 'action_done' as const, - actionId: 'image.filter.sketch', - actionVersion: 'v1', - status: 'succeeded' as const, - result: { url: 'data:image/png;base64,aW1hZ2U=' }, - }; - })(); - } - } - const actionRun = { - create: async () => ({ id: 'run-1' }), - markRunning: async (id: string) => ({ id, status: 'running' }), - complete: async (id: string, input: unknown) => { - completedRuns.push({ id, input }); - return { id, input }; - }, - }; - const plans = { - buildImagePlan: async (model: { modelId?: string }) => { - t.deepEqual(model, { modelId: 'gpt-image-1' }); - return { - nativeDispatch: { - image: { - routes: [{ provider: 'openai', modelId: 'gpt-image-1' }], - }, - }, - }; - }, - }; - const bridge = new TestActionRuntimeBridge( - { copilotActionRun: actionRun } as unknown as Models, - stubTurnPersistence(), - plans as any - ); - const events = []; - - for await (const event of bridge.runStream({ - userId: 'user-1', - workspaceId: 'workspace-1', - actionId: 'image.filter.sketch', - actionVersion: 'v1', - prepareImageRoutes: { - stepId: 'generate-image', - modelId: 'gpt-image-1', - messages: [{ role: 'user', content: 'draw' }], - }, - persistAttachment: async attachment => ({ - ...(attachment as Record), - url: 'affine://image-result', - }), - })) { - events.push(event); - } - - const nativeInput = capturedInputs[0] as { - input: { preparedRoutes: Record }; - }; - t.deepEqual(nativeInput.input.preparedRoutes['generate-image'], [ - { provider: 'openai', modelId: 'gpt-image-1' }, - ]); - t.deepEqual(events[0].attachment, { url: 'affine://image-result' }); - t.like((completedRuns[0] as { input: Record }).input, { - artifacts: [{ url: 'affine://image-result' }], - }); -}); - -test('ActionRuntimeBridge should persist aborted status from abort signal', async t => { - class TestActionRuntimeBridge extends ActionRuntimeBridge { - protected override runNativeStream() { - return (async function* () { - yield { - type: 'action_start' as const, - actionId: 'mindmap.generate', - actionVersion: 'v1', - status: 'running' as const, - }; - })(); - } - } - const completedRuns: unknown[] = []; - const actionRun = { - create: async () => ({ id: 'run-1' }), - markRunning: async (id: string) => ({ id, status: 'running' }), - complete: async (id: string, input: unknown) => { - completedRuns.push({ id, input }); - return { id, ...(input as Record) }; - }, - }; - const abort = new AbortController(); - abort.abort(); - const bridge = new TestActionRuntimeBridge( - { copilotActionRun: actionRun } as unknown as Models, - stubTurnPersistence(), - undefined - ); - - for await (const event of bridge.runStream({ - userId: 'user-1', - workspaceId: 'workspace-1', - actionId: 'mindmap.generate', - actionVersion: 'v1', - signal: abort.signal, - })) { - void event; - } - - t.like((completedRuns[0] as { input: Record }).input, { - status: 'aborted', - errorCode: undefined, - }); -}); - -test('ActionStreamHost should project native action events into ChatEvent envelope', t => { - t.deepEqual( - projectActionEventToChatEvent('message-1', { - type: 'attachment', - actionId: 'mindmap.generate', - actionVersion: 'v1', - runId: 'run-1', - attachment: { url: 'https://example.com/a.png' }, - }), - { - type: 'attachment', - id: 'message-1', - data: { url: 'https://example.com/a.png' }, - } - ); - t.deepEqual( - projectActionEventToChatEvent('message-1', { - type: 'action_done', - actionId: 'mindmap.generate', - actionVersion: 'v1', - runId: 'run-1', - status: 'succeeded', - }), - { - type: 'event', - id: 'message-1', - data: { - type: 'action_done', - actionId: 'mindmap.generate', - actionVersion: 'v1', - runId: 'run-1', - status: 'succeeded', - }, - } - ); -}); - -test('ActionStreamHost should prepare action turn and bridge native stream', async t => { - const bridgeInputs: unknown[] = []; - const session = { - config: { - sessionId: 'session-1', - workspaceId: 'workspace-1', - docId: 'doc-1', - promptName: 'mindmap.generate', - promptConfig: {}, - }, - finish: Sinon.stub().returns([{ role: 'user', content: 'make a map' }]), - }; - const conversations = { - prepareTurn: Sinon.stub().resolves({ - messageId: 'submission-1', - params: { topic: 'planning' }, - session, - latestTurn: { id: 'turn-1' }, - }), - buildLatestTurnPromptParams: Sinon.stub().returns({ - content: 'make a map', - }), - }; - const prompts = { - get: Sinon.stub().resolves({ - model: 'prompt-model', - config: {}, - }), - finish: Sinon.stub().returns([{ role: 'user', content: 'make a map' }]), - }; - const bridge = { - runStream: (input: unknown) => { - bridgeInputs.push(input); - return (async function* () { - yield { - type: 'action_done' as const, - actionId: 'mindmap.generate', - actionVersion: 'v1', - status: 'succeeded' as const, - runId: 'run-1', - result: { content: 'ok' }, - }; - })(); - }, - }; - const host = new ActionStreamHost( - conversations as any, - bridge as unknown as ActionRuntimeBridge, - prompts as any, - {} as any - ); - - const prepared = await host.stream('user-1', 'session-1', { - actionId: 'mindmap.generate', - actionVersion: 'v1', - modelId: 'model-1', - retry: 'true', - runId: 'run-1', - messageId: 'submission-1', - }); - const events = []; - for await (const event of prepared.stream) { - events.push(event); - } - - t.is(prepared.messageId, 'submission-1'); - t.is(prepared.actionId, 'mindmap.generate'); - t.is(prepared.actionVersion, 'v1'); - t.is(events.at(-1)?.type, 'action_done'); - Sinon.assert.calledOnceWithExactly( - conversations.prepareTurn, - 'user-1', - 'session-1', - { - actionId: 'mindmap.generate', - actionVersion: 'v1', - modelId: 'model-1', - retry: 'true', - runId: 'run-1', - messageId: 'submission-1', - } - ); - t.like(bridgeInputs[0] as Record, { - userId: 'user-1', - workspaceId: 'workspace-1', - docId: 'doc-1', - userMessageId: 'turn-1', - compatSubmissionId: 'submission-1', - actionId: 'mindmap.generate', - actionVersion: 'v1', - retryOf: 'run-1', - }); - t.like( - (bridgeInputs[0] as { prepareStructuredRoutes: Record }) - .prepareStructuredRoutes, - { - stepId: 'generate', - modelId: 'model-1', - messages: [{ role: 'user', content: 'make a map' }], - responseSchemaJson: { - type: 'object', - properties: { - result: { type: 'string' }, - }, - required: ['result'], - additionalProperties: false, - }, - } - ); - Sinon.assert.calledOnceWithExactly(prompts.get, 'mindmap.generate'); -}); - -test('ActionStreamHost should prepare image action routes and persist native attachments', async t => { - const bridgeInputs: any[] = []; - const imageResults = { - persistNativeArtifact: Sinon.stub().resolves('affine://image-result'), - }; - const session = { - config: { - sessionId: 'session-1', - workspaceId: 'workspace-1', - docId: 'doc-1', - promptName: 'image.filter.sketch', - promptConfig: {}, - }, - finish: Sinon.stub().returns([{ role: 'user', content: 'fallback' }]), - }; - const conversations = { - prepareTurn: Sinon.stub().resolves({ - messageId: 'submission-1', - params: { content: 'make a sketch' }, - session, - latestTurn: { id: 'turn-1' }, - }), - buildLatestTurnPromptParams: Sinon.stub().returns({}), - }; - const prompts = { - get: Sinon.stub().resolves({ - model: 'gpt-image-1', - config: { quality: 'high' }, - }), - finish: Sinon.stub().returns([{ role: 'user', content: 'make a sketch' }]), - }; - const bridge = { - runStream: (input: any) => { - bridgeInputs.push(input); - return (async function* () { - const attachment = await input.persistAttachment({ - data_base64: 'aW1hZ2U=', - media_type: 'image/png', - }); - yield { - type: 'attachment' as const, - actionId: 'image.filter.sketch', - actionVersion: 'v1', - runId: 'run-1', - attachment, - }; - })(); - }, - }; - const host = new ActionStreamHost( - conversations as any, - bridge as unknown as ActionRuntimeBridge, - prompts as any, - imageResults as any - ); - - const prepared = await host.stream('user-1', 'session-1', { - actionId: 'image.filter.sketch', - modelId: 'chat-model', - }); - const events = []; - for await (const event of prepared.stream) { - events.push(event); - } - - t.is(bridgeInputs[0].prepareStructuredRoutes, undefined); - t.like(bridgeInputs[0].prepareImageRoutes, { - stepId: 'generate-image', - modelId: 'gpt-image-1', - messages: [{ role: 'user', content: 'make a sketch' }], - }); - t.like(bridgeInputs[0].prepareImageRoutes.options, { - quality: 'high', - user: 'user-1', - workspace: 'workspace-1', - session: 'session-1', - }); - t.deepEqual(events[0].attachment, { - url: 'affine://image-result', - mimeType: 'image/png', - }); - Sinon.assert.calledOnceWithExactly( - imageResults.persistNativeArtifact, - 'user-1', - 'workspace-1', - { - data_base64: 'aW1hZ2U=', - media_type: 'image/png', - } - ); -}); - -test('attachment materialization planner should keep admitted bytes inline', async t => { - const host = new AttachmentAdmissionHost({ - fetchRemoteAttachment: Sinon.stub(), - } as unknown as AttachmentMaterializer); - const admitted = await host.admitPromptAttachment( - { - kind: 'bytes', - data: Buffer.from('image-bytes', 'utf8').toString('base64'), - mimeType: 'image/png', - }, - { - userId: 'user-1', - workspaceId: 'workspace-1', - } - ); - t.deepEqual(planAdmittedAttachmentMaterialization(admitted), { - mode: 'inline', - reason: 'admitted_bytes', - attachment: { - kind: 'bytes', - data: Buffer.from('image-bytes', 'utf8').toString('base64'), - encoding: 'base64', - mimeType: 'image/png', - fileName: undefined, - providerHint: undefined, - }, - }); -}); - -test('attachment materialization planner should separate Gemini remote reference and inline prefetch', async t => { - const backendConfig = { - base_url: 'https://generativelanguage.googleapis.com/v1beta', - auth_token: 'test-key', - request_layer: 'gemini_api' as const, - }; - - const inlinePlan = await planHostUrlAttachmentMaterialization( - 'gemini', - backendConfig, - { - attachmentId: 'att-inline', - url: 'https://example.com/a.mp3', - expectedMime: 'audio/mpeg', - maxSize: 64 * 1024 * 1024, - } - ); - const remotePlan = await planHostUrlAttachmentMaterialization( - 'gemini', - backendConfig, - { - attachmentId: 'att-file', - url: 'https://generativelanguage.googleapis.com/v1beta/files/file-123', - expectedMime: 'application/pdf', - maxSize: 64 * 1024 * 1024, - } - ); - - t.like(inlinePlan, { - mode: 'materialization_request', - reason: 'gemini_api_inline_http_url', - }); - t.like( - inlinePlan.mode === 'materialization_request' - ? inlinePlan.request - : undefined, - { - attachmentId: 'att-inline', - target: 'bytes', - expectedMime: 'audio/mpeg', - redirectPolicy: 'follow-safe', - } - ); - t.like(remotePlan, { - mode: 'remote_reference', - reason: 'gemini_api_file_uri', - url: 'https://generativelanguage.googleapis.com/v1beta/files/file-123', - }); -}); diff --git a/packages/backend/server/src/__tests__/copilot/native-provider.spec.ts b/packages/backend/server/src/__tests__/copilot/native-provider.spec.ts deleted file mode 100644 index b287c0fb16..0000000000 --- a/packages/backend/server/src/__tests__/copilot/native-provider.spec.ts +++ /dev/null @@ -1,2000 +0,0 @@ -import serverNativeModule from '@affine/server-native'; -import test from 'ava'; -import { z } from 'zod'; - -import { CopilotPromptInvalid, CopilotProviderSideError } from '../../base'; -import { - type LlmBackendConfig, - type LlmEmbeddingRequest, - type LlmRequest, - type LlmRerankRequest, - type LlmStructuredRequest, - type LlmStructuredResponse, - type LlmToolLoopStreamEvent, - parseNativeStructuredOutput, -} from '../../native'; -import { - type NodeTextMiddleware, - ProviderMiddlewareConfig, -} from '../../plugins/copilot/config'; -import { GeminiProvider } from '../../plugins/copilot/providers/gemini/gemini'; -import { GeminiVertexProvider } from '../../plugins/copilot/providers/gemini/vertex'; -import { OpenAIProvider } from '../../plugins/copilot/providers/openai'; -import { - CopilotProviderType, - type PromptMessage, - type StreamObject, -} from '../../plugins/copilot/providers/types'; -import { getVertexGoogleBaseUrl } from '../../plugins/copilot/providers/utils'; -import { - buildPromptStructuredResponseFromFields, - buildStructuredResponseContract, - buildToolContracts, - type RequiredStructuredOutputContract, - requireStructuredOutputContract, -} from '../../plugins/copilot/runtime/contracts'; -import { - buildCanonicalNativeRequest, - buildCanonicalNativeStructuredRequest, - buildNativeRequest, - buildNativeStructuredRequest, -} from '../../plugins/copilot/runtime/native-request-runtime'; -import { getProviderRuntimeHost } from '../../plugins/copilot/runtime/provider-runtime-context'; -import type { ToolLoopBackend } from '../../plugins/copilot/runtime/tool/bridge'; -import { createToolExecutionCallback } from '../../plugins/copilot/runtime/tool/bridge'; -import { NativeProviderAdapter } from '../../plugins/copilot/runtime/tool/native-adapter'; -import { NativeRuntimeAdapter } from '../../plugins/copilot/runtime/tool/native-runtime-adapter'; -import type { - CopilotToolExecuteOptions, - CopilotToolSet, -} from '../../plugins/copilot/tools'; -import { defineTool } from '../../plugins/copilot/tools/tool'; -import { - jsonOnlyPromptMessages, - nativeMessages, - nativeUserText, - promptMessages, - systemPrompt, - userPrompt, -} from './prompt-test-helper'; - -const mockDispatch = () => - (async function* (): AsyncIterableIterator { - yield { type: 'text_delta', text: 'Use [^1] now' }; - yield { type: 'citation', index: 1, url: 'https://affine.pro' }; - yield { type: 'done', finish_reason: 'stop' }; - })(); - -function stream( - factory: () => LlmToolLoopStreamEvent[] -): AsyncIterableIterator { - return (async function* () { - for (const event of factory()) { - yield event; - } - })(); -} - -async function collectChunks(iterable: AsyncIterable) { - const chunks: T[] = []; - for await (const chunk of iterable) { - chunks.push(chunk); - } - return chunks; -} - -function structuredOptions( - schema: z.ZodTypeAny, - extra?: Record -) { - const { responseSchemaJson, schemaHash } = - buildStructuredResponseContract(schema); - return { - responseSchemaJson, - schemaHash, - ...extra, - }; -} - -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; -} - -function normalizeToolExecuteOptions( - signalOrOptions?: AbortSignal | CopilotToolExecuteOptions, - maybeMessages?: PromptMessage[] -): CopilotToolExecuteOptions { - if ( - signalOrOptions && - typeof signalOrOptions === 'object' && - 'aborted' in signalOrOptions - ) { - return { - signal: signalOrOptions, - messages: maybeMessages, - }; - } - - if (!signalOrOptions) { - return maybeMessages ? { messages: maybeMessages } : {}; - } - - return { - ...signalOrOptions, - signal: signalOrOptions.signal, - messages: signalOrOptions.messages ?? maybeMessages, - }; -} - -function createTestToolLoopBridge( - dispatch: ( - request: LlmRequest, - signal?: AbortSignal - ) => AsyncIterableIterator, - tools: CopilotToolSet, - maxSteps = 20 -) { - return async function* ( - request: LlmRequest, - signalOrOptions?: AbortSignal | CopilotToolExecuteOptions, - maybeMessages?: PromptMessage[] - ): AsyncIterableIterator { - const toolExecuteOptions = normalizeToolExecuteOptions( - signalOrOptions, - maybeMessages - ); - const execute = createToolExecutionCallback(tools, toolExecuteOptions); - const messages = request.messages.map(message => ({ - ...message, - content: [...message.content], - })); - - for (let step = 0; step < maxSteps; step++) { - const toolCalls: Array< - Extract - > = []; - let finalDone: Extract | null = - null; - - for await (const event of dispatch( - { ...request, stream: true, messages }, - toolExecuteOptions.signal - )) { - if (event.type === 'tool_call') { - toolCalls.push(event); - yield event; - continue; - } - if (event.type === 'done') { - finalDone = event; - continue; - } - if (event.type === 'error') { - throw new Error(event.message); - } - yield event; - } - - if (!toolCalls.length) { - if (finalDone) { - yield finalDone; - } - return; - } - - if (step === maxSteps - 1) { - throw new Error('ToolCallLoop max steps reached'); - } - - messages.push({ - role: 'assistant', - content: toolCalls.map(call => ({ - type: 'tool_call', - call_id: call.call_id, - name: call.name, - arguments: call.arguments, - arguments_text: call.arguments_text, - arguments_error: call.arguments_error, - thought: call.thought, - })), - }); - - for (const call of toolCalls) { - const result = await execute({ - callId: call.call_id, - name: call.name, - args: call.arguments as Record, - rawArgumentsText: call.arguments_text, - argumentParseError: call.arguments_error, - }); - messages.push({ - role: 'tool', - content: [ - { - 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, - }, - ], - }); - yield { - 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, - }; - } - } - }; -} - -function installNativeDispatchRecorder( - owner: Partial<{ - structuredRequests: LlmStructuredRequest[]; - structuredFactory: (request: LlmStructuredRequest) => LlmStructuredResponse; - embeddingRequests: LlmEmbeddingRequest[]; - embeddingFactory: (request: LlmEmbeddingRequest) => { - model: string; - embeddings: number[][]; - usage?: { - prompt_tokens: number; - total_tokens: number; - }; - }; - rerankRequests: LlmRerankRequest[]; - rerankFactory: (request: LlmRerankRequest) => { - model: string; - scores: number[]; - }; - }> -) { - const originalStructured = (serverNativeModule as any).llmStructuredDispatch; - const originalEmbedding = (serverNativeModule as any).llmEmbeddingDispatch; - const originalRerank = (serverNativeModule as any).llmRerankDispatch; - - if (owner.structuredRequests && owner.structuredFactory) { - (serverNativeModule as any).llmStructuredDispatch = ( - _protocol: string, - _backendConfigJson: string, - requestJson: string - ) => { - const request = JSON.parse(requestJson) as LlmStructuredRequest; - owner.structuredRequests!.push(request); - return JSON.stringify(owner.structuredFactory!(request)); - }; - } - - if (owner.embeddingRequests && owner.embeddingFactory) { - (serverNativeModule as any).llmEmbeddingDispatch = ( - _protocol: string, - _backendConfigJson: string, - requestJson: string - ) => { - const request = JSON.parse(requestJson) as LlmEmbeddingRequest; - owner.embeddingRequests!.push(request); - return JSON.stringify(owner.embeddingFactory!(request)); - }; - } - - if (owner.rerankRequests && owner.rerankFactory) { - (serverNativeModule as any).llmRerankDispatch = ( - _protocol: string, - _backendConfigJson: string, - requestJson: string - ) => { - const request = JSON.parse(requestJson) as LlmRerankRequest; - owner.rerankRequests!.push(request); - return JSON.stringify(owner.rerankFactory!(request)); - }; - } - - return () => { - (serverNativeModule as any).llmStructuredDispatch = originalStructured; - (serverNativeModule as any).llmEmbeddingDispatch = originalEmbedding; - (serverNativeModule as any).llmRerankDispatch = originalRerank; - }; -} - -function installRemoteAttachmentMaterializer(owner: { - remoteAttachmentRequests: string[]; - remoteAttachmentSignals: Array; - remoteAttachmentResponses: Map; -}) { - return { - fetchRemoteAttachment: async ( - url: string, - options: { signal?: AbortSignal } - ) => { - owner.remoteAttachmentRequests.push(url); - owner.remoteAttachmentSignals.push(options.signal); - const response = owner.remoteAttachmentResponses.get(url); - if (!response) { - throw new Error(`missing remote attachment stub for ${url}`); - } - return response; - }, - }; -} - -class TestGeminiProvider extends GeminiProvider<{ apiKey: string }> { - override readonly type = CopilotProviderType.Gemini; - readonly dispatchRequests: LlmRequest[] = []; - readonly structuredRequests: LlmStructuredRequest[] = []; - readonly embeddingRequests: LlmEmbeddingRequest[] = []; - readonly remoteAttachmentRequests: string[] = []; - readonly remoteAttachmentSignals: Array = []; - readonly retryDelays: number[] = []; - remoteAttachmentResponses = new Map< - string, - { data: string; mimeType: string } - >(); - testTools: CopilotToolSet = {}; - testMiddleware: ProviderMiddlewareConfig = { - rust: { - request: ['normalize_messages', 'tool_schema_rewrite'], - stream: ['stream_event_normalize', 'citation_indexing'], - }, - node: { - text: ['citation_footnote', 'callout'], - }, - }; - dispatchFactory: (request: LlmRequest) => LlmToolLoopStreamEvent[] = () => [ - { type: 'text_delta', text: 'native' }, - { type: 'done', finish_reason: 'stop' }, - ]; - structuredFactory: (request: LlmStructuredRequest) => LlmStructuredResponse = - () => ({ - id: 'structured_1', - model: 'gemini-3.6-flash', - output_text: '{"summary":"AFFiNE native"}', - output_json: { summary: 'AFFiNE native' }, - usage: { - prompt_tokens: 4, - completion_tokens: 3, - total_tokens: 7, - }, - finish_reason: 'stop', - }); - embeddingFactory: (request: LlmEmbeddingRequest) => { - model: string; - embeddings: number[][]; - usage?: { - prompt_tokens: number; - total_tokens: number; - }; - } = request => ({ - model: request.model, - embeddings: request.inputs.map((_, index) => [index + 0.1, index + 0.2]), - usage: { - prompt_tokens: request.inputs.length, - total_tokens: request.inputs.length, - }, - }); - protected override readonly attachmentMaterializer = - installRemoteAttachmentMaterializer(this) as any; - - override configured() { - return true; - } - - protected override async createNativeConfig(): Promise { - return { - base_url: 'https://generativelanguage.googleapis.com/v1beta', - auth_token: 'api-key', - request_layer: 'gemini_api', - }; - } - - private createTestDispatch(_backendConfig: LlmBackendConfig) { - return (request: LlmRequest) => { - this.dispatchRequests.push(request); - return stream(() => this.dispatchFactory(request)); - }; - } - - override createNativeAdapter( - backend: ToolLoopBackend, - tools: CopilotToolSet, - nodeTextMiddleware?: NodeTextMiddleware[] - ) { - if (!('backendConfig' in backend)) { - throw new Error('expected direct backend config for test adapter'); - } - return new NativeProviderAdapter( - createTestToolLoopBridge( - this.createTestDispatch(backend.backendConfig), - tools, - this.MAX_STEPS - ), - { nodeTextMiddleware } - ); - } - - protected override async waitForStructuredRetry(delayMs: number) { - this.retryDelays.push(delayMs); - } - - override getActiveProviderMiddleware(): ProviderMiddlewareConfig { - return this.testMiddleware; - } - - override async getTools(): Promise { - return this.testTools; - } -} - -class TestGeminiVertexProvider extends GeminiVertexProvider { - testConfig = { - location: 'us-central1', - project: 'p1', - googleAuthOptions: {}, - } as any; - readonly dispatchRequests: LlmRequest[] = []; - readonly remoteAttachmentRequests: string[] = []; - readonly remoteAttachmentSignals: Array = []; - remoteAttachmentResponses = new Map< - string, - { data: string; mimeType: string } - >(); - testTools: CopilotToolSet = {}; - testMiddleware: ProviderMiddlewareConfig = { - rust: { - request: ['normalize_messages', 'tool_schema_rewrite'], - stream: ['stream_event_normalize', 'citation_indexing'], - }, - node: { - text: ['citation_footnote', 'callout'], - }, - }; - protected override readonly attachmentMaterializer = - installRemoteAttachmentMaterializer(this) as any; - - override get config() { - return this.testConfig; - } - - override configured() { - return true; - } - - protected override async resolveVertexAuth() { - return { - baseUrl: 'https://vertex.example', - headers: () => ({ - Authorization: 'Bearer vertex-token', - 'x-goog-user-project': 'p1', - }), - fetch: undefined, - } as const; - } - - private createTestDispatch(_backendConfig: LlmBackendConfig) { - return (request: LlmRequest) => { - this.dispatchRequests.push(request); - return stream(() => [ - { type: 'text_delta', text: 'vertex native' }, - { type: 'done', finish_reason: 'stop' }, - ]); - }; - } - - // oxlint-disable-next-line sonarjs/no-identical-functions - override createNativeAdapter( - backend: ToolLoopBackend, - tools: CopilotToolSet, - nodeTextMiddleware?: NodeTextMiddleware[] - ) { - if (!('backendConfig' in backend)) { - throw new Error('expected direct backend config for test adapter'); - } - return new NativeProviderAdapter( - createTestToolLoopBridge( - this.createTestDispatch(backend.backendConfig), - tools, - this.MAX_STEPS - ), - { nodeTextMiddleware } - ); - } - - override getActiveProviderMiddleware(): ProviderMiddlewareConfig { - return this.testMiddleware; - } - - override async getTools(): Promise { - return this.testTools; - } - - async exposeNativeConfig() { - return await this.createNativeConfig(); - } -} - -class TestOpenAIProvider extends OpenAIProvider { - readonly structuredRequests: LlmStructuredRequest[] = []; - readonly embeddingRequests: LlmEmbeddingRequest[] = []; - readonly rerankRequests: LlmRerankRequest[] = []; - structuredFactory: (request: LlmStructuredRequest) => LlmStructuredResponse = - request => ({ - id: 'structured_openai_1', - model: request.model, - output_text: '{"summary":"AFFiNE structured"}', - output_json: { summary: 'AFFiNE structured' }, - usage: { - prompt_tokens: 4, - completion_tokens: 3, - total_tokens: 7, - }, - finish_reason: 'stop', - }); - embeddingFactory: (request: LlmEmbeddingRequest) => { - model: string; - embeddings: number[][]; - usage?: { - prompt_tokens: number; - total_tokens: number; - }; - } = request => ({ - model: request.model, - embeddings: request.inputs.map(() => [0.4, 0.5]), - usage: { - prompt_tokens: request.inputs.length, - total_tokens: request.inputs.length, - }, - }); - rerankFactory: (request: LlmRerankRequest) => { - model: string; - scores: number[]; - } = request => ({ - model: request.model, - scores: request.candidates.map(() => 0.8), - }); - testMiddleware: ProviderMiddlewareConfig = { - rust: { - request: ['normalize_messages', 'tool_schema_rewrite'], - stream: ['stream_event_normalize', 'citation_indexing'], - }, - }; - - override get config() { - return { - apiKey: 'openai-key', - baseURL: 'https://api.openai.com/v1', - }; - } - - override configured() { - return true; - } - - override getActiveProviderMiddleware(): ProviderMiddlewareConfig { - return this.testMiddleware; - } -} - -test('NativeProviderAdapter should append citation and attachment footnotes', async t => { - const dispatch = () => - (async function* (): AsyncIterableIterator { - yield { - type: 'tool_result', - call_id: 'call_1', - name: 'blob_read', - arguments: { blob_id: 'blob_1' }, - output: { - blobId: 'blob_1', - fileName: 'a.txt', - fileType: 'text/plain', - content: 'A', - }, - }; - yield { - type: 'tool_result', - call_id: 'call_2', - name: 'blob_read', - arguments: { blob_id: 'blob_2' }, - output: { - blobId: 'blob_2', - fileName: 'b.txt', - fileType: 'text/plain', - content: 'B', - }, - }; - yield { type: 'text_delta', text: 'Answer from files.' }; - yield { type: 'done', finish_reason: 'stop' }; - })(); - const dispatchWithModelReference = () => - (async function* (): AsyncIterableIterator { - yield { - type: 'tool_result', - call_id: 'call_1', - name: 'doc_semantic_search', - arguments: { query: 'A' }, - output: [ - { - blobId: 'blob_1', - name: 'a.txt', - mimeType: 'text/plain', - content: 'A', - }, - ], - }; - yield { type: 'text_delta', text: 'Answer from file.[^1]' }; - yield { type: 'done', finish_reason: 'stop' }; - })(); - - const cases = [ - { - title: 'streamText citation footnotes', - run: async () => { - const adapter = new NativeProviderAdapter( - createTestToolLoopBridge(mockDispatch, {}, 3) - ); - return ( - await collectChunks( - adapter.streamText({ - model: 'gpt-5-mini', - stream: true, - messages: nativeMessages(nativeUserText('hi')), - }) - ) - ).join(''); - }, - verify: (text: string) => { - t.true(text.includes('Use [^1] now')); - t.true( - text.includes('[^1]: {"type":"url","url":"https%3A%2F%2Faffine.pro"}') - ); - }, - }, - { - title: 'streamObject citation footnotes', - run: async () => { - const adapter = new NativeProviderAdapter( - createTestToolLoopBridge(mockDispatch, {}, 3) - ); - const chunks = await collectChunks( - adapter.streamObject({ - model: 'gpt-5-mini', - stream: true, - messages: nativeMessages(nativeUserText('hi')), - }) - ); - t.deepEqual( - chunks.map(chunk => chunk.type), - ['text-delta', 'text-delta'], - 'streamObject citation chunk types' - ); - return chunks - .filter(chunk => chunk.type === 'text-delta') - .map(chunk => chunk.textDelta) - .join(''); - }, - verify: (text: string) => { - t.true(text.includes('Use [^1] now')); - t.true( - text.includes('[^1]: {"type":"url","url":"https%3A%2F%2Faffine.pro"}') - ); - }, - }, - { - title: 'streamObject attachment footnotes', - run: async () => { - const adapter = new NativeProviderAdapter( - createTestToolLoopBridge(dispatch, {}, 3) - ); - const chunks = await collectChunks( - adapter.streamObject({ - model: 'gpt-5-mini', - stream: true, - messages: nativeMessages(nativeUserText('hi')), - }) - ); - return chunks - .filter(chunk => chunk.type === 'text-delta') - .map(chunk => chunk.textDelta) - .join(''); - }, - verify: (text: string) => { - t.true(text.includes('Answer from files.')); - t.true(text.includes('[^1][^2]')); - t.true( - text.includes( - '[^1]: {"type":"attachment","blobId":"blob_1","fileName":"a.txt","fileType":"text/plain"}' - ) - ); - t.true( - text.includes( - '[^2]: {"type":"attachment","blobId":"blob_2","fileName":"b.txt","fileType":"text/plain"}' - ) - ); - }, - }, - { - title: 'streamObject attachment definitions for model references', - run: async () => { - const adapter = new NativeProviderAdapter( - createTestToolLoopBridge(dispatchWithModelReference, {}, 3) - ); - const chunks = await collectChunks( - adapter.streamObject({ - model: 'gpt-5-mini', - stream: true, - messages: nativeMessages(nativeUserText('hi')), - }) - ); - return chunks - .filter(chunk => chunk.type === 'text-delta') - .map(chunk => chunk.textDelta) - .join(''); - }, - verify: (text: string) => { - t.true(text.includes('Answer from file.[^1]')); - t.true( - text.includes( - '[^1]: {"type":"attachment","blobId":"blob_1","fileName":"a.txt","fileType":"text/plain"}' - ) - ); - }, - }, - ] as const; - - for (const testCase of cases) { - testCase.verify(await testCase.run()); - } -}); - -test('NativeProviderAdapter streamObject should map tool and text events', async t => { - let round = 0; - const dispatch = (_request: LlmRequest) => - (async function* (): AsyncIterableIterator { - round += 1; - if (round === 1) { - yield { - type: 'tool_call', - call_id: 'call_1', - name: 'doc_read', - arguments: { doc_id: 'a1' }, - }; - yield { type: 'done', finish_reason: 'tool_calls' }; - return; - } - yield { type: 'text_delta', text: 'ok' }; - yield { type: 'done', finish_reason: 'stop' }; - })(); - - const adapter = new NativeProviderAdapter( - createTestToolLoopBridge( - dispatch, - { - doc_read: { - inputSchema: z.object({ doc_id: z.string() }), - execute: async () => ({ markdown: '# a1' }), - }, - }, - 4 - ) - ); - - const events = []; - for await (const event of adapter.streamObject({ - model: 'gpt-5-mini', - stream: true, - messages: nativeMessages(nativeUserText('read')), - })) { - events.push(event); - } - - t.deepEqual( - events.map(event => event.type), - ['tool-call', 'tool-result', 'text-delta'] - ); - t.snapshot(events); -}); - -test('NativeProviderAdapter streamObject should finalize usage with selected provider', async t => { - const usageEvents: Array<{ - providerId: string; - model?: string; - usage?: { - prompt_tokens: number; - completion_tokens: number; - total_tokens: number; - cached_tokens?: number; - }; - }> = []; - const adapter = new NativeProviderAdapter( - () => - stream(() => [ - { type: 'message_start', model: 'gpt-5-mini' }, - { type: 'text_delta', text: 'ok' }, - { - type: 'done', - finish_reason: 'stop', - usage: { prompt_tokens: 2, completion_tokens: 3, total_tokens: 5 }, - }, - { - type: 'provider_selected', - provider_id: 'byok-aaaaaaaaaaaa-openai-server-key1', - }, - ]), - { - onUsage: input => { - usageEvents.push(input); - }, - } - ); - - const events = await collectChunks( - adapter.streamObject({ - model: 'gpt-5-mini', - stream: true, - messages: nativeMessages(nativeUserText('hi')), - }) - ); - - t.deepEqual(events, [{ type: 'text-delta', textDelta: 'ok' }]); - t.deepEqual(usageEvents, [ - { - providerId: 'byok-aaaaaaaaaaaa-openai-server-key1', - model: 'gpt-5-mini', - usage: { prompt_tokens: 2, completion_tokens: 3, total_tokens: 5 }, - }, - ]); -}); - -test('NativeProviderAdapter streamObject should keep streaming when usage callback fails', async t => { - const adapter = new NativeProviderAdapter( - () => - stream(() => [ - { type: 'message_start', model: 'gpt-5-mini' }, - { type: 'text_delta', text: 'ok' }, - { - type: 'done', - finish_reason: 'stop', - usage: { prompt_tokens: 2, completion_tokens: 3, total_tokens: 5 }, - }, - { - type: 'provider_selected', - provider_id: 'byok-aaaaaaaaaaaa-openai-server-key1', - }, - ]), - { - onUsage: () => { - throw new Error('usage callback failed'); - }, - } - ); - - const events = await collectChunks( - adapter.streamObject({ - model: 'gpt-5-mini', - stream: true, - messages: nativeMessages(nativeUserText('hi')), - }) - ); - - t.deepEqual(events, [{ type: 'text-delta', textDelta: 'ok' }]); -}); - -test('NativeRuntimeAdapter streamObject should keep raw runtime stream objects only', async t => { - const adapter = new NativeRuntimeAdapter( - createTestToolLoopBridge(mockDispatch, {}, 3) - ); - - const chunks: StreamObject[] = []; - for await (const chunk of adapter.streamObject({ - model: 'gpt-5-mini', - stream: true, - messages: nativeMessages(nativeUserText('hi')), - })) { - chunks.push(chunk); - } - - t.deepEqual(chunks, [{ type: 'text-delta', textDelta: 'Use [^1] now' }]); -}); - -test('structured response contract helpers should normalize explicit fields only', t => { - const schemaJson = { - type: 'object', - properties: { - summary: { type: 'string' }, - }, - required: ['summary'], - additionalProperties: false, - }; - const reorderedSchemaJson = { - required: ['summary'], - additionalProperties: false, - properties: { - summary: { type: 'string' }, - }, - type: 'object', - }; - - const explicit = buildPromptStructuredResponseFromFields({ - responseSchemaJson: schemaJson, - }); - const reordered = buildPromptStructuredResponseFromFields({ - responseSchemaJson: reorderedSchemaJson, - }); - - t.truthy(explicit); - t.is(explicit?.schemaHash, reordered?.schemaHash); - t.deepEqual(explicit, { - responseSchemaJson: schemaJson, - schemaHash: reordered?.schemaHash, - }); -}); - -test('buildNativeRequest should include rust middleware from profile', async t => { - const { request } = await buildNativeRequest({ - model: 'gpt-5-mini', - messages: promptMessages(userPrompt('hello')), - toolContracts: [], - middleware: { - rust: { - request: ['normalize_messages', 'clamp_max_tokens'], - stream: ['stream_event_normalize', 'citation_indexing'], - }, - node: { - text: ['callout'], - }, - }, - }); - - t.deepEqual(request.middleware, { - request: ['normalize_messages', 'clamp_max_tokens'], - stream: ['stream_event_normalize', 'citation_indexing'], - }); -}); - -test('buildCanonicalNativeRequest should only use explicit structured contract inputs', async t => { - const schema = z.object({ - summary: z.string(), - }); - - const { request } = await buildCanonicalNativeRequest({ - model: 'gpt-4.1', - messages: promptMessages( - systemPrompt('Return valid JSON.'), - userPrompt('Summarize AFFiNE.') - ), - responseContract: buildStructuredResponseContract(schema), - }); - - t.snapshot(request.responseSchema); -}); - -test('buildCanonicalNativeStructuredRequest should accept schema-only explicit structured response contracts', async t => { - const { request } = await buildCanonicalNativeStructuredRequest({ - model: 'gpt-4.1', - messages: [ - { - role: 'system', - content: 'Return JSON only.', - responseFormat: { - type: 'json_schema', - responseSchemaJson: { - type: 'object', - properties: { summary: { type: 'string' } }, - required: ['summary'], - additionalProperties: false, - }, - strict: false, - }, - }, - { role: 'user', content: 'Summarize AFFiNE.' }, - ], - responseContract: { - responseSchemaJson: { - type: 'object', - properties: { summary: { type: 'string' } }, - required: ['summary'], - additionalProperties: false, - }, - }, - }); - - t.snapshot({ - schema: request.schema, - strict: request.strict, - }); -}); - -test('buildCanonicalNativeStructuredRequest should honor explicit structured options contract before system responseFormat', async t => { - const responseContract = buildPromptStructuredResponseFromFields({ - responseSchemaJson: { - type: 'object', - properties: { ok: { type: 'boolean' } }, - required: ['ok'], - additionalProperties: false, - }, - schemaHash: 'ok-v1', - strict: true, - }); - const { request } = await buildCanonicalNativeStructuredRequest({ - model: 'gpt-4.1', - messages: [ - { - role: 'system', - content: 'Return JSON only.', - responseFormat: { - type: 'json_schema', - responseSchemaJson: { - type: 'object', - properties: { summary: { type: 'string' } }, - required: ['summary'], - additionalProperties: false, - }, - strict: false, - }, - }, - { role: 'user', content: 'Summarize AFFiNE.' }, - ], - options: { - responseSchemaJson: { - type: 'object', - properties: { ok: { type: 'boolean' } }, - required: ['ok'], - additionalProperties: false, - }, - schemaHash: 'ok-v1', - strict: true, - }, - responseContract: responseContract!, - }); - - t.snapshot({ - schema: request.schema, - strict: request.strict, - }); -}); - -test('buildCanonicalNativeStructuredRequest should honor explicit responseSchema for array outputs', async t => { - const schema = z.array(z.object({ speaker: z.string(), text: z.string() })); - const { request } = await buildCanonicalNativeStructuredRequest({ - model: 'gemini-3.6-flash', - messages: jsonOnlyPromptMessages('Transcribe this audio.'), - options: {}, - responseContract: buildStructuredResponseContract(schema), - }); - - t.snapshot(request.schema); -}); - -test('buildCanonicalNativeStructuredRequest should consume explicit structured response contract without options.schema', async t => { - const schema = z.object({ summary: z.string() }); - const responseContract = buildStructuredResponseContract(schema); - const { request } = await buildCanonicalNativeStructuredRequest({ - model: 'gemini-3.6-flash', - messages: jsonOnlyPromptMessages('Summarize AFFiNE.'), - options: { strict: false }, - responseContract, - }); - - t.snapshot({ schema: request.schema, strict: request.strict }); -}); - -test('buildCanonicalNativeStructuredRequest should accept explicit schema contracts without schemaHash', async t => { - const { request } = await buildCanonicalNativeStructuredRequest({ - model: 'gpt-4.1', - messages: jsonOnlyPromptMessages('Summarize AFFiNE.'), - responseContract: { - responseSchemaJson: { - type: 'object', - properties: { summary: { type: 'string' } }, - required: ['summary'], - additionalProperties: false, - }, - }, - }); - - t.snapshot({ - schema: request.schema, - strict: request.strict, - }); -}); - -test('buildNativeRequest should canonicalize Gemini attachments', async t => { - const cases: Array<{ - title: string; - input: Parameters[0]; - }> = [ - { - title: 'remote file url', - input: { - model: 'gemini-3.6-flash', - messages: [ - { - role: 'user' as const, - content: 'summarize this attachment', - attachments: ['https://example.com/a.pdf'], - params: { mimetype: 'application/pdf' }, - }, - ], - }, - }, - { - title: 'remote image url', - input: { - model: 'gemini-3.6-flash', - messages: [ - { - role: 'user' as const, - content: 'describe this image', - attachments: ['https://example.com/cat.png'], - }, - ], - }, - }, - { - title: 'data url', - input: { - model: 'gemini-3.6-flash', - messages: [ - { - role: 'user' as const, - content: 'read this note', - attachments: ['data:text/plain,hello%20world'], - params: { mimetype: 'text/plain' }, - }, - ], - }, - }, - { - title: 'remote audio url', - input: { - model: 'gemini-3.6-flash', - messages: [ - { - role: 'user' as const, - content: 'transcribe this clip', - attachments: ['https://example.com/a.mp3'], - params: { mimetype: 'audio/mpeg' }, - }, - ], - }, - }, - { - title: 'bytes and file handle', - input: { - model: 'gemini-3.6-flash', - messages: [ - { - role: 'user' as const, - content: 'inspect these assets', - attachments: [ - { - kind: 'bytes' as const, - data: Buffer.from('hello', 'utf8').toString('base64'), - mimeType: 'text/plain', - fileName: 'hello.txt', - }, - { - kind: 'file_handle' as const, - fileHandle: 'file_123', - mimeType: 'application/pdf', - fileName: 'report.pdf', - }, - ], - }, - ], - attachmentCapability: { - kinds: ['image', 'audio', 'file'], - sourceKinds: ['bytes', 'file_handle'], - }, - }, - }, - ]; - - for (const testCase of cases) { - const { request } = await buildNativeRequest(testCase.input); - t.snapshot(request.messages[0]?.content, testCase.title); - } -}); - -test('buildNativeRequest should reject attachments outside native admission matrix', async t => { - const error = await t.throwsAsync( - buildNativeRequest({ - model: 'gpt-4o', - messages: [ - { - role: 'user', - content: 'summarize this attachment', - attachments: ['https://example.com/a.pdf'], - params: { mimetype: 'application/pdf' }, - }, - ], - attachmentCapability: { - kinds: ['image'], - sourceKinds: ['url', 'data'], - allowRemoteUrls: true, - }, - }) - ); - - t.true(error instanceof CopilotPromptInvalid); - t.regex(error.message, /does not support file attachments/i); -}); - -test('buildNativeStructuredRequest should prefer explicit schema option', async t => { - const provider = new TestOpenAIProvider(); - t.teardown(installNativeDispatchRecorder(provider)); - const schema = z.object({ summary: z.string() }); - - await getProviderRuntimeHost(provider).run.structured( - { modelId: 'gpt-4.1' }, - jsonOnlyPromptMessages('Summarize AFFiNE in one sentence.'), - structuredOptions(schema), - structuredContract(schema) - ); - - t.snapshot(provider.structuredRequests[0]?.schema); -}); - -test('buildNativeStructuredRequest should preserve caller strictness override', async t => { - const provider = new TestOpenAIProvider(); - t.teardown(installNativeDispatchRecorder(provider)); - - await getProviderRuntimeHost(provider).run.structured( - { modelId: 'gpt-4.1' }, - jsonOnlyPromptMessages('Summarize AFFiNE in one sentence.'), - structuredOptions(z.object({ summary: z.string() }), { strict: false }), - structuredContract(z.object({ summary: z.string() })) - ); - - t.is(provider.structuredRequests[0]?.strict, false); -}); - -test('buildNativeStructuredRequest should ignore legacy params.schema fallback when explicit schema contract exists', async t => { - const { request } = await buildNativeStructuredRequest({ - model: 'gpt-4.1', - messages: promptMessages( - systemPrompt('Return JSON only.', { - params: { - schema: z.object({ summary: z.string() }), - }, - }), - userPrompt('Summarize AFFiNE in one sentence.') - ), - responseContract: { - responseSchemaJson: { - type: 'object', - properties: { summary: { type: 'string' } }, - required: ['summary'], - additionalProperties: false, - }, - }, - }); - - t.snapshot({ - schema: request.schema, - strict: request.strict, - }); -}); - -test('buildNativeStructuredRequest should reject legacy options.schema fallback', async t => { - const provider = new TestOpenAIProvider(); - - const error = await t.throwsAsync(() => - getProviderRuntimeHost(provider).run.structured( - { modelId: 'gpt-4.1' }, - jsonOnlyPromptMessages('Summarize AFFiNE in one sentence.'), - { - schema: z.object({ summary: z.string() }), - } as never - ) - ); - - t.true(error instanceof CopilotPromptInvalid); - t.regex((error as Error).message, /Schema is required/); -}); - -test('buildNativeRequest should preserve tool schemas and defer Gemini rewrite to native request layer', async t => { - const schema = z.object({ - doc_id: z.string(), - options: z.object({ mode: z.enum(['full', 'summary']) }), - }); - - const [{ request: geminiRequest }, { request: openaiRequest }] = - await Promise.all([ - buildNativeRequest({ - model: 'gemini-3.6-flash', - messages: promptMessages(userPrompt('read doc')), - toolContracts: buildToolContracts({ - doc_read: defineTool({ - inputSchema: schema, - execute: async () => ({ markdown: '# doc' }), - }), - }), - }), - buildNativeRequest({ - model: 'gpt-4.1', - messages: promptMessages(userPrompt('read doc')), - toolContracts: buildToolContracts({ - doc_read: defineTool({ - inputSchema: schema, - execute: async () => ({ markdown: '# doc' }), - }), - }), - }), - ]); - - t.true( - JSON.stringify(geminiRequest.tools?.[0]?.parameters).includes( - 'additionalProperties' - ) - ); - t.true( - JSON.stringify(openaiRequest.tools?.[0]?.parameters).includes( - 'additionalProperties' - ) - ); -}); - -test('defineTool should precompute json schema at definition time', t => { - const tool = defineTool({ - description: 'Read a doc', - inputSchema: z.object({ - docId: z.string(), - includeChildren: z.boolean().optional(), - }), - execute: async () => ({ ok: true }), - }); - - t.snapshot(tool.jsonSchema); -}); - -test('buildNativeStructuredRequest should preserve schemas and defer Gemini rewrite to native request layer', async t => { - const schema = z.object({ - summary: z.string(), - metadata: z.object({ format: z.enum(['short', 'long']) }), - }); - - const [{ request: geminiRequest }, { request: openaiRequest }] = - await Promise.all([ - buildNativeStructuredRequest({ - model: 'gemini-3.6-flash', - messages: promptMessages(userPrompt('Summarize AFFiNE.')), - responseContract: buildStructuredResponseContract(schema), - }), - buildNativeStructuredRequest({ - model: 'gpt-4.1', - messages: promptMessages(userPrompt('Summarize AFFiNE.')), - responseContract: buildStructuredResponseContract(schema), - }), - ]); - - for (const [title, request] of [ - ['gemini', geminiRequest], - ['openai', openaiRequest], - ] as const) { - t.true( - JSON.stringify(request.schema).includes('additionalProperties'), - title - ); - } -}); - -test('NativeProviderAdapter streamText should skip citation footnotes when disabled', async t => { - const adapter = new NativeProviderAdapter( - createTestToolLoopBridge(mockDispatch, {}, 3), - { nodeTextMiddleware: ['callout'] } - ); - const chunks: string[] = []; - for await (const chunk of adapter.streamText({ - model: 'gpt-5-mini', - stream: true, - messages: nativeMessages(nativeUserText('hi')), - })) { - chunks.push(chunk); - } - - const text = chunks.join(''); - t.true(text.includes('Use [^1] now')); - t.false( - text.includes('[^1]: {"type":"url","url":"https%3A%2F%2Faffine.pro"}') - ); -}); - -test('GeminiProvider should use native path for text-only requests', async t => { - const provider = new TestGeminiProvider(); - - const result = await getProviderRuntimeHost(provider).run.text( - { modelId: 'gemini-3.6-flash' }, - promptMessages(userPrompt('hello')), - { reasoning: true } - ); - - t.is(result, 'native'); - t.is(provider.dispatchRequests.length, 1); - t.snapshot({ - remoteAttachmentRequests: provider.remoteAttachmentRequests, - include: provider.dispatchRequests[0]?.include, - reasoning: provider.dispatchRequests[0]?.reasoning, - middleware: provider.dispatchRequests[0]?.middleware, - }); -}); - -test('GeminiProvider should use native path for structured requests', async t => { - const provider = new TestGeminiProvider(); - t.teardown(installNativeDispatchRecorder(provider)); - - const schema = z.object({ summary: z.string() }); - const result = await getProviderRuntimeHost(provider).run.structured( - { modelId: 'gemini-3.6-flash' }, - jsonOnlyPromptMessages('Summarize AFFiNE in one short sentence.'), - structuredOptions(schema), - structuredContract(schema) - ); - - t.is(provider.structuredRequests.length, 1); - t.snapshot({ - request: provider.structuredRequests[0], - result: JSON.parse(result), - }); -}); - -test('GeminiProvider should retry when native structured dispatch returns invalid_structured_output', async t => { - const provider = new TestGeminiProvider(); - t.teardown(installNativeDispatchRecorder(provider)); - let attempts = 0; - provider.structuredFactory = () => { - attempts += 1; - if (attempts === 1) { - throw Object.assign( - new Error( - 'structured response did not contain valid JSON: summary: missing' - ), - { code: 'invalid_structured_output' as const } - ); - } - return { - id: `structured_retry_${attempts}`, - model: 'gemini-3.6-flash', - output_text: '{"summary":"ok"}', - output_json: { summary: 'ok' }, - usage: { - prompt_tokens: 4, - completion_tokens: 3, - total_tokens: 7, - }, - finish_reason: 'stop', - }; - }; - - const result = await getProviderRuntimeHost(provider).run.structured( - { modelId: 'gemini-3.6-flash' }, - jsonOnlyPromptMessages('Summarize AFFiNE in one short sentence.'), - structuredOptions(z.object({ summary: z.string() }), { maxRetries: 2 }), - structuredContract(z.object({ summary: z.string() })) - ); - - t.is(attempts, 2); - t.deepEqual(JSON.parse(result), { summary: 'ok' }); -}); - -test('GeminiProvider should treat maxRetries as retry count for backend failures', async t => { - const provider = new TestGeminiProvider(); - t.teardown(installNativeDispatchRecorder(provider)); - let attempts = 0; - provider.structuredFactory = () => { - attempts += 1; - throw new Error('backend down'); - }; - - const error = await t.throwsAsync( - getProviderRuntimeHost(provider).run.structured( - { modelId: 'gemini-3.6-flash' }, - jsonOnlyPromptMessages('Summarize AFFiNE in one short sentence.'), - structuredOptions(z.object({ summary: z.string() }), { maxRetries: 2 }), - structuredContract(z.object({ summary: z.string() })) - ) - ); - - t.is(attempts, 3); - t.deepEqual(provider.retryDelays, [2_000, 4_000]); - t.regex(error.message, /backend down/); -}); - -test('GeminiProvider should use native structured path for audio attachments', async t => { - const provider = new TestGeminiProvider(); - t.teardown(installNativeDispatchRecorder(provider)); - const inlineData = Buffer.from('audio-bytes', 'utf8').toString('base64'); - provider.remoteAttachmentResponses.set('https://example.com/a.mp3', { - data: inlineData, - mimeType: 'audio/mpeg', - }); - provider.structuredFactory = () => ({ - id: 'structured_audio_1', - model: 'gemini-3.6-flash', - output_text: '[{"a":"Speaker 1","s":0,"e":1,"t":"Hello"}]', - output_json: [{ a: 'Speaker 1', s: 0, e: 1, t: 'Hello' }], - usage: { prompt_tokens: 4, completion_tokens: 3, total_tokens: 7 }, - finish_reason: 'stop', - }); - - const result = await getProviderRuntimeHost(provider).run.structured( - { modelId: 'gemini-3.6-flash' }, - promptMessages( - systemPrompt('Return JSON only.'), - userPrompt('transcribe the audio', { - attachments: ['https://example.com/a.mp3'], - params: { mimetype: 'audio/mpeg' }, - }) - ), - structuredOptions( - z.array( - z.object({ a: z.string(), s: z.number(), e: z.number(), t: z.string() }) - ) - ), - structuredContract( - z.array( - z.object({ a: z.string(), s: z.number(), e: z.number(), t: z.string() }) - ) - ) - ); - - t.is(provider.structuredRequests.length, 1); - t.snapshot({ - content: provider.structuredRequests[0]?.messages[1]?.content, - remoteAttachmentRequests: provider.remoteAttachmentRequests, - result: JSON.parse(result), - }); -}); - -test('GeminiProvider should use native path for embeddings', async t => { - const provider = new TestGeminiProvider(); - t.teardown(installNativeDispatchRecorder(provider)); - - const result = await getProviderRuntimeHost(provider).run.embedding( - { modelId: 'gemini-embedding-001' }, - ['first', 'second'], - { dimensions: 3 } - ); - - t.is(provider.embeddingRequests.length, 1); - t.snapshot({ result, request: provider.embeddingRequests[0] }); -}); - -test('GeminiProvider should canonicalize native text attachments', async t => { - const cases = [ - { - title: 'remote file attachment', - setup(provider: TestGeminiProvider) { - const inlineData = Buffer.from('pdf-bytes', 'utf8').toString('base64'); - provider.remoteAttachmentResponses.set('https://example.com/a.pdf', { - data: inlineData, - mimeType: 'application/pdf', - }); - }, - messages: [ - { - role: 'user' as const, - content: 'summarize this file', - attachments: ['https://example.com/a.pdf'], - params: { mimetype: 'application/pdf' }, - }, - ] satisfies PromptMessage[], - }, - { - title: 'remote image attachment', - setup(provider: TestGeminiProvider) { - const inlineData = Buffer.from('image-bytes', 'utf8').toString( - 'base64' - ); - provider.remoteAttachmentResponses.set('https://example.com/a.jpg', { - data: inlineData, - mimeType: 'image/jpeg', - }); - }, - messages: [ - { - role: 'user' as const, - content: 'describe this image', - attachments: ['https://example.com/a.jpg'], - }, - ] satisfies PromptMessage[], - }, - { - title: 'downloaded audio webm attachment', - setup(provider: TestGeminiProvider) { - const inlineData = Buffer.from('audio-bytes', 'utf8').toString( - 'base64' - ); - provider.remoteAttachmentResponses.set('https://example.com/a.webm', { - data: inlineData, - mimeType: 'audio/webm', - }); - }, - messages: [ - { - role: 'user' as const, - content: 'transcribe this clip', - attachments: ['https://example.com/a.webm'], - }, - ] satisfies PromptMessage[], - }, - { - title: 'google file url attachment', - setup() {}, - messages: [ - { - role: 'user' as const, - content: 'summarize this file', - attachments: [ - 'https://generativelanguage.googleapis.com/v1beta/files/file-123', - ], - params: { mimetype: 'application/pdf' }, - }, - ] satisfies PromptMessage[], - }, - ] as const; - - for (const testCase of cases) { - const provider = new TestGeminiProvider(); - testCase.setup(provider); - - const result = await getProviderRuntimeHost(provider).run.text( - { modelId: 'gemini-3.6-flash' }, - testCase.messages - ); - - t.is(result, 'native', testCase.title); - t.snapshot( - { - remoteAttachmentRequests: provider.remoteAttachmentRequests, - content: provider.dispatchRequests[0]?.messages[0]?.content, - }, - testCase.title - ); - } -}); - -test('GeminiProvider should pass abort signal to remote attachment prefetch', async t => { - const provider = new TestGeminiProvider(); - provider.remoteAttachmentResponses.set('https://example.com/a.jpg', { - data: Buffer.from('image-bytes', 'utf8').toString('base64'), - mimeType: 'image/jpeg', - }); - const controller = new AbortController(); - - await getProviderRuntimeHost(provider).run.text( - { modelId: 'gemini-3.6-flash' }, - [ - { - role: 'user', - content: 'describe this image', - attachments: ['https://example.com/a.jpg'], - }, - ], - { signal: controller.signal } - ); - - t.deepEqual(provider.remoteAttachmentRequests, ['https://example.com/a.jpg']); - t.is(provider.remoteAttachmentSignals[0], controller.signal); -}); - -test('GeminiProvider should not pass materialized inline attachment URL to native request', async t => { - const provider = new TestGeminiProvider(); - const inlineData = Buffer.from('image-bytes', 'utf8').toString('base64'); - provider.remoteAttachmentResponses.set('https://example.com/a.jpg', { - data: inlineData, - mimeType: 'image/jpeg', - }); - - await getProviderRuntimeHost(provider).run.text( - { modelId: 'gemini-3.6-flash' }, - [ - { - role: 'user', - content: 'describe this image', - attachments: ['https://example.com/a.jpg'], - }, - ], - { - user: 'user-1', - workspace: 'workspace-1', - session: 'session-1', - } - ); - - const content = provider.dispatchRequests[0]?.messages[0]?.content as Array<{ - type: string; - source?: Record; - }>; - const attachmentPart = content.find(part => part.type === 'image'); - - t.deepEqual(provider.remoteAttachmentRequests, ['https://example.com/a.jpg']); - t.is(attachmentPart?.source?.data, inlineData); - t.is(attachmentPart?.source?.media_type, 'image/jpeg'); - t.false('url' in (attachmentPart?.source ?? {})); -}); - -test('GeminiProvider should reject unsupported attachment schemes at input validation', async t => { - const provider = new TestGeminiProvider(); - - const error = await t.throwsAsync( - getProviderRuntimeHost(provider).run.text( - { modelId: 'gemini-3.6-flash' }, - [ - { - role: 'user', - content: 'read this attachment', - attachments: ['blob:https://example.com/file-id'], - params: { mimetype: 'application/pdf' }, - }, - ], - {} - ) - ); - - t.true(error instanceof CopilotPromptInvalid); - t.regex(error.message, /attachments must use https\?:\/\/, gs:\/\/ or data:/); - t.is(provider.dispatchRequests.length, 0); -}); - -test('GeminiProvider should validate malformed attachments before canonicalization', async t => { - const provider = new TestGeminiProvider(); - - const error = await t.throwsAsync( - getProviderRuntimeHost(provider).run.text( - { modelId: 'gemini-3.6-flash' }, - [ - { - role: 'user', - content: 'read this attachment', - attachments: [{ kind: 'url' }], - }, - ] as any, - {} - ) - ); - - t.true(error instanceof CopilotPromptInvalid); - t.regex(error.message, /attachments\[0\]/); - t.is(provider.dispatchRequests.length, 0); -}); - -test('GeminiProvider should drive tool loop on native path', async t => { - const provider = new TestGeminiProvider(); - provider.testTools = { - doc_read: defineTool({ - inputSchema: z.object({ doc_id: z.string() }), - execute: async args => ({ markdown: `# ${(args as any).doc_id}` }), - }), - }; - provider.dispatchFactory = request => { - const hasToolResult = request.messages.some( - message => message.role === 'tool' - ); - if (!hasToolResult) { - return [ - { - type: 'tool_call', - call_id: 'call_1', - name: 'doc_read', - arguments: { doc_id: 'a1' }, - }, - { type: 'done', finish_reason: 'tool_calls' }, - ]; - } - - return [ - { type: 'text_delta', text: 'after tool' }, - { type: 'done', finish_reason: 'stop' }, - ]; - }; - - const result = await getProviderRuntimeHost(provider).run.text( - { modelId: 'gemini-3.6-flash' }, - [{ role: 'user', content: 'read doc a1' }], - {} - ); - - t.true(result.includes('after tool')); - t.is(provider.dispatchRequests.length, 2); - t.true( - provider.dispatchRequests[1]?.messages.some( - message => message.role === 'tool' - ) - ); -}); - -test('GeminiVertexProvider should prefetch bearer token for native config', async t => { - const provider = new TestGeminiVertexProvider(); - const config = await provider.exposeNativeConfig(); - t.snapshot(config); -}); - -test('GeminiVertexProvider should build project scoped Vertex base URL', t => { - t.is( - getVertexGoogleBaseUrl({ - project: 'p1', - location: 'us-central1', - googleAuthOptions: {}, - }), - 'https://us-central1-aiplatform.googleapis.com/v1/projects/p1/locations/us-central1/publishers/google' - ); -}); - -test('GeminiVertexProvider should materialize remote attachments before native text path', async t => { - const cases = [ - { - title: 'remote http url', - url: 'https://example.com/a.mp3', - data: Buffer.from('audio-bytes', 'utf8').toString('base64'), - mimeType: 'audio/mpeg', - }, - { - title: 'gs url', - url: 'gs://bucket/audio.opus', - data: Buffer.from('opus-bytes', 'utf8').toString('base64'), - mimeType: 'audio/opus', - }, - ] as const; - - for (const testCase of cases) { - const provider = new TestGeminiVertexProvider(); - provider.remoteAttachmentResponses.set(testCase.url, { - data: testCase.data, - mimeType: testCase.mimeType, - }); - - const result = await getProviderRuntimeHost(provider).run.text( - { modelId: 'gemini-3.6-flash' }, - [ - { - role: 'user', - content: 'transcribe the audio', - attachments: [testCase.url], - }, - ], - {} - ); - - t.is(result, 'vertex native', testCase.title); - t.snapshot( - { - remoteAttachmentRequests: provider.remoteAttachmentRequests, - content: provider.dispatchRequests[0]?.messages[0]?.content, - }, - testCase.title - ); - } -}); - -test('OpenAIProvider should use native structured dispatch', async t => { - const provider = new TestOpenAIProvider(); - t.teardown(installNativeDispatchRecorder(provider)); - const schema = z.object({ summary: z.string() }); - - const result = await getProviderRuntimeHost(provider).run.structured( - { modelId: 'gpt-4.1' }, - jsonOnlyPromptMessages('Summarize AFFiNE in one sentence.'), - structuredOptions(schema), - structuredContract(schema) - ); - - t.is(provider.structuredRequests.length, 1); - t.snapshot({ - result: JSON.parse(result), - request: provider.structuredRequests[0], - }); -}); - -test('parseNativeStructuredOutput should require native output_json', t => { - const error = t.throws(() => - parseNativeStructuredOutput({ - output_text: '{"summary":"AFFiNE"}', - }) - ); - - t.true(error instanceof Error); - const structuredError = error as Error & { - code?: string; - name?: string; - }; - t.is(structuredError.name, 'StructuredResponseParseError'); - t.is(structuredError.code, 'invalid_structured_output'); - t.regex(structuredError.message, /missing required output_json/); -}); - -test('OpenAIProvider should prefer native output_json for structured dispatch', async t => { - const provider = new TestOpenAIProvider(); - t.teardown(installNativeDispatchRecorder(provider)); - provider.structuredFactory = request => ({ - id: 'structured_openai_output_json', - model: request.model, - output_text: 'not-json-anymore', - output_json: { summary: 'AFFiNE structured' }, - usage: { prompt_tokens: 4, completion_tokens: 3, total_tokens: 7 }, - finish_reason: 'stop', - }); - - const result = await getProviderRuntimeHost(provider).run.structured( - { modelId: 'gpt-4.1' }, - jsonOnlyPromptMessages('Summarize AFFiNE in one sentence.'), - structuredOptions(z.object({ summary: z.string() })), - structuredContract(z.object({ summary: z.string() })) - ); - - t.snapshot(JSON.parse(result)); -}); - -test('OpenAIProvider should use native embedding dispatch', async t => { - const provider = new TestOpenAIProvider(); - t.teardown(installNativeDispatchRecorder(provider)); - - const result = await getProviderRuntimeHost(provider).run.embedding( - { modelId: 'text-embedding-3-small' }, - ['alpha', 'beta'], - { dimensions: 8 } - ); - - t.is(provider.embeddingRequests.length, 1); - t.snapshot({ - result, - request: provider.embeddingRequests[0], - }); -}); - -test('OpenAIProvider should use native rerank dispatch', async t => { - const provider = new TestOpenAIProvider(); - t.teardown(installNativeDispatchRecorder(provider)); - - const scores = await getProviderRuntimeHost(provider).run.rerank( - { modelId: 'gpt-4.1' }, - { - query: 'programming', - candidates: [ - { id: 'react', text: 'React is a UI library.' }, - { id: 'weather', text: 'The park is sunny today.' }, - ], - } - ); - - t.is(provider.rerankRequests.length, 1); - t.snapshot({ scores, request: provider.rerankRequests[0] }); -}); - -test('OpenAIProvider rerank should normalize native dispatch errors', async t => { - class ErroringOpenAIProvider extends TestOpenAIProvider { - override rerankFactory = () => { - throw new Error('native rerank exploded'); - }; - } - - const provider = new ErroringOpenAIProvider(); - t.teardown(installNativeDispatchRecorder(provider)); - - const error = await t.throwsAsync( - getProviderRuntimeHost(provider).run.rerank( - { modelId: 'gpt-4.1' }, - { - query: 'programming', - candidates: [{ id: 'react', text: 'React is a UI library.' }], - } - ) - ); - - t.true(error instanceof CopilotProviderSideError); - t.regex(error.message, /native rerank exploded/i); -}); diff --git a/packages/backend/server/src/__tests__/copilot/native-runtime-contract.spec.ts b/packages/backend/server/src/__tests__/copilot/native-runtime-contract.spec.ts new file mode 100644 index 0000000000..d221c54e92 --- /dev/null +++ b/packages/backend/server/src/__tests__/copilot/native-runtime-contract.spec.ts @@ -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))); + } +}); diff --git a/packages/backend/server/src/__tests__/copilot/provider-middleware.spec.ts b/packages/backend/server/src/__tests__/copilot/provider-middleware.spec.ts deleted file mode 100644 index 51eca69b82..0000000000 --- a/packages/backend/server/src/__tests__/copilot/provider-middleware.spec.ts +++ /dev/null @@ -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']); -}); diff --git a/packages/backend/server/src/__tests__/copilot/provider-native.spec.ts b/packages/backend/server/src/__tests__/copilot/provider-native.spec.ts deleted file mode 100644 index 9f1e7b11ea..0000000000 --- a/packages/backend/server/src/__tests__/copilot/provider-native.spec.ts +++ /dev/null @@ -1,3242 +0,0 @@ -import serverNativeModule from '@affine/server-native'; -import test from 'ava'; -import Sinon from 'sinon'; -import { z } from 'zod'; - -import { - CopilotPromptInvalid, - CopilotQuotaExceeded, - NoCopilotProviderAvailable, -} from '../../base'; -import { - type LlmBackendConfig, - type LlmEmbeddingRequest, - type LlmImageRequest, - llmMatchModelCapabilities, - type LlmPreparedDispatchRoute, - type LlmPreparedEmbeddingDispatchRoute, - type LlmPreparedImageDispatchRoute, - type LlmPreparedRerankDispatchRoute, - type LlmPreparedStructuredDispatchRoute, - type LlmProtocol, - type LlmRequest, - type LlmRerankRequest, - llmResolveRequestedModelMatch, - type LlmStructuredRequest, -} from '../../native'; -import type { - CopilotProviderProfile, - ProviderMiddlewareConfig, -} from '../../plugins/copilot/config'; -import { CopilotProviderFactory } from '../../plugins/copilot/providers/factory'; -import { OpenAIProvider } from '../../plugins/copilot/providers/openai'; -import { CopilotProvider } from '../../plugins/copilot/providers/provider'; -import { buildProviderRegistry } from '../../plugins/copilot/providers/provider-registry'; -import { - type CopilotProviderExecution, - type NativeExecutionRoute, - type ProviderDriverSpec, -} from '../../plugins/copilot/providers/provider-runtime-contract'; -import { - type CopilotProviderModel, - CopilotProviderType, - type ModelFullConditions, - ModelInputType, - ModelOutputType, -} from '../../plugins/copilot/providers/types'; -import { CapabilityRuntime } from '../../plugins/copilot/runtime/capability-runtime'; -import { - buildStructuredResponseContract, - parseCapabilityMatchRequest, - parseExecutionPlan, - parseProviderDriverSpec, - parseRequestedModelMatchRequest, - type RequiredStructuredOutputContract, - requireStructuredOutputContract, -} from '../../plugins/copilot/runtime/contracts'; -import { ExecutionPlanBuilder } from '../../plugins/copilot/runtime/execution-plan'; -import { NativeExecutionEngine } from '../../plugins/copilot/runtime/native-execution-engine'; -import { buildNativeRequest } from '../../plugins/copilot/runtime/native-request-runtime'; -import { getProviderRuntimeHost } from '../../plugins/copilot/runtime/provider-runtime-context'; -import { defineTool } from '../../plugins/copilot/tools/tool'; -import { - nativeMessages, - nativeUserText, - promptMessages, - singleUserPromptMessages, - systemPrompt, - userPrompt, -} from './prompt-test-helper'; - -function createNativeExecutionEngine() { - return new NativeExecutionEngine({ - recordUsage: Sinon.stub().resolves(), - recordProviderFailure: Sinon.stub().resolves(), - } as never); -} - -function structuredOptions( - schema: z.ZodTypeAny, - extra?: Record -) { - const { responseSchemaJson, schemaHash } = - buildStructuredResponseContract(schema); - return { - responseSchemaJson, - schemaHash, - ...extra, - }; -} - -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 TestOpenAIProvider extends CopilotProvider<{ apiKey: string }> { - readonly type = CopilotProviderType.OpenAI; - protected resolveModelBackendKind() { - return 'openai_responses' as const; - } - - configured() { - return true; - } - - async text(_cond: any, _messages: any[], _options?: any) { - return ''; - } - - async *streamText(_cond: any, _messages: any[], _options?: any) { - yield ''; - } - - exposeMetricLabels(execution?: CopilotProviderExecution) { - return this.metricLabels('gpt-5-mini', {}, execution); - } - - exposeMiddleware(execution?: CopilotProviderExecution) { - return this.getActiveProviderMiddleware(execution); - } -} - -class DriverOnlyProvider extends CopilotProvider<{ apiKey: string }> { - readonly type = CopilotProviderType.OpenAI; - protected resolveModelBackendKind() { - return 'openai_responses' as const; - } - - configured() { - return true; - } - - override getDriverSpec(): ProviderDriverSpec { - return { - createBackendConfig: async () => ({ - base_url: 'https://api.openai.com', - auth_token: 'test-key', - }), - mapError: (error: unknown) => error, - structured: {}, - embedding: {}, - rerank: {}, - image: {}, - }; - } -} - -async function collectAsync(iterable: AsyncIterable) { - const items: T[] = []; - for await (const item of iterable) { - items.push(item); - } - return items; -} - -const OPENAI_BASE_URL = 'https://api.openai.com'; -const GEMINI_BASE_URL = 'https://generativelanguage.googleapis.com'; - -function summarizePreparedDispatchRoutes(routes: unknown) { - if (!Array.isArray(routes)) { - return routes; - } - - return routes.map(route => { - const request = - route && typeof route === 'object' && 'request' in route - ? (route as Record).request - : undefined; - const firstContent = - request?.messages?.[0]?.content?.find?.( - (part: { type?: string }) => part?.type === 'text' - )?.text ?? null; - - const requestShape: Record = { - keys: request ? Object.keys(request).sort() : [], - firstContent, - schemaKeys: request?.schema?.properties - ? Object.keys(request.schema.properties).sort() - : undefined, - inputCount: Array.isArray(request?.inputs) ? request.inputs.length : 0, - query: request?.query, - candidateCount: Array.isArray(request?.candidates) - ? request.candidates.length - : 0, - toolNames: Array.isArray(request?.tools) - ? request.tools.map((tool: { name?: string }) => tool.name) - : [], - }; - - if (request && typeof request === 'object' && 'prompt' in request) { - requestShape.prompt = request.prompt; - } - if (request && typeof request === 'object' && 'images' in request) { - requestShape.imageCount = Array.isArray(request.images) - ? request.images.length - : 0; - } - - return { - providerId: - route && typeof route === 'object' && 'provider_id' in route - ? (route as Record).provider_id - : undefined, - model: - route && typeof route === 'object' && 'model' in route - ? (route as Record).model - : undefined, - requestShape, - }; - }); -} - -function nativeBackendConfig( - authToken: string, - baseUrl: string = OPENAI_BASE_URL -): LlmBackendConfig { - return { base_url: baseUrl, auth_token: authToken }; -} - -type NativeRouteOptions = { - providerId: string; - request: TRequest; - authToken: string; - protocol?: LlmProtocol; - model?: string; - baseUrl?: string; -}; - -function nativeRoute( - options: NativeRouteOptions -): LlmPreparedDispatchRoute; -function nativeRoute( - options: NativeRouteOptions -): LlmPreparedStructuredDispatchRoute; -function nativeRoute( - options: NativeRouteOptions -): LlmPreparedEmbeddingDispatchRoute; -function nativeRoute( - options: NativeRouteOptions -): LlmPreparedRerankDispatchRoute; -function nativeRoute( - options: NativeRouteOptions -): LlmPreparedImageDispatchRoute; -function nativeRoute({ - providerId, - request, - authToken, - protocol = 'openai_chat', - model = 'gpt-5-mini', - baseUrl = OPENAI_BASE_URL, -}: NativeRouteOptions< - | LlmRequest - | LlmStructuredRequest - | LlmEmbeddingRequest - | LlmRerankRequest - | LlmImageRequest ->) { - return { - provider_id: providerId, - protocol, - model, - config: nativeBackendConfig(authToken, baseUrl), - request, - }; -} - -function preparedRoute({ - providerId, - authToken, - protocol = 'openai_chat', - model = 'gpt-5-mini', - baseUrl = OPENAI_BASE_URL, -}: { - providerId: string; - authToken: string; - protocol?: LlmProtocol; - model?: string; - baseUrl?: string; -}): NativeExecutionRoute & { providerId: string } { - return { - providerId, - protocol, - model, - backendConfig: nativeBackendConfig(authToken, baseUrl), - }; -} - -function nativeTextRequest( - text: string, - model: string = 'gpt-5-mini' -): LlmRequest { - return { model, messages: nativeMessages(nativeUserText(text)) }; -} - -function nativeStructuredRequest( - text: string, - schema: Record, - model: string = 'gpt-5-mini' -): LlmStructuredRequest { - return { ...nativeTextRequest(text, model), schema }; -} - -function nativeEmbeddingRequest( - input: string, - model: string = 'text-embedding-3-small' -): LlmEmbeddingRequest { - return { model, inputs: [input] }; -} - -function nativeRerankRequest( - query: string, - candidates: Array<{ id?: string; text: string }>, - model: string = 'gpt-4o-mini' -): LlmRerankRequest { - return { model, query, candidates }; -} - -function nativeImageRequest( - prompt: string, - model: string = 'gpt-image-1' -): LlmImageRequest { - return { model, prompt, operation: 'generate', images: [] }; -} - -function createProvider(profileMiddleware?: ProviderMiddlewareConfig) { - const provider = new TestOpenAIProvider(); - (provider as any).AFFiNEConfig = { - copilot: { - providers: { - profiles: [ - { - id: 'openai-main', - type: CopilotProviderType.OpenAI, - config: { apiKey: 'test' }, - middleware: profileMiddleware, - }, - ], - defaults: {}, - openai: { apiKey: 'legacy' }, - }, - }, - }; - return provider; -} - -function createExecution( - provider: TestOpenAIProvider -): CopilotProviderExecution { - const registry = buildProviderRegistry( - (provider as any).AFFiNEConfig.copilot.providers - ); - const profile = registry.profiles.get('openai-main'); - if (!profile) { - throw new Error('missing openai-main profile'); - } - return { - providerId: 'openai-main', - profile, - }; -} - -test('metricLabels should include active provider id', t => { - const provider = createProvider(); - const labels = provider.exposeMetricLabels(createExecution(provider)); - t.is(labels.providerId, 'openai-main'); -}); - -test('CapabilityRuntime should route capability plans through plan builder and native engine', async t => { - const plans = { - buildTextPlan: Sinon.stub().resolves({ kind: 'text-plan' }), - buildStreamTextPlan: Sinon.stub().resolves({ kind: 'stream-text-plan' }), - buildStreamObjectPlan: Sinon.stub().resolves({ - kind: 'stream-object-plan', - }), - buildStructuredPlan: Sinon.stub().resolves({ - kind: 'structured-plan', - routePolicy: { fallbackOrder: ['openai-primary'] }, - }), - buildEmbeddingPlan: Sinon.stub().resolves({ - kind: 'embedding-plan', - routePolicy: { fallbackOrder: ['openai-primary'] }, - }), - buildRerankPlan: Sinon.stub().resolves({ - kind: 'rerank-plan', - routePolicy: { fallbackOrder: ['openai-primary'] }, - }), - }; - const engine = { - execute: Sinon.stub().callsFake( - async (plan: { - kind: string; - routePolicy?: { fallbackOrder: string[] }; - }) => { - switch (plan.kind) { - case 'text-plan': - return 'done'; - case 'structured-plan': - return '{"ok":true}'; - case 'embedding-plan': - return [[0.1, 0.2]]; - case 'rerank-plan': - return [0.9, 0.1]; - default: - throw new Error(`unexpected execute plan: ${plan.kind}`); - } - } - ), - executeStream: Sinon.stub().callsFake((plan: { kind: string }) => { - switch (plan.kind) { - case 'stream-text-plan': - return (async function* () { - yield 'chunk'; - })(); - case 'stream-object-plan': - return (async function* () { - yield { type: 'text-delta', textDelta: 'chunk' } as const; - })(); - default: - throw new Error(`unexpected executeStream plan: ${plan.kind}`); - } - }), - }; - const runtime = new CapabilityRuntime(plans as never, engine as never); - const schema = z.object({ ok: z.boolean() }); - const cases = [ - { - title: 'text', - planBuilder: plans.buildTextPlan, - execute: () => - runtime.text( - { modelId: 'gpt-5-mini' }, - promptMessages(userPrompt('hi')) - ), - expected: 'done', - executionStub: engine.execute, - expectedPlan: { kind: 'text-plan' }, - }, - { - title: 'streamText', - planBuilder: plans.buildStreamTextPlan, - execute: () => - collectAsync( - runtime.streamText( - { modelId: 'gpt-5-mini' }, - promptMessages(userPrompt('hi')) - ) - ), - expected: ['chunk'], - executionStub: engine.executeStream, - expectedPlan: { kind: 'stream-text-plan' }, - }, - { - title: 'streamObject', - planBuilder: plans.buildStreamObjectPlan, - execute: () => - collectAsync( - runtime.streamObject( - { modelId: 'gpt-5-mini' }, - promptMessages(userPrompt('hi')) - ) - ), - expected: [{ type: 'text-delta', textDelta: 'chunk' }], - executionStub: engine.executeStream, - expectedPlan: { kind: 'stream-object-plan' }, - }, - { - title: 'structured', - planBuilder: plans.buildStructuredPlan, - execute: () => - runtime.generateStructured( - { modelId: 'gpt-5-mini' }, - promptMessages(userPrompt('hi')), - structuredOptions(schema), - undefined, - structuredContract(schema) - ), - expected: '{"ok":true}', - executionStub: engine.execute, - expectedPlan: { - kind: 'structured-plan', - routePolicy: { fallbackOrder: ['openai-primary'] }, - }, - }, - { - title: 'embedding', - planBuilder: plans.buildEmbeddingPlan, - execute: () => runtime.embed('text-embedding-3-small', 'hello world'), - expected: [[0.1, 0.2]], - executionStub: engine.execute, - expectedPlan: { - kind: 'embedding-plan', - routePolicy: { fallbackOrder: ['openai-primary'] }, - }, - }, - { - title: 'rerank', - planBuilder: plans.buildRerankPlan, - execute: () => - runtime.rerank('gpt-4o-mini', { - query: 'programming', - candidates: [{ text: 'React is a UI library.' }], - }), - expected: [0.9, 0.1], - executionStub: engine.execute, - expectedPlan: { - kind: 'rerank-plan', - routePolicy: { fallbackOrder: ['openai-primary'] }, - }, - }, - ] as const; - - for (const testCase of cases) { - t.deepEqual(await testCase.execute(), testCase.expected, testCase.title); - Sinon.assert.calledOnce(testCase.planBuilder); - Sinon.assert.calledWith(testCase.executionStub, testCase.expectedPlan); - } -}); - -test('CapabilityRuntime should defer no-route embedding plans to native engine', async t => { - const plans = { - buildEmbeddingPlan: Sinon.stub().resolves({ - kind: 'embedding-plan', - routePolicy: { fallbackOrder: [] }, - routes: [{}], - }), - }; - const engine = { - execute: Sinon.stub().rejects( - new NoCopilotProviderAvailable({ - modelId: 'text-embedding-3-small', - }) - ), - }; - const runtime = new CapabilityRuntime(plans as never, engine as never); - - const error = await t.throwsAsync(() => - runtime.embed('text-embedding-3-small', 'hello world') - ); - - t.true(error instanceof NoCopilotProviderAvailable); - Sinon.assert.calledOnce(engine.execute); - Sinon.assert.calledWith(engine.execute, { - kind: 'embedding-plan', - routePolicy: { fallbackOrder: [] }, - routes: [{}], - }); -}); - -test('NativeExecutionEngine should expose execute/executeStream as the single plan entrypoints', async t => { - const engine = createNativeExecutionEngine(); - let dispatchCalls = 0; - let streamCalls = 0; - - const originalDispatch = (serverNativeModule as any).llmDispatchPrepared; - const originalStream = (serverNativeModule as any).llmDispatchPreparedStream; - (serverNativeModule as any).llmDispatchPrepared = () => { - dispatchCalls += 1; - return JSON.stringify({ - provider_id: 'openai-primary', - response: { - id: 'chat_execute', - model: 'gpt-5-mini', - message: { - role: 'assistant', - content: [{ type: 'text', text: 'execute-ok' }], - }, - usage: { - prompt_tokens: 1, - completion_tokens: 1, - total_tokens: 2, - }, - finish_reason: 'stop', - }, - }); - }; - (serverNativeModule as any).llmDispatchPreparedStream = ( - _routesJson: string, - callback: (error: Error | null, arg: string) => void - ) => { - streamCalls += 1; - callback(null, JSON.stringify({ type: 'text_delta', text: 'stream-ok' })); - callback(null, '__AFFINE_LLM_STREAM_END__'); - return { abort() {} }; - }; - t.teardown(() => { - (serverNativeModule as any).llmDispatchPrepared = originalDispatch; - (serverNativeModule as any).llmDispatchPreparedStream = originalStream; - }); - - const text = await engine.execute({ - nativeDispatch: { - chat: { - routes: [ - nativeRoute({ - providerId: 'openai-primary', - authToken: 'primary-key', - request: nativeTextRequest('hello'), - }), - ], - prepared: { - route: preparedRoute({ - providerId: 'openai-primary', - authToken: 'primary-key', - }), - request: nativeTextRequest('hello'), - tools: {}, - postprocess: { nodeTextMiddleware: [] }, - }, - hasTools: false, - }, - }, - request: { - kind: 'text', - cond: { modelId: 'gpt-5-mini' }, - messages: singleUserPromptMessages('hello'), - options: undefined, - }, - routePolicy: { fallbackOrder: ['openai-primary'] }, - runtimePolicy: {}, - attachmentPolicy: { materializeRemoteAttachments: true }, - responsePostprocess: { mode: 'text' }, - hostPersistence: { persistAssistantTurn: true, outputKind: 'text' }, - hostContext: {}, - }); - const chunks = await collectAsync( - engine.executeStream({ - nativeDispatch: { - chat: { - routes: [ - nativeRoute({ - providerId: 'openai-primary', - authToken: 'primary-key', - request: nativeTextRequest('hello'), - }), - ], - prepared: { - route: preparedRoute({ - providerId: 'openai-primary', - authToken: 'primary-key', - }), - request: nativeTextRequest('hello'), - tools: {}, - postprocess: { nodeTextMiddleware: [] }, - }, - hasTools: false, - }, - }, - request: { - kind: 'streamText', - cond: { modelId: 'gpt-5-mini' }, - messages: singleUserPromptMessages('hello'), - options: undefined, - }, - routePolicy: { fallbackOrder: ['openai-primary'] }, - runtimePolicy: {}, - attachmentPolicy: { materializeRemoteAttachments: true }, - responsePostprocess: { mode: 'streamText' }, - hostPersistence: { persistAssistantTurn: true, outputKind: 'streamText' }, - hostContext: {}, - }) - ); - - t.is(text, 'execute-ok'); - t.deepEqual(chunks, ['stream-ok']); - t.is(dispatchCalls, 1); - t.is(streamCalls, 1); -}); - -test('NativeExecutionEngine should record BYOK usage when stream finalizes with selected provider', async t => { - const byok = { - recordUsage: Sinon.stub().resolves(), - }; - const engine = new NativeExecutionEngine(byok as never); - const providerId = 'byok-aaaaaaaaaaaa-openai-server-key1'; - - const originalStream = (serverNativeModule as any).llmDispatchPreparedStream; - (serverNativeModule as any).llmDispatchPreparedStream = ( - _routesJson: string, - callback: (error: Error | null, arg: string) => void - ) => { - callback( - null, - JSON.stringify({ - type: 'message_start', - model: 'gpt-5-mini', - }) - ); - callback(null, JSON.stringify({ type: 'text_delta', text: 'ok' })); - callback( - null, - JSON.stringify({ - type: 'done', - finish_reason: 'stop', - usage: { - prompt_tokens: 2, - completion_tokens: 3, - total_tokens: 5, - }, - }) - ); - callback( - null, - JSON.stringify({ - type: 'provider_selected', - provider_id: providerId, - }) - ); - callback(null, '__AFFINE_LLM_STREAM_END__'); - return { abort() {} }; - }; - t.teardown(() => { - (serverNativeModule as any).llmDispatchPreparedStream = originalStream; - }); - - const chunks = await collectAsync( - engine.executeStream({ - nativeDispatch: { - chat: { - routes: [ - nativeRoute({ - providerId, - authToken: 'byok-key', - request: nativeTextRequest('hello'), - }), - ], - prepared: { - route: preparedRoute({ - providerId, - authToken: 'byok-key', - }), - request: nativeTextRequest('hello'), - tools: {}, - postprocess: { nodeTextMiddleware: [] }, - }, - hasTools: false, - }, - }, - request: { - kind: 'streamText', - cond: { modelId: 'gpt-5-mini' }, - messages: singleUserPromptMessages('hello'), - options: { - workspace: 'workspace-1', - user: 'user-1', - session: 'session-1', - featureKind: 'chat', - }, - }, - routePolicy: { fallbackOrder: [providerId] }, - runtimePolicy: {}, - attachmentPolicy: { materializeRemoteAttachments: true }, - responsePostprocess: { mode: 'streamText' }, - hostPersistence: { persistAssistantTurn: true, outputKind: 'streamText' }, - hostContext: {}, - }) - ); - - t.deepEqual(chunks, ['ok']); - Sinon.assert.calledOnceWithMatch(byok.recordUsage, { - workspaceId: 'workspace-1', - userId: 'user-1', - sessionId: 'session-1', - featureKind: 'chat', - providerId, - model: 'gpt-5-mini', - usage: { - prompt_tokens: 2, - completion_tokens: 3, - total_tokens: 5, - }, - }); -}); - -test('NativeExecutionEngine should record plain text BYOK usage as chat by default', async t => { - const byok = { - recordUsage: Sinon.stub().resolves(), - recordProviderFailure: Sinon.stub().resolves(), - }; - const engine = new NativeExecutionEngine(byok as never); - const providerId = 'byok-aaaaaaaaaaaa-openai-server-key1'; - - const originalDispatch = (serverNativeModule as any).llmDispatchPrepared; - (serverNativeModule as any).llmDispatchPrepared = () => { - return JSON.stringify({ - provider_id: providerId, - response: { - id: 'chat_execute', - model: 'gpt-5-mini', - message: { - role: 'assistant', - content: [{ type: 'text', text: 'execute-ok' }], - }, - usage: { - prompt_tokens: 1, - completion_tokens: 2, - total_tokens: 3, - }, - finish_reason: 'stop', - }, - }); - }; - t.teardown(() => { - (serverNativeModule as any).llmDispatchPrepared = originalDispatch; - }); - - const text = await engine.execute({ - nativeDispatch: { - chat: { - routes: [ - nativeRoute({ - providerId, - authToken: 'byok-key', - request: nativeTextRequest('hello'), - }), - ], - prepared: { - route: preparedRoute({ - providerId, - authToken: 'byok-key', - }), - request: nativeTextRequest('hello'), - tools: {}, - postprocess: { nodeTextMiddleware: [] }, - }, - hasTools: false, - }, - }, - request: { - kind: 'text', - cond: { modelId: 'gpt-5-mini' }, - messages: singleUserPromptMessages('hello'), - options: { - workspace: 'workspace-1', - user: 'user-1', - session: 'session-1', - }, - }, - routePolicy: { fallbackOrder: [providerId] }, - runtimePolicy: {}, - attachmentPolicy: { materializeRemoteAttachments: true }, - responsePostprocess: { mode: 'text' }, - hostPersistence: { persistAssistantTurn: true, outputKind: 'text' }, - hostContext: {}, - }); - - t.is(text, 'execute-ok'); - Sinon.assert.calledOnceWithMatch(byok.recordUsage, { - workspaceId: 'workspace-1', - userId: 'user-1', - sessionId: 'session-1', - featureKind: 'chat', - providerId, - model: 'gpt-5-mini', - usage: { - prompt_tokens: 1, - completion_tokens: 2, - total_tokens: 3, - }, - }); -}); - -test('NativeExecutionEngine should not fail stream when BYOK usage recording fails', async t => { - const byok = { - recordUsage: Sinon.stub().rejects(new Error('usage db down')), - }; - const engine = new NativeExecutionEngine(byok as never); - const providerId = 'byok-aaaaaaaaaaaa-openai-server-key1'; - - const originalStream = (serverNativeModule as any).llmDispatchPreparedStream; - (serverNativeModule as any).llmDispatchPreparedStream = ( - _routesJson: string, - callback: (error: Error | null, arg: string) => void - ) => { - callback( - null, - JSON.stringify({ type: 'message_start', model: 'gpt-5-mini' }) - ); - callback(null, JSON.stringify({ type: 'text_delta', text: 'ok' })); - callback( - null, - JSON.stringify({ - type: 'done', - finish_reason: 'stop', - usage: { prompt_tokens: 2, completion_tokens: 3, total_tokens: 5 }, - }) - ); - callback( - null, - JSON.stringify({ type: 'provider_selected', provider_id: providerId }) - ); - callback(null, '__AFFINE_LLM_STREAM_END__'); - return { abort() {} }; - }; - t.teardown(() => { - (serverNativeModule as any).llmDispatchPreparedStream = originalStream; - }); - - const chunks = await collectAsync( - engine.executeStream({ - nativeDispatch: { - chat: { - routes: [ - nativeRoute({ - providerId, - authToken: 'byok-key', - request: nativeTextRequest('hello'), - }), - ], - prepared: { - route: preparedRoute({ providerId, authToken: 'byok-key' }), - request: nativeTextRequest('hello'), - tools: {}, - postprocess: { nodeTextMiddleware: [] }, - }, - hasTools: false, - }, - }, - request: { - kind: 'streamText', - cond: { modelId: 'gpt-5-mini' }, - messages: singleUserPromptMessages('hello'), - options: { - workspace: 'workspace-1', - user: 'user-1', - session: 'session-1', - featureKind: 'chat', - }, - }, - routePolicy: { fallbackOrder: [providerId] }, - runtimePolicy: {}, - attachmentPolicy: { materializeRemoteAttachments: true }, - responsePostprocess: { mode: 'streamText' }, - hostPersistence: { persistAssistantTurn: true, outputKind: 'streamText' }, - hostContext: {}, - }) - ); - - t.deepEqual(chunks, ['ok']); - Sinon.assert.calledOnce(byok.recordUsage); -}); - -test('CopilotProviderFactory should return no prepared routes when native prepare returns null', async t => { - const provider = new DriverOnlyProvider(); - (provider as any).AFFiNEConfig = { copilot: { providers: { openai: {} } } }; - (provider as any).toolExecutorHost = { - createNativeAdapter: () => { - throw new Error('native adapter should not be used'); - }, - getTools: async () => ({}), - }; - const runtimeHost = getProviderRuntimeHost(provider); - runtimeHost.prepare.chat = async () => null; - runtimeHost.prepare.structured = async () => null; - runtimeHost.prepare.embedding = async () => null; - runtimeHost.prepare.rerank = async () => null; - - const registryService = { - getRegistry: () => - buildProviderRegistry({ - profiles: [ - { - id: 'openai-main', - type: CopilotProviderType.OpenAI, - config: { apiKey: 'test-key' }, - }, - ], - defaults: {}, - openai: { apiKey: 'test-key' }, - }), - }; - const server = { - enableFeature: Sinon.stub(), - disableFeature: Sinon.stub(), - }; - const access = { - resolveRouteAccess: Sinon.stub().resolves({ - byokProfiles: [], - quotaBackedRoutesAvailable: true, - }), - }; - const factory = new CopilotProviderFactory( - server as never, - registryService as never, - access as never - ); - factory.register('openai-main', provider); - - const chatRoutes = await factory.prepareRoutes( - 'text', - { - modelId: 'gpt-5-mini', - outputType: ModelOutputType.Text, - }, - singleUserPromptMessages('hello') - ); - const structuredRoutes = await factory.prepareStructuredRoutes( - { - modelId: 'gpt-5-mini', - outputType: ModelOutputType.Structured, - }, - singleUserPromptMessages('hello'), - structuredOptions(z.object({ ok: z.boolean() })), - {}, - structuredContract(z.object({ ok: z.boolean() })) - ); - const embeddingRoutes = await factory.prepareEmbeddingRoutes( - 'text-embedding-3-small', - 'hello world' - ); - const rerankRoutes = await factory.prepareRerankRoutes('gpt-4o-mini', { - query: 'programming', - candidates: [{ text: 'React is a UI library.' }], - }); - - t.snapshot({ - chat: { - length: chatRoutes.length, - providerId: chatRoutes[0]?.providerId, - prepared: chatRoutes[0]?.prepared, - }, - structured: { - length: structuredRoutes.length, - prepared: structuredRoutes[0]?.preparedStructured, - }, - embedding: { - length: embeddingRoutes.length, - prepared: embeddingRoutes[0]?.preparedEmbedding, - }, - rerank: { - length: rerankRoutes.length, - prepared: rerankRoutes[0]?.preparedRerank, - }, - }); -}); - -test('driver-only provider should use base native driver templates', async t => { - const provider = new DriverOnlyProvider(); - (provider as any).AFFiNEConfig = { copilot: { providers: { openai: {} } } }; - (provider as any).toolExecutorHost = { - createNativeAdapter: () => ({ - text: async () => 'driver text', - streamText: async function* () { - yield 'driver stream'; - }, - streamObject: async function* () { - yield { type: 'text-delta', textDelta: 'driver object' }; - }, - }), - getTools: async () => ({}), - }; - 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 - ) => - JSON.stringify({ - id: 'structured_1', - model: 'gpt-5-mini', - output_text: '{"ok":true}', - output_json: { ok: true }, - usage: { prompt_tokens: 1, completion_tokens: 1, total_tokens: 2 }, - finish_reason: 'stop', - }); - (serverNativeModule as any).llmEmbeddingDispatch = ( - _protocol: string, - _backendConfigJson: string, - _requestJson: string - ) => JSON.stringify({ embeddings: [[0.1, 0.2]] }); - (serverNativeModule as any).llmRerankDispatch = ( - _protocol: string, - _backendConfigJson: string, - _requestJson: string - ) => JSON.stringify({ scores: [0.9, 0.1] }); - t.teardown(() => { - (serverNativeModule as any).llmStructuredDispatch = originalStructured; - (serverNativeModule as any).llmEmbeddingDispatch = originalEmbedding; - (serverNativeModule as any).llmRerankDispatch = originalRerank; - }); - - const runtimeHost = getProviderRuntimeHost(provider); - const schema = z.object({ ok: z.boolean() }); - const helloPrompt = promptMessages(userPrompt('hello')); - const cases = [ - { - title: 'text', - run: () => runtimeHost.run.text({ modelId: 'gpt-5-mini' }, helloPrompt), - expected: 'driver text', - }, - { - title: 'streamText', - run: () => - collectAsync( - runtimeHost.run.streamText({ modelId: 'gpt-5-mini' }, helloPrompt) - ), - expected: ['driver stream'], - }, - { - title: 'streamObject', - run: () => - collectAsync( - runtimeHost.run.streamObject({ modelId: 'gpt-5-mini' }, helloPrompt) - ), - expected: [{ type: 'text-delta', textDelta: 'driver object' }], - }, - { - title: 'structured', - run: () => - runtimeHost.run.structured( - { modelId: 'gpt-5-mini' }, - helloPrompt, - structuredOptions(schema), - structuredContract(schema) - ), - expected: '{"ok":true}', - }, - { - title: 'embedding', - run: () => - runtimeHost.run.embedding( - { modelId: 'text-embedding-3-small' }, - 'hello world' - ), - expected: [[0.1, 0.2]], - }, - { - title: 'rerank', - run: () => - runtimeHost.run.rerank( - { modelId: 'gpt-4o-mini' }, - { - query: 'programming', - candidates: [{ text: 'React is a UI library.' }], - } - ), - expected: [0.9, 0.1], - }, - ] as const; - - for (const testCase of cases) { - t.deepEqual(await testCase.run(), testCase.expected, testCase.title); - } -}); - -test('driver-only provider should require explicit structured response contracts', async t => { - const provider = new DriverOnlyProvider(); - (provider as any).AFFiNEConfig = { copilot: { providers: { openai: {} } } }; - (provider as any).toolExecutorHost = { - createNativeAdapter: () => { - throw new Error( - 'chat adapter should not be used in non-chat driver test' - ); - }, - getTools: async () => ({}), - }; - - const schemaJson = { - type: 'object', - properties: { - ok: { type: 'boolean' }, - }, - additionalProperties: false, - required: ['ok'], - }; - let capturedRequest: - | { - schema?: unknown; - strict?: boolean; - messages?: Array<{ - response_format?: { - response_schema_json?: unknown; - strict?: boolean; - }; - }>; - } - | undefined; - - const original = (serverNativeModule as any) - .llmBuildCanonicalStructuredRequest; - (serverNativeModule as any).llmBuildCanonicalStructuredRequest = ( - requestJson: string - ) => { - capturedRequest = JSON.parse(requestJson); - return original(requestJson); - }; - t.teardown(() => { - (serverNativeModule as any).llmBuildCanonicalStructuredRequest = original; - }); - - const error = await t.throwsAsync(() => - getProviderRuntimeHost(provider).prepare.structured( - { modelId: 'gpt-5-mini' }, - [ - systemPrompt('Return JSON only.', { - responseFormat: { - type: 'json_schema', - responseSchemaJson: schemaJson, - strict: false, - }, - }), - userPrompt('hello'), - ] - ) - ); - - t.true(error instanceof CopilotPromptInvalid); - t.is(capturedRequest, undefined); -}); - -test('getActiveProviderMiddleware should merge defaults with profile override', t => { - const provider = createProvider({ - rust: { request: ['clamp_max_tokens'] }, - node: { text: ['thinking_format'] }, - }); - - const middleware = provider.exposeMiddleware(createExecution(provider)); - - t.snapshot(middleware); -}); - -test('llmMatchModelCapabilities should honor structured attachment capability and remote rules', t => { - const contract = parseCapabilityMatchRequest({ - models: [ - { - id: 'structured-file', - capabilities: [ - { - input: ['text', 'file'], - output: ['structured'], - attachments: { - kinds: ['image'], - sourceKinds: ['url'], - allowRemoteUrls: true, - }, - structuredAttachments: { - kinds: ['file'], - sourceKinds: ['file_handle'], - allowRemoteUrls: false, - }, - defaultForOutputType: true, - }, - ], - }, - ], - cond: { - modelId: 'structured-file', - outputType: 'structured', - inputTypes: ['text', 'file'], - attachmentKinds: ['file'], - attachmentSourceKinds: ['file_handle'], - hasRemoteAttachments: false, - }, - }); - - const modelId = llmMatchModelCapabilities( - contract.models.map(model => ({ - ...model, - capabilities: model.capabilities.map(capability => ({ - ...capability, - input: capability.input.map(input => input as ModelInputType), - output: capability.output.map(output => output as ModelOutputType), - })), - })), - { - modelId: contract.cond.modelId, - outputType: contract.cond.outputType as ModelOutputType, - inputTypes: contract.cond.inputTypes as ModelInputType[], - attachmentKinds: contract.cond.attachmentKinds, - attachmentSourceKinds: contract.cond.attachmentSourceKinds, - hasRemoteAttachments: contract.cond.hasRemoteAttachments, - } - ); - - t.is(modelId, 'structured-file'); - t.is( - llmMatchModelCapabilities( - [ - { - id: 'structured-file', - capabilities: [ - { - input: [ModelInputType.Text, ModelInputType.File], - output: [ModelOutputType.Structured], - structuredAttachments: { - kinds: ['file'], - sourceKinds: ['file_handle'], - allowRemoteUrls: false, - }, - defaultForOutputType: true, - }, - ], - }, - ], - { - modelId: 'structured-file', - outputType: ModelOutputType.Structured, - inputTypes: [ModelInputType.Text, ModelInputType.File], - attachmentKinds: ['file'], - attachmentSourceKinds: ['url'], - hasRemoteAttachments: true, - } - ), - undefined - ); -}); - -test('llmMatchModelCapabilities should cover capability matrix combinations', t => { - const models: CopilotProviderModel[] = [ - { - id: 'text-default', - capabilities: [ - { - input: [ModelInputType.Text], - output: [ModelOutputType.Text], - defaultForOutputType: true, - }, - ], - }, - { - id: 'vision-remote', - capabilities: [ - { - input: [ModelInputType.Text, ModelInputType.Image], - output: [ModelOutputType.Text], - attachments: { - kinds: ['image'], - sourceKinds: ['url'], - allowRemoteUrls: true, - }, - }, - ], - }, - { - id: 'structured-file', - capabilities: [ - { - input: [ModelInputType.Text, ModelInputType.File], - output: [ModelOutputType.Structured], - structuredAttachments: { - kinds: ['file'], - sourceKinds: ['file_handle'], - allowRemoteUrls: false, - }, - defaultForOutputType: true, - }, - ], - }, - ]; - - const cases: Array<{ - title: string; - cond: ModelFullConditions; - expected?: string; - }> = [ - { - title: 'default text model', - cond: { - outputType: ModelOutputType.Text, - inputTypes: [ModelInputType.Text], - }, - expected: 'text-default', - }, - { - title: 'explicit multimodal override', - cond: { - modelId: 'vision-remote', - outputType: ModelOutputType.Text, - inputTypes: [ModelInputType.Text, ModelInputType.Image], - attachmentKinds: ['image'], - attachmentSourceKinds: ['url'], - hasRemoteAttachments: true, - }, - expected: 'vision-remote', - }, - { - title: 'structured file capability', - cond: { - outputType: ModelOutputType.Structured, - inputTypes: [ModelInputType.Text, ModelInputType.File], - attachmentKinds: ['file'], - attachmentSourceKinds: ['file_handle'], - }, - expected: 'structured-file', - }, - { - title: 'remote attachment rejected when capability is stricter', - cond: { - modelId: 'structured-file', - outputType: ModelOutputType.Structured, - inputTypes: [ModelInputType.Text, ModelInputType.File], - attachmentKinds: ['file'], - attachmentSourceKinds: ['url'], - hasRemoteAttachments: true, - }, - expected: undefined, - }, - ]; - - for (const entry of cases) { - t.is( - llmMatchModelCapabilities(models, entry.cond), - entry.expected, - entry.title - ); - } -}); - -test('checkParams should infer remote image capability from url extension without host mime inference', async t => { - const provider = new TestOpenAIProvider(); - - const cond = await provider.checkParams({ - cond: { - modelId: 'gpt-4.1', - outputType: ModelOutputType.Text, - inputTypes: [ModelInputType.Text], - }, - messages: [ - { - role: 'user', - content: 'describe this image', - attachments: ['https://example.com/cat.png'], - }, - ], - }); - - t.snapshot({ - inputTypes: cond.inputTypes, - attachmentKinds: cond.attachmentKinds, - attachmentSourceKinds: cond.attachmentSourceKinds, - }); - t.is(cond.hasRemoteAttachments, true); -}); - -test('llmResolveRequestedModelMatch should preserve provider-prefixed optional matches', t => { - const request = parseRequestedModelMatchRequest({ - providerIds: ['openai-default', 'gemini-default'], - defaultModel: 'gpt-5.6-luna', - optionalModels: ['gpt-5.6-luna', 'gpt-5.6-terra'], - requestedModelId: 'openai-default/gpt-5.6-terra', - }); - - t.snapshot(llmResolveRequestedModelMatch(request), 'prefixed optional hit'); - t.snapshot( - llmResolveRequestedModelMatch({ - ...request, - requestedModelId: 'openai-default/not-in-optional', - }), - 'prefixed optional miss' - ); -}); - -test('CopilotProviderFactory should resolve legacy model ids through native registry without migration', async t => { - const provider = createProvider(); - const registryService = { - getRegistry: () => - buildProviderRegistry({ - profiles: [ - { - id: 'openai-main', - type: CopilotProviderType.OpenAI, - config: { apiKey: 'test-key' }, - }, - ], - defaults: {}, - openai: { apiKey: 'test-key' }, - }), - }; - const server = { - enableFeature: Sinon.stub(), - disableFeature: Sinon.stub(), - }; - const access = { - resolveRouteAccess: Sinon.stub().resolves({ - byokProfiles: [], - quotaBackedRoutesAvailable: true, - }), - }; - const factory = new CopilotProviderFactory( - server as never, - registryService as never, - access as never - ); - factory.register('openai-main', provider); - - const resolvedProvider = await factory.getProviderByModel('gpt-5-2025-08-07'); - t.is(resolvedProvider, provider); - t.is(provider.resolveModel('gpt-5-2025-08-07')?.id, 'gpt-5'); -}); - -const BYOK_OPENAI_PROFILE: CopilotProviderProfile = { - id: 'byok-aaaaaaaaaaaa-openai-server-key1', - type: CopilotProviderType.OpenAI, - priority: 10_000, - config: { apiKey: 'byok-key' }, -}; - -const BYOK_FAL_PROFILE: CopilotProviderProfile = { - id: 'byok-aaaaaaaaaaaa-fal-server-key1', - type: CopilotProviderType.FAL, - priority: 10_000, - config: { apiKey: 'byok-key' }, -}; - -function createProviderFactoryWithByokRoutes({ - byokProfiles = [BYOK_OPENAI_PROFILE], - hasQuota = true, -}: { - byokProfiles?: CopilotProviderProfile[]; - hasQuota?: boolean; -} = {}) { - const provider = createProvider(); - const registryService = { - getRegistry: () => - buildProviderRegistry({ - profiles: [ - { - id: 'openai-main', - type: CopilotProviderType.OpenAI, - priority: 1, - config: { apiKey: 'test-key' }, - }, - ], - defaults: {}, - }), - }; - const server = { - enableFeature: Sinon.stub(), - disableFeature: Sinon.stub(), - }; - const byok = { - getProfiles: Sinon.stub().resolves(byokProfiles), - }; - const access = { - resolveRouteAccess: Sinon.stub().callsFake(async context => ({ - byokProfiles: await byok.getProfiles(context), - quotaBackedRoutesAvailable: context.quotaBackedRoutesAllowed ?? hasQuota, - })), - }; - const factory = new CopilotProviderFactory( - server as never, - registryService as never, - access as never - ); - factory.register('openai-main', provider); - - return { factory, byok }; -} - -test('CopilotProviderFactory should use matching BYOK routes before quota-backed routes', async t => { - const { factory } = createProviderFactoryWithByokRoutes(); - - const routes = await factory.resolveRoutes( - { modelId: 'gpt-5-mini', outputType: ModelOutputType.Text }, - {}, - { userId: 'user-1', workspaceId: 'workspace-1' } - ); - - t.deepEqual( - routes.map(route => route.providerId), - ['byok-aaaaaaaaaaaa-openai-server-key1'] - ); -}); - -test('CopilotProviderFactory should skip unsupported BYOK profiles and use quota-backed fallback', async t => { - const { factory } = createProviderFactoryWithByokRoutes({ - byokProfiles: [BYOK_FAL_PROFILE], - }); - - const routes = await factory.resolveRoutes( - { modelId: 'gpt-5-mini', outputType: ModelOutputType.Text }, - {}, - { userId: 'user-1', workspaceId: 'workspace-1' } - ); - - t.deepEqual( - routes.map(route => route.providerId), - ['openai-main'] - ); -}); - -test('CopilotProviderFactory should resolve BYOK embedding routes with workspace context', async t => { - const { factory, byok } = createProviderFactoryWithByokRoutes(); - - const routes = await factory.resolveRoutes( - { - modelId: 'text-embedding-3-small', - outputType: ModelOutputType.Embedding, - }, - {}, - { workspaceId: 'workspace-1', featureKind: 'workspace_indexing' } - ); - - t.deepEqual( - routes.map(route => route.providerId), - ['byok-aaaaaaaaaaaa-openai-server-key1'] - ); - Sinon.assert.calledOnceWithMatch(byok.getProfiles, { - workspaceId: 'workspace-1', - featureKind: 'workspace_indexing', - }); -}); - -test('CopilotProviderFactory should treat embedding preparation as embedding feature by default', async t => { - const { factory, byok } = createProviderFactoryWithByokRoutes(); - - await factory.prepareEmbeddingRoutes('text-embedding-3-small', 'hello', { - workspace: 'workspace-1', - }); - - t.true(byok.getProfiles.calledOnce); - Sinon.assert.calledOnceWithMatch(byok.getProfiles, { - workspaceId: 'workspace-1', - featureKind: 'embedding', - }); -}); - -test('CopilotProviderFactory should resolve BYOK rerank routes before quota-backed routes', async t => { - const { factory, byok } = createProviderFactoryWithByokRoutes(); - - const preparedRoutes = await factory.prepareRerankRoutes( - 'gpt-4o-mini', - { - query: 'programming', - candidates: [{ text: 'React is a UI library.' }], - }, - { workspace: 'workspace-1' } - ); - const resolvedRoutes = await factory.resolveRoutes( - { modelId: 'gpt-4o-mini', outputType: ModelOutputType.Rerank }, - {}, - { workspaceId: 'workspace-1', featureKind: 'rerank' } - ); - - t.deepEqual( - preparedRoutes.map(route => route.providerId), - [] - ); - t.deepEqual( - resolvedRoutes.map(route => route.providerId), - ['byok-aaaaaaaaaaaa-openai-server-key1'] - ); - Sinon.assert.calledWithMatch(byok.getProfiles, { - workspaceId: 'workspace-1', - featureKind: 'rerank', - }); -}); - -test('CopilotProviderFactory should treat image preparation as image feature by default', async t => { - const { factory, byok } = createProviderFactoryWithByokRoutes(); - - await factory.prepareImageRoutes( - { modelId: 'gpt-image-1', outputType: ModelOutputType.Image }, - singleUserPromptMessages('draw a cat'), - { workspace: 'workspace-1' } - ); - - t.true(byok.getProfiles.calledOnce); - Sinon.assert.calledOnceWithMatch(byok.getProfiles, { - workspaceId: 'workspace-1', - featureKind: 'image', - }); -}); - -test('CopilotProviderFactory should omit quota-backed routes when quota is exhausted', async t => { - const { factory } = createProviderFactoryWithByokRoutes({ hasQuota: false }); - - const routes = await factory.resolveRoutes( - { modelId: 'gpt-5-mini', outputType: ModelOutputType.Text }, - {}, - { userId: 'user-1', workspaceId: 'workspace-1' } - ); - - t.deepEqual( - routes.map(route => route.providerId), - ['byok-aaaaaaaaaaaa-openai-server-key1'] - ); -}); - -test('CopilotProviderFactory should raise quota exceeded when only quota-backed routes match', async t => { - const { factory } = createProviderFactoryWithByokRoutes({ - byokProfiles: [], - hasQuota: false, - }); - - await t.throwsAsync( - factory.resolveRoutes( - { modelId: 'gpt-5-mini', outputType: ModelOutputType.Text }, - {}, - { userId: 'user-1', workspaceId: 'workspace-1' } - ), - { instanceOf: CopilotQuotaExceeded } - ); -}); - -test('CopilotProviderFactory should not report quota exhausted when quota-backed routes are disabled', async t => { - const { factory } = createProviderFactoryWithByokRoutes({ - byokProfiles: [], - hasQuota: true, - }); - - const routes = await factory.resolveRoutes( - { modelId: 'gpt-5-mini', outputType: ModelOutputType.Text }, - {}, - { - userId: 'user-1', - workspaceId: 'workspace-1', - quotaBackedRoutesAllowed: false, - } - ); - - t.deepEqual(routes, []); -}); - -test('selectModel should reject unknown models without online fallback', t => { - const provider = new TestOpenAIProvider(); - t.is(provider.resolveModel('online-preview'), undefined); - - const error = t.throws(() => - provider.selectModel({ - modelId: 'online-preview', - outputType: ModelOutputType.Text, - }) - ); - - t.truthy(error); - t.regex((error as Error).message, /does not support|No model supports/); -}); - -test('OpenAI oldApiStyle should resolve chat backend variants from native registry', async t => { - class LegacyOpenAIProvider extends OpenAIProvider { - override get config() { - return { - apiKey: 'test-key', - baseURL: 'https://api.openai.com/v1', - oldApiStyle: true, - }; - } - - override configured() { - return true; - } - } - - const provider = new LegacyOpenAIProvider(); - (provider as any).toolExecutorHost = { - createNativeAdapter: () => { - throw new Error('native adapter should not be used'); - }, - getTools: async () => ({}), - }; - - const prepared = await getProviderRuntimeHost(provider).prepare.chat( - 'text', - { - modelId: 'o3', - }, - singleUserPromptMessages('hello') - ); - - t.is(prepared?.route.model, 'o3'); - t.is(prepared?.route.protocol, 'openai_chat'); - t.is(prepared?.route.requestLayer, 'chat_completions'); -}); - -test('OpenAI image driver should host-materialize remote edit inputs', async t => { - const provider = new OpenAIProvider(); - (provider as any).AFFiNEConfig = { - copilot: { - providers: { - profiles: [], - defaults: {}, - openai: { apiKey: 'test-key' }, - }, - }, - }; - (provider as any).attachmentAdmissionHost = { - admitPromptAttachment: async (_attachment: unknown, context: any) => { - t.is(context.userId, 'user-1'); - t.is(context.workspaceId, 'workspace-1'); - t.is(context.sessionId, 'session-1'); - return { - id: 'att_1', - kind: 'bytes', - mimeType: 'image/png', - size: 5, - hash: 'hash', - data: 'aW1hZ2U=', - encoding: 'base64', - }; - }, - }; - const driver = provider.getExecutionDrivers()?.image; - const messages = await driver?.prepareMessages?.( - [ - { - role: 'user', - content: 'edit this', - attachments: ['https://example.com/input.png'], - }, - ], - { base_url: 'https://api.openai.com', auth_token: 'test-key' }, - { user: 'user-1', workspace: 'workspace-1', session: 'session-1' } - ); - - t.deepEqual(messages?.[0].attachments, [ - { - kind: 'bytes', - data: 'aW1hZ2U=', - encoding: 'base64', - mimeType: 'image/png', - fileName: undefined, - providerHint: undefined, - }, - ]); -}); - -test('OpenAI native request should preserve caller sampling options and defer compatibility to rust middleware', async t => { - const provider = createProvider(); - const middleware = provider.exposeMiddleware(createExecution(provider)); - - const { request } = await buildNativeRequest({ - model: 'gpt-5.4', - messages: singleUserPromptMessages('hello'), - options: { - temperature: 0.7, - topP: 0.8, - presencePenalty: 0.2, - frequencyPenalty: 0.1, - maxTokens: 128, - }, - middleware, - }); - - t.is(request.temperature, 0.7); - t.is(request.middleware, undefined); -}); - -test('ExecutionPlan should serialize routed request state and reject host-only signal', t => { - const plan = parseExecutionPlan({ - routes: [ - { - providerId: 'openai-main', - protocol: 'openai_chat', - model: 'gpt-5-mini', - backendConfig: { - base_url: 'https://api.openai.com/v1', - auth_token: 'test-key', - }, - }, - ], - request: { - kind: 'text', - cond: { modelId: 'gpt-5-mini', outputType: ModelOutputType.Text }, - messages: singleUserPromptMessages('hello'), - options: { temperature: 0.3, reasoning: true }, - }, - transport: { - kind: 'chat', - request: { - model: 'gpt-5-mini', - messages: [ - { - role: 'user', - content: [{ type: 'text', text: 'hello' }], - }, - ], - }, - }, - routePolicy: { fallbackOrder: ['openai-main'] }, - runtimePolicy: { prefer: CopilotProviderType.OpenAI, maxSteps: 4 }, - attachmentPolicy: { materializeRemoteAttachments: true }, - responsePostprocess: { mode: 'text' }, - hostContext: { currentMessages: singleUserPromptMessages('hello') }, - }); - - t.snapshot({ - fallbackOrder: plan.routePolicy.fallbackOrder, - transport: plan.transport, - }); - - const error = t.throws(() => - parseExecutionPlan({ - ...plan, - request: { - kind: 'text', - cond: { modelId: 'gpt-5-mini' }, - messages: singleUserPromptMessages('hello'), - options: { signal: new AbortController().signal }, - }, - }) - ); - - t.truthy(error); - - const hostContextError = t.throws(() => - parseExecutionPlan({ - ...plan, - hostContext: { - currentMessages: singleUserPromptMessages('hello'), - currentSessionId: 'session-1', - }, - }) - ); - - t.truthy(hostContextError); -}); - -test('ProviderDriverSpec should freeze declarative driver shape', t => { - const spec = parseProviderDriverSpec({ - driverId: 'openai-default', - providerType: CopilotProviderType.OpenAI, - models: ['gpt-5-mini'], - routes: [ - { - kind: 'text', - protocol: 'openai_chat', - requestLayer: 'chat_completions', - supportsNativeFallback: true, - requestMiddlewares: ['normalize_messages', 'openai_request_compat'], - streamMiddlewares: ['stream_event_normalize'], - }, - ], - hostOnly: { - errorMapper: 'openai', - structuredRetry: true, - }, - }); - - t.is(spec.routes[0]?.kind, 'text'); - - const error = t.throws(() => - parseProviderDriverSpec({ - ...spec, - routes: [ - { - kind: 'text', - protocol: 'openai_chat', - passthroughHelper: 'not-allowed', - }, - ], - }) - ); - - t.truthy(error); -}); - -test('NativeExecutionEngine should dispatch prepared text routes through native fallback', async t => { - const engine = createNativeExecutionEngine(); - const registry = buildProviderRegistry({ - profiles: [ - { - id: 'openai-primary', - type: CopilotProviderType.OpenAI, - config: { apiKey: '1' }, - }, - { - id: 'openai-fallback', - type: CopilotProviderType.OpenAI, - config: { apiKey: '2' }, - }, - ], - }); - const primaryProfile = registry.profiles.get('openai-primary'); - const fallbackProfile = registry.profiles.get('openai-fallback'); - if (!primaryProfile || !fallbackProfile) { - throw new Error('missing test provider profiles'); - } - - let capturedRoutes: unknown; - let called = false; - const original = (serverNativeModule as any).llmDispatchPrepared; - (serverNativeModule as any).llmDispatchPrepared = (routesJson: string) => { - called = true; - capturedRoutes = JSON.parse(routesJson); - return JSON.stringify({ - provider_id: 'openai-fallback', - response: { - id: 'chat_2', - model: 'gpt-5-mini', - message: { - role: 'assistant', - content: [{ type: 'text', text: 'fallback-ok' }], - }, - usage: { - prompt_tokens: 1, - completion_tokens: 1, - total_tokens: 2, - }, - finish_reason: 'stop', - }, - }); - }; - t.teardown(() => { - (serverNativeModule as any).llmDispatchPrepared = original; - }); - - const result = await engine.execute({ - nativeDispatch: { - chat: { - routes: [ - nativeRoute({ - providerId: 'openai-primary', - authToken: 'primary-key', - request: nativeTextRequest('hello from primary'), - }), - nativeRoute({ - providerId: 'openai-fallback', - authToken: 'fallback-key', - protocol: 'gemini', - baseUrl: GEMINI_BASE_URL, - request: nativeTextRequest('hello from fallback'), - }), - ], - prepared: { - route: preparedRoute({ - providerId: 'openai-primary', - authToken: 'primary-key', - }), - request: nativeTextRequest('hello from primary'), - tools: {}, - postprocess: { nodeTextMiddleware: [] }, - }, - hasTools: false, - }, - }, - request: { - kind: 'text', - cond: { modelId: 'gpt-5-mini' }, - messages: singleUserPromptMessages('hello'), - options: undefined, - }, - routePolicy: { fallbackOrder: ['openai-primary', 'openai-fallback'] }, - runtimePolicy: {}, - attachmentPolicy: { materializeRemoteAttachments: true }, - responsePostprocess: { mode: 'text' }, - hostPersistence: { persistAssistantTurn: true, outputKind: 'text' }, - hostContext: { - currentMessages: singleUserPromptMessages('hello'), - }, - }); - - t.is(result, 'fallback-ok'); - t.true(called); - t.snapshot(summarizePreparedDispatchRoutes(capturedRoutes)); -}); - -test('NativeExecutionEngine should record single BYOK route dispatch failure', async t => { - const byok = { - recordProviderFailure: Sinon.stub().resolves(), - recordUsage: Sinon.stub().resolves(), - }; - const engine = new NativeExecutionEngine(byok as never); - const providerId = 'byok-aaaaaaaaaaaa-openai-server-key1'; - - const original = (serverNativeModule as any).llmDispatchPrepared; - (serverNativeModule as any).llmDispatchPrepared = () => { - throw new Error('401 invalid sk-test-primary'); - }; - t.teardown(() => { - (serverNativeModule as any).llmDispatchPrepared = original; - }); - - const error = await t.throwsAsync( - engine.execute({ - nativeDispatch: { - chat: { - routes: [ - nativeRoute({ - providerId, - authToken: 'primary-key', - request: nativeTextRequest('hello'), - }), - ], - prepared: { - route: preparedRoute({ - providerId, - authToken: 'primary-key', - }), - request: nativeTextRequest('hello'), - tools: {}, - postprocess: { nodeTextMiddleware: [] }, - }, - hasTools: false, - }, - }, - request: { - kind: 'text', - cond: { modelId: 'gpt-5-mini' }, - messages: singleUserPromptMessages('hello'), - options: { - workspace: 'workspace-1', - user: 'user-1', - session: 'session-1', - featureKind: 'chat', - }, - }, - routePolicy: { fallbackOrder: [providerId] }, - runtimePolicy: {}, - attachmentPolicy: { materializeRemoteAttachments: true }, - responsePostprocess: { mode: 'text' }, - hostPersistence: { persistAssistantTurn: true, outputKind: 'text' }, - hostContext: { - currentMessages: singleUserPromptMessages('hello'), - }, - }) - ); - - t.truthy(error); - Sinon.assert.calledOnceWithMatch(byok.recordProviderFailure, { - workspaceId: 'workspace-1', - providerId, - featureKind: 'chat', - }); - Sinon.assert.notCalled(byok.recordUsage); -}); - -test('NativeExecutionEngine should reject single-route plans when no native route is prepared', async t => { - const engine = createNativeExecutionEngine(); - - const error = await t.throwsAsync( - engine.execute({ - request: { - kind: 'text', - cond: { modelId: 'gpt-5-mini' }, - messages: promptMessages(userPrompt('hello')), - options: undefined, - }, - routePolicy: { fallbackOrder: ['openai-primary'] }, - runtimePolicy: {}, - attachmentPolicy: { materializeRemoteAttachments: true }, - responsePostprocess: { mode: 'text' }, - hostPersistence: { persistAssistantTurn: true, outputKind: 'text' }, - hostContext: { - currentMessages: singleUserPromptMessages('hello'), - }, - }), - { - instanceOf: NoCopilotProviderAvailable, - } - ); - - t.true(error instanceof NoCopilotProviderAvailable); -}); - -test('NativeExecutionEngine should prefer prepared native fallback dispatch for explicit routes', async t => { - const engine = createNativeExecutionEngine(); - let capturedRoutes: unknown; - let called = false; - - const original = (serverNativeModule as any).llmDispatchPrepared; - (serverNativeModule as any).llmDispatchPrepared = (routesJson: string) => { - called = true; - capturedRoutes = JSON.parse(routesJson); - return JSON.stringify({ - provider_id: 'openai-fallback', - response: { - id: 'chat_1', - model: 'gpt-5-mini', - message: { - role: 'assistant', - content: [{ type: 'text', text: 'native-fallback-ok' }], - }, - usage: { - prompt_tokens: 1, - completion_tokens: 1, - total_tokens: 2, - }, - finish_reason: 'stop', - }, - }); - }; - t.teardown(() => { - (serverNativeModule as any).llmDispatchPrepared = original; - }); - - const result = await engine.execute({ - nativeDispatch: { - chat: { - routes: [ - nativeRoute({ - providerId: 'openai-primary', - authToken: 'primary-key', - request: nativeTextRequest('hello'), - }), - nativeRoute({ - providerId: 'openai-fallback', - authToken: 'fallback-key', - request: nativeTextRequest('hello'), - }), - ], - prepared: { - route: preparedRoute({ - providerId: 'openai-primary', - authToken: 'primary-key', - }), - request: nativeTextRequest('hello'), - tools: {}, - postprocess: { - nodeTextMiddleware: [], - }, - }, - hasTools: false, - }, - }, - request: { - kind: 'text', - cond: { modelId: 'gpt-5-mini' }, - messages: singleUserPromptMessages('hello'), - options: undefined, - }, - routePolicy: { fallbackOrder: ['openai-primary'] }, - runtimePolicy: {}, - attachmentPolicy: { materializeRemoteAttachments: true }, - responsePostprocess: { mode: 'text' }, - hostPersistence: { persistAssistantTurn: true, outputKind: 'text' }, - hostContext: {}, - }); - - t.is(result, 'native-fallback-ok'); - t.true(called); - t.snapshot(summarizePreparedDispatchRoutes(capturedRoutes)); -}); - -test('NativeExecutionEngine should stream through prepared native fallback dispatch', async t => { - const engine = createNativeExecutionEngine(); - let called = false; - - const original = (serverNativeModule as any).llmDispatchPreparedStream; - (serverNativeModule as any).llmDispatchPreparedStream = ( - _routesJson: string, - callback: (error: Error | null, arg: string) => void - ) => { - called = true; - callback( - null, - JSON.stringify({ type: 'text_delta', text: 'stream-native-ok' }) - ); - callback(null, '__AFFINE_LLM_STREAM_END__'); - return { abort() {} }; - }; - t.teardown(() => { - (serverNativeModule as any).llmDispatchPreparedStream = original; - }); - - const chunks: string[] = []; - for await (const chunk of engine.executeStream({ - nativeDispatch: { - chat: { - routes: [ - nativeRoute({ - providerId: 'openai-primary', - authToken: 'primary-key', - request: nativeTextRequest('hello'), - }), - nativeRoute({ - providerId: 'openai-fallback', - authToken: 'fallback-key', - request: nativeTextRequest('hello'), - }), - ], - prepared: { - route: preparedRoute({ - providerId: 'openai-primary', - authToken: 'primary-key', - }), - request: nativeTextRequest('hello'), - tools: {}, - postprocess: { - nodeTextMiddleware: [], - }, - }, - hasTools: false, - }, - }, - request: { - kind: 'streamText', - cond: { modelId: 'gpt-5-mini' }, - messages: promptMessages(userPrompt('hello')), - options: undefined, - }, - routePolicy: { fallbackOrder: ['openai-primary'] }, - runtimePolicy: {}, - attachmentPolicy: { materializeRemoteAttachments: true }, - responsePostprocess: { mode: 'streamText' }, - hostPersistence: { - persistAssistantTurn: true, - outputKind: 'streamText', - }, - hostContext: {}, - })) { - chunks.push(chunk); - } - - t.true(called); - t.deepEqual(chunks, ['stream-native-ok']); -}); - -test('ExecutionPlanBuilder should keep tool-loop chat routes on prepared dispatch path', async t => { - const provider = new TestOpenAIProvider(); - const toolSchema = { - answer: { - name: 'answer', - description: 'Answer', - parameters: { - type: 'object', - properties: { - value: { type: 'string' }, - }, - required: ['value'], - }, - }, - }; - const noopTool = { - answer: defineTool({ - description: 'Answer', - inputSchema: z.object({ value: z.string() }), - execute: async () => ({ ok: true }), - }), - }; - const providers = { - prepareRoutes: Sinon.stub().resolves([ - { - providerId: 'openai-primary', - provider, - execution: { providerId: 'openai-primary', profile: {} as any }, - profile: {} as any, - modelId: 'gpt-5-mini', - prepared: { - route: preparedRoute({ - providerId: 'openai-primary', - authToken: 'primary-key', - }), - request: { - ...nativeTextRequest('hello'), - tools: [toolSchema.answer], - } as LlmRequest, - tools: noopTool, - maxSteps: 4, - postprocess: { - nodeTextMiddleware: [], - }, - }, - }, - { - providerId: 'openai-fallback', - provider, - execution: { providerId: 'openai-fallback', profile: {} as any }, - profile: {} as any, - modelId: 'gpt-5-mini', - prepared: { - route: preparedRoute({ - providerId: 'openai-fallback', - authToken: 'fallback-key', - protocol: 'gemini', - baseUrl: GEMINI_BASE_URL, - }), - request: { - ...nativeTextRequest('hello'), - tools: [toolSchema.answer], - } as LlmRequest, - tools: noopTool, - maxSteps: 4, - postprocess: { - nodeTextMiddleware: [], - }, - }, - }, - ]), - }; - const metrics = { recordPlan: Sinon.stub() }; - const builder = new ExecutionPlanBuilder( - providers as never, - metrics as never - ); - - const plan = await builder.buildTextPlan({ modelId: 'gpt-5-mini' }, [ - userPrompt('hello'), - ]); - - t.is(plan.nativeDispatch?.chat?.routes.length, 2); - t.true(plan.nativeDispatch?.chat?.hasTools ?? false); - t.snapshot({ - transport: plan.transport, - preparedTools: plan.nativeDispatch?.chat?.prepared.request.tools?.map( - tool => tool.name - ), - }); -}); - -test('ExecutionPlanBuilder should keep single-route tool chat plans on prepared_routes path', async t => { - const provider = new TestOpenAIProvider(); - const toolSchema = { - answer: { - name: 'answer', - description: 'Answer', - parameters: { - type: 'object', - properties: { - value: { type: 'string' }, - }, - required: ['value'], - }, - }, - }; - const noopTool = { - answer: defineTool({ - description: 'Answer', - inputSchema: z.object({ value: z.string() }), - execute: async () => ({ ok: true }), - }), - }; - const providers = { - prepareRoutes: Sinon.stub().resolves([ - { - providerId: 'openai-primary', - provider, - execution: { providerId: 'openai-primary', profile: {} as any }, - profile: {} as any, - modelId: 'gpt-5-mini', - prepared: { - route: preparedRoute({ - providerId: 'openai-primary', - authToken: 'primary-key', - }), - request: { - ...nativeTextRequest('hello'), - tools: [toolSchema.answer], - } as LlmRequest, - tools: noopTool, - maxSteps: 4, - postprocess: { - nodeTextMiddleware: [], - }, - }, - }, - ]), - }; - const metrics = { recordPlan: Sinon.stub() }; - const builder = new ExecutionPlanBuilder( - providers as never, - metrics as never - ); - - const plan = await builder.buildTextPlan({ modelId: 'gpt-5-mini' }, [ - userPrompt('hello'), - ]); - - t.is(plan.nativeDispatch?.chat?.routes.length, 1); - t.true(plan.nativeDispatch?.chat?.hasTools ?? false); - t.snapshot(plan.transport); -}); - -test('NativeExecutionEngine should route tool-loop chat prepared routes through native dispatch', async t => { - const engine = createNativeExecutionEngine(); - let capturedRoutes: unknown; - let called = false; - let toolCallbackCount = 0; - - const original = (serverNativeModule as any) - .llmDispatchToolLoopStreamPrepared; - (serverNativeModule as any).llmDispatchToolLoopStreamPrepared = async ( - routesJson: string, - maxSteps: number, - callback: (error: Error | null, eventJson: string) => void, - toolCallback: (error: Error | null, requestJson: string) => Promise - ) => { - called = true; - capturedRoutes = JSON.parse(routesJson); - t.is(maxSteps, 4); - - const toolResult = JSON.parse( - await toolCallback( - null, - JSON.stringify({ - callId: 'call_1', - name: 'answer', - args: { value: 'native-tool-ok' }, - }) - ) - ) as { - callId: string; - name: string; - args: Record; - output: unknown; - isError?: boolean; - }; - toolCallbackCount += 1; - - callback( - null, - JSON.stringify({ - type: 'tool_call', - call_id: 'call_1', - name: 'answer', - arguments: { value: 'native-tool-ok' }, - }) - ); - callback( - null, - JSON.stringify({ - type: 'tool_result', - call_id: 'call_1', - name: toolResult.name, - arguments: toolResult.args, - output: toolResult.output, - }) - ); - callback( - null, - JSON.stringify({ type: 'text_delta', text: 'native-tool-ok' }) - ); - callback(null, JSON.stringify({ type: 'done' })); - callback(null, '__AFFINE_LLM_STREAM_END__'); - - return { abort() {} }; - }; - t.teardown(() => { - (serverNativeModule as any).llmDispatchToolLoopStreamPrepared = original; - }); - - const result = await engine.execute({ - nativeDispatch: { - chat: { - routes: [ - nativeRoute({ - providerId: 'openai-primary', - authToken: 'primary-key', - request: { - ...nativeTextRequest('hello'), - tools: [ - { - name: 'answer', - parameters: { - type: 'object', - properties: { value: { type: 'string' } }, - required: ['value'], - }, - }, - ], - }, - }), - nativeRoute({ - providerId: 'openai-fallback', - authToken: 'fallback-key', - protocol: 'gemini', - baseUrl: GEMINI_BASE_URL, - request: { - ...nativeTextRequest('hello from fallback'), - tools: [ - { - name: 'answer', - parameters: { - type: 'object', - properties: { value: { type: 'string' } }, - required: ['value'], - }, - }, - ], - }, - }), - ], - prepared: { - route: preparedRoute({ - providerId: 'openai-primary', - authToken: 'primary-key', - }), - request: { - ...nativeTextRequest('hello'), - tools: [ - { - name: 'answer', - parameters: { - type: 'object', - properties: { - value: { type: 'string' }, - }, - required: ['value'], - }, - }, - ], - }, - tools: { - answer: defineTool({ - description: 'Answer', - inputSchema: z.object({ value: z.string() }), - execute: async args => ({ value: String(args.value) }), - }), - }, - maxSteps: 4, - postprocess: { - nodeTextMiddleware: [], - }, - }, - hasTools: true, - }, - }, - request: { - kind: 'text', - cond: { modelId: 'gpt-5-mini' }, - messages: singleUserPromptMessages('hello'), - options: undefined, - }, - routePolicy: { fallbackOrder: ['openai-primary'] }, - runtimePolicy: {}, - attachmentPolicy: { materializeRemoteAttachments: true }, - responsePostprocess: { mode: 'text' }, - hostPersistence: { persistAssistantTurn: true, outputKind: 'text' }, - hostContext: { - currentMessages: singleUserPromptMessages('hello'), - }, - }); - - t.is(result, 'native-tool-ok'); - t.true(called); - t.is(toolCallbackCount, 1); - t.snapshot(summarizePreparedDispatchRoutes(capturedRoutes)); -}); - -test('ExecutionPlanBuilder should build native prepared routes for structured, image, embedding and rerank', async t => { - const provider = new TestOpenAIProvider(); - const providers = { - prepareStructuredRoutes: Sinon.stub().resolves([ - { - providerId: 'openai-primary', - provider, - execution: { providerId: 'openai-primary', profile: {} as any }, - profile: {} as any, - modelId: 'gpt-5-mini', - preparedStructured: { - route: preparedRoute({ - providerId: 'openai-primary', - authToken: 'primary-key', - }), - request: nativeStructuredRequest('hello', { - type: 'object', - properties: { ok: { type: 'boolean' } }, - required: ['ok'], - }), - }, - }, - { - providerId: 'openai-fallback', - provider, - execution: { providerId: 'openai-fallback', profile: {} as any }, - profile: {} as any, - modelId: 'gpt-5-mini', - preparedStructured: { - route: preparedRoute({ - providerId: 'openai-fallback', - authToken: 'fallback-key', - }), - request: nativeStructuredRequest('hello', { - type: 'object', - properties: { ok: { type: 'boolean' } }, - required: ['ok'], - }), - }, - }, - ]), - prepareEmbeddingRoutes: Sinon.stub().resolves([ - { - providerId: 'openai-primary', - provider, - execution: { providerId: 'openai-primary', profile: {} as any }, - profile: {} as any, - modelId: 'text-embedding-3-small', - preparedEmbedding: { - route: preparedRoute({ - providerId: 'openai-primary', - authToken: 'primary-key', - model: 'text-embedding-3-small', - }), - request: nativeEmbeddingRequest('hello'), - }, - }, - { - providerId: 'openai-fallback', - provider, - execution: { providerId: 'openai-fallback', profile: {} as any }, - profile: {} as any, - modelId: 'text-embedding-3-small', - preparedEmbedding: { - route: preparedRoute({ - providerId: 'openai-fallback', - authToken: 'fallback-key', - model: 'text-embedding-3-small', - }), - request: nativeEmbeddingRequest('hello'), - }, - }, - ]), - prepareImageRoutes: Sinon.stub().resolves([ - { - providerId: 'openai-default', - provider, - execution: { providerId: 'openai-default', profile: {} as any }, - profile: {} as any, - modelId: 'gpt-image-1', - preparedImage: { - route: preparedRoute({ - providerId: 'openai-default', - authToken: 'image-key', - protocol: 'openai_images', - model: 'gpt-image-1', - }), - request: nativeImageRequest('draw a cat'), - }, - }, - ]), - prepareRerankRoutes: Sinon.stub().resolves([ - { - providerId: 'openai-primary', - provider, - execution: { providerId: 'openai-primary', profile: {} as any }, - profile: {} as any, - modelId: 'gpt-4o-mini', - preparedRerank: { - route: preparedRoute({ - providerId: 'openai-primary', - authToken: 'primary-key', - model: 'gpt-4o-mini', - }), - request: nativeRerankRequest('programming', [ - { text: 'React is a UI library.' }, - ]), - }, - }, - { - providerId: 'openai-fallback', - provider, - execution: { providerId: 'openai-fallback', profile: {} as any }, - profile: {} as any, - modelId: 'gpt-4o-mini', - preparedRerank: { - route: preparedRoute({ - providerId: 'openai-fallback', - authToken: 'fallback-key', - model: 'gpt-4o-mini', - }), - request: nativeRerankRequest('programming', [ - { text: 'React is a UI library.' }, - ]), - }, - }, - ]), - }; - const metrics = { recordPlan: Sinon.stub() }; - const builder = new ExecutionPlanBuilder( - providers as never, - metrics as never - ); - - const structuredPlan = await builder.buildStructuredPlan( - { modelId: 'gpt-5-mini' }, - singleUserPromptMessages('hello'), - structuredOptions(z.object({ ok: z.boolean() })), - undefined, - structuredContract(z.object({ ok: z.boolean() })) - ); - const imagePlan = await builder.buildImagePlan({ modelId: 'gpt-image-1' }, [ - userPrompt('draw a cat'), - ]); - const signal = new AbortController().signal; - const embeddingPlan = await builder.buildEmbeddingPlan( - 'text-embedding-3-small', - 'hello', - { signal, dimensions: 256 } - ); - const rerankPlan = await builder.buildRerankPlan('gpt-4o-mini', { - query: 'programming', - candidates: [{ text: 'React is a UI library.' }], - }); - - t.snapshot({ - structured: { - routes: structuredPlan.nativeDispatch?.structured?.routes.length, - transport: structuredPlan.transport, - }, - image: imagePlan.nativeDispatch?.image, - embedding: { - routes: embeddingPlan.nativeDispatch?.embedding?.routes.length, - transport: embeddingPlan.transport, - }, - rerank: { - routes: rerankPlan.nativeDispatch?.rerank?.routes.length, - transport: rerankPlan.transport, - }, - }); - - t.is(embeddingPlan.hostContext.signal, signal); - t.truthy(embeddingPlan.serializable); - const serializable = embeddingPlan.serializable!; - t.deepEqual(serializable.request.options, { - dimensions: 256, - }); - t.is(serializable.routes.length, 2); - t.deepEqual(serializable.routePolicy.fallbackOrder, [ - 'openai-primary', - 'openai-fallback', - ]); -}); - -test('NativeExecutionEngine should dispatch structured prepared routes through native execution', async t => { - const engine = createNativeExecutionEngine(); - let capturedRoutes: unknown; - let called = false; - - const original = (serverNativeModule as any).llmStructuredDispatchPrepared; - (serverNativeModule as any).llmStructuredDispatchPrepared = ( - routesJson: string - ) => { - called = true; - capturedRoutes = JSON.parse(routesJson); - return JSON.stringify({ - provider_id: 'openai-fallback', - response: { - id: 'structured_1', - model: 'gpt-5-mini', - output_text: '{"ok":true}', - output_json: { ok: true }, - usage: { - prompt_tokens: 1, - completion_tokens: 1, - total_tokens: 2, - }, - finish_reason: 'stop', - }, - }); - }; - t.teardown(() => { - (serverNativeModule as any).llmStructuredDispatchPrepared = original; - }); - - const result = await engine.execute({ - nativeDispatch: { - structured: { - routes: [ - nativeRoute({ - providerId: 'openai-primary', - authToken: 'primary-key', - request: nativeStructuredRequest('hello', { - type: 'object', - properties: { ok: { type: 'boolean' } }, - required: ['ok'], - }), - }), - nativeRoute({ - providerId: 'openai-fallback', - authToken: 'fallback-key', - protocol: 'gemini', - baseUrl: GEMINI_BASE_URL, - request: nativeStructuredRequest('hello from fallback', { - type: 'object', - properties: { ok: { type: 'boolean' } }, - required: ['ok'], - }), - }), - ], - prepared: { - route: preparedRoute({ - providerId: 'openai-primary', - authToken: 'primary-key', - }), - request: nativeStructuredRequest('hello', { - type: 'object', - properties: { ok: { type: 'boolean' } }, - required: ['ok'], - }), - }, - }, - }, - request: { - kind: 'structured', - cond: { modelId: 'gpt-5-mini' }, - messages: singleUserPromptMessages('hello'), - options: structuredOptions(z.object({ ok: z.boolean() })), - }, - routePolicy: { fallbackOrder: ['openai-primary'] }, - runtimePolicy: {}, - attachmentPolicy: { materializeRemoteAttachments: true }, - responsePostprocess: { mode: 'structured' }, - hostPersistence: { persistAssistantTurn: true, outputKind: 'structured' }, - hostContext: {}, - }); - - t.is(result, '{"ok":true}'); - t.true(called); - t.snapshot(summarizePreparedDispatchRoutes(capturedRoutes)); -}); - -test('NativeExecutionEngine should dispatch embedding prepared routes through native execution', async t => { - const engine = createNativeExecutionEngine(); - let capturedRoutes: unknown; - let called = false; - - const original = (serverNativeModule as any).llmEmbeddingDispatchPrepared; - (serverNativeModule as any).llmEmbeddingDispatchPrepared = ( - routesJson: string - ) => { - called = true; - capturedRoutes = JSON.parse(routesJson); - return JSON.stringify({ - provider_id: 'openai-fallback', - response: { - model: 'text-embedding-3-small', - embeddings: [[0.1, 0.2]], - }, - }); - }; - t.teardown(() => { - (serverNativeModule as any).llmEmbeddingDispatchPrepared = original; - }); - - const result = await engine.execute({ - nativeDispatch: { - embedding: { - routes: [ - nativeRoute({ - providerId: 'openai-primary', - authToken: 'primary-key', - model: 'text-embedding-3-small', - request: nativeEmbeddingRequest('hello'), - }), - nativeRoute({ - providerId: 'openai-fallback', - authToken: 'fallback-key', - protocol: 'gemini', - model: 'text-embedding-3-small', - baseUrl: GEMINI_BASE_URL, - request: nativeEmbeddingRequest('hello fallback'), - }), - ], - prepared: { - route: preparedRoute({ - providerId: 'openai-primary', - authToken: 'primary-key', - model: 'text-embedding-3-small', - }), - request: nativeEmbeddingRequest('hello'), - }, - }, - }, - request: { - kind: 'embedding', - cond: { modelId: 'text-embedding-3-small' }, - modelId: 'text-embedding-3-small', - input: 'hello', - options: undefined, - }, - routePolicy: { fallbackOrder: ['openai-primary'] }, - runtimePolicy: {}, - attachmentPolicy: { materializeRemoteAttachments: false }, - responsePostprocess: { mode: 'embedding' }, - hostPersistence: { persistAssistantTurn: false, outputKind: 'embedding' }, - hostContext: {}, - }); - - t.snapshot({ - called, - result, - routes: summarizePreparedDispatchRoutes(capturedRoutes), - }); -}); - -test('NativeExecutionEngine should dispatch rerank prepared routes through native execution', async t => { - const engine = createNativeExecutionEngine(); - let capturedRoutes: unknown; - let called = false; - - const original = (serverNativeModule as any).llmRerankDispatchPrepared; - (serverNativeModule as any).llmRerankDispatchPrepared = ( - routesJson: string - ) => { - called = true; - capturedRoutes = JSON.parse(routesJson); - return JSON.stringify({ - provider_id: 'openai-fallback', - response: { - model: 'gpt-4o-mini', - scores: [0.9, 0.1], - }, - }); - }; - t.teardown(() => { - (serverNativeModule as any).llmRerankDispatchPrepared = original; - }); - - const result = await engine.execute({ - nativeDispatch: { - rerank: { - routes: [ - nativeRoute({ - providerId: 'openai-primary', - authToken: 'primary-key', - model: 'gpt-4o-mini', - request: nativeRerankRequest('programming', [ - { text: 'React is a UI library.' }, - ]), - }), - nativeRoute({ - providerId: 'openai-fallback', - authToken: 'fallback-key', - protocol: 'gemini', - model: 'gpt-4o-mini', - baseUrl: GEMINI_BASE_URL, - request: nativeRerankRequest('programming fallback', [ - { text: 'Vue is a UI framework.' }, - ]), - }), - ], - prepared: { - route: preparedRoute({ - providerId: 'openai-primary', - authToken: 'primary-key', - model: 'gpt-4o-mini', - }), - request: nativeRerankRequest('programming', [ - { text: 'React is a UI library.' }, - ]), - }, - }, - }, - request: { - kind: 'rerank', - cond: { modelId: 'gpt-4o-mini' }, - modelId: 'gpt-4o-mini', - request: { - query: 'programming', - candidates: [{ text: 'React is a UI library.' }], - }, - options: undefined, - }, - routePolicy: { fallbackOrder: ['openai-primary'] }, - runtimePolicy: {}, - attachmentPolicy: { materializeRemoteAttachments: false }, - responsePostprocess: { mode: 'rerank' }, - hostPersistence: { persistAssistantTurn: false, outputKind: 'rerank' }, - hostContext: {}, - }); - - t.snapshot({ - called, - result, - routes: summarizePreparedDispatchRoutes(capturedRoutes), - }); -}); - -test('NativeExecutionEngine should dispatch image plans through prepared native routes', async t => { - const engine = createNativeExecutionEngine(); - let capturedRoutes: unknown; - const original = (serverNativeModule as any).llmImageDispatchPrepared; - (serverNativeModule as any).llmImageDispatchPrepared = ( - routesJson: string - ) => { - capturedRoutes = JSON.parse(routesJson); - return JSON.stringify({ - provider_id: 'openai-image', - response: { - images: [ - { - data_base64: 'aW1hZ2U=', - media_type: 'image/webp', - }, - { - url: 'https://cdn.example.com/image.png', - media_type: 'image/png', - }, - ], - }, - }); - }; - t.teardown(() => { - (serverNativeModule as any).llmImageDispatchPrepared = original; - }); - - const request = nativeImageRequest('draw a cat'); - const imageArtifacts = await collectAsync( - engine.executeImageArtifacts({ - nativeDispatch: { - image: { - routes: [ - nativeRoute({ - providerId: 'openai-image', - authToken: 'image-key', - protocol: 'openai_images', - model: 'gpt-image-1', - request, - }), - ], - prepared: { - route: preparedRoute({ - providerId: 'openai-image', - authToken: 'image-key', - protocol: 'openai_images', - model: 'gpt-image-1', - }), - request, - }, - }, - }, - request: { - kind: 'image', - cond: { modelId: 'gpt-image-1' }, - messages: singleUserPromptMessages('draw a cat'), - options: undefined, - }, - routePolicy: { fallbackOrder: ['openai-image'] }, - runtimePolicy: {}, - attachmentPolicy: { materializeRemoteAttachments: true }, - responsePostprocess: { mode: 'image' }, - hostPersistence: { persistAssistantTurn: true, outputKind: 'image' }, - hostContext: {}, - }) - ); - - t.deepEqual(imageArtifacts, [ - { - data_base64: 'aW1hZ2U=', - media_type: 'image/webp', - }, - { - url: 'https://cdn.example.com/image.png', - media_type: 'image/png', - }, - ]); - t.snapshot(summarizePreparedDispatchRoutes(capturedRoutes)); -}); - -test('NativeExecutionEngine should record zero-token BYOK image usage without provider usage', async t => { - const byok = { - recordUsage: Sinon.stub().resolves(), - }; - const engine = new NativeExecutionEngine(byok as never); - const providerId = 'byok-aaaaaaaaaaaa-fal-server-key1'; - - const original = (serverNativeModule as any).llmImageDispatchPrepared; - (serverNativeModule as any).llmImageDispatchPrepared = () => { - return JSON.stringify({ - provider_id: providerId, - response: { - images: [ - { - url: 'https://cdn.example.com/image.png', - media_type: 'image/png', - }, - ], - }, - }); - }; - t.teardown(() => { - (serverNativeModule as any).llmImageDispatchPrepared = original; - }); - - const request = nativeImageRequest('draw a cat'); - const imageArtifacts = await collectAsync( - engine.executeImageArtifacts({ - nativeDispatch: { - image: { - routes: [ - nativeRoute({ - providerId, - authToken: 'image-key', - protocol: 'fal_image', - model: 'fal-ai/fast-sdxl', - request, - }), - ], - prepared: { - route: preparedRoute({ - providerId, - authToken: 'image-key', - protocol: 'fal_image', - model: 'fal-ai/fast-sdxl', - }), - request, - }, - }, - }, - request: { - kind: 'image', - cond: { modelId: 'fal-ai/fast-sdxl' }, - messages: singleUserPromptMessages('draw a cat'), - options: { - workspace: 'workspace-1', - user: 'user-1', - session: 'session-1', - featureKind: 'image', - }, - }, - routePolicy: { fallbackOrder: [providerId] }, - runtimePolicy: {}, - attachmentPolicy: { materializeRemoteAttachments: true }, - responsePostprocess: { mode: 'image' }, - hostPersistence: { persistAssistantTurn: true, outputKind: 'image' }, - hostContext: {}, - }) - ); - - t.is(imageArtifacts.length, 1); - Sinon.assert.calledOnceWithMatch(byok.recordUsage, { - workspaceId: 'workspace-1', - userId: 'user-1', - sessionId: 'session-1', - featureKind: 'image', - providerId, - model: 'fal-ai/fast-sdxl', - usage: undefined, - }); -}); - -test('NativeExecutionEngine should reject image plans without native dispatch', async t => { - const engine = createNativeExecutionEngine(); - - await t.throwsAsync( - collectAsync( - engine.executeImageArtifacts({ - request: { - kind: 'image', - cond: { modelId: 'gpt-image-1' }, - messages: singleUserPromptMessages('draw a cat'), - options: undefined, - }, - routePolicy: { fallbackOrder: [] }, - runtimePolicy: {}, - attachmentPolicy: { materializeRemoteAttachments: true }, - responsePostprocess: { mode: 'image' }, - hostPersistence: { persistAssistantTurn: true, outputKind: 'image' }, - hostContext: {}, - }) - ), - { instanceOf: NoCopilotProviderAvailable } - ); -}); diff --git a/packages/backend/server/src/__tests__/copilot/provider-registry.spec.ts b/packages/backend/server/src/__tests__/copilot/provider-registry.spec.ts deleted file mode 100644 index 5625b8dc02..0000000000 --- a/packages/backend/server/src/__tests__/copilot/provider-registry.spec.ts +++ /dev/null @@ -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', - ]); -}); diff --git a/packages/backend/server/src/__tests__/copilot/provider-template.spec.ts b/packages/backend/server/src/__tests__/copilot/provider-template.spec.ts deleted file mode 100644 index 319b6f3116..0000000000 --- a/packages/backend/server/src/__tests__/copilot/provider-template.spec.ts +++ /dev/null @@ -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, - }, - ]); -}); diff --git a/packages/backend/server/src/__tests__/copilot/runtime-boundaries.spec.ts b/packages/backend/server/src/__tests__/copilot/runtime-boundaries.spec.ts new file mode 100644 index 0000000000..29baa97f12 --- /dev/null +++ b/packages/backend/server/src/__tests__/copilot/runtime-boundaries.spec.ts @@ -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(); + 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 { + 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' } + ); +}); diff --git a/packages/backend/server/src/__tests__/copilot/tool-call-loop.spec.ts b/packages/backend/server/src/__tests__/copilot/tool-call-loop.spec.ts deleted file mode 100644 index 614795293e..0000000000 --- a/packages/backend/server/src/__tests__/copilot/tool-call-loop.spec.ts +++ /dev/null @@ -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( - 'toolCallbackResponse', - { - callId: 'call_1', - name: 'doc_read', - args: { doc_id: 'a1' }, - output: { markdown: '# a1' }, - } - ); - const request = llmValidateContract( - '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 | 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 - ) => { - 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; - 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[1]; - - let semanticSearchCalled = false; - const contextService = { - matchWorkspaceAll: async () => { - semanticSearchCalled = true; - return []; - }, - } as unknown as Parameters[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[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[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.', - }); -}); diff --git a/packages/backend/server/src/__tests__/copilot/transcript-contract.spec.ts b/packages/backend/server/src/__tests__/copilot/transcript-contract.spec.ts index 681e4f5592..8f2cfaca93 100644 --- a/packages/backend/server/src/__tests__/copilot/transcript-contract.spec.ts +++ b/packages/backend/server/src/__tests__/copilot/transcript-contract.spec.ts @@ -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, { blobId: 'blob-1', - recipeId: 'transcript.audio.gemini', + recipeId: 'transcript.audio', }); t.like(queuedJobs[0] as Record, { 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, { - actionId: 'transcript.audio.gemini', + actionId: 'transcript.audio', actionVersion: 'v1', }); - t.like( - (bridgeInputs[0] as { prepareStructuredRoutes: Record }) - .prepareStructuredRoutes, - { - stepId: 'transcribe', - modelId: 'gemini-3.5-flash-lite', - } - ); + t.like((bridgeInputs[0] as { step: Record }).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/ } ); diff --git a/packages/backend/server/src/__tests__/mocks/copilot.mock.ts b/packages/backend/server/src/__tests__/mocks/copilot.mock.ts index 79f54cb725..285faa4da8 100644 --- a/packages/backend/server/src/__tests__/mocks/copilot.mock.ts +++ b/packages/backend/server/src/__tests__/mocks/copilot.mock.ts @@ -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; + 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) ?? {}).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; + 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; - 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; - 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; - 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; - }, - 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 - ) { - 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[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 { - yield* this.streamText(cond, messages, options); - } - - private async *streamObjectRuntime( - cond: ModelConditions, - messages: PromptMessage[], - options?: CopilotChatOptions - ): AsyncIterableIterator { - yield* this.streamObject(cond, messages, options); - } - - async text( - cond: ModelConditions, - messages: PromptMessage[], - options: CopilotChatOptions = {} - ): Promise { - 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 { - 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 { - 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 { - 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 { - 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; - } - } - } -} diff --git a/packages/backend/server/src/__tests__/mocks/index.ts b/packages/backend/server/src/__tests__/mocks/index.ts index 291af092c5..9def4fe0ce 100644 --- a/packages/backend/server/src/__tests__/mocks/index.ts +++ b/packages/backend/server/src/__tests__/mocks/index.ts @@ -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, diff --git a/packages/backend/server/src/__tests__/mocks/prompt-service.mock.ts b/packages/backend/server/src/__tests__/mocks/prompt-service.mock.ts deleted file mode 100644 index 4572857531..0000000000 --- a/packages/backend/server/src/__tests__/mocks/prompt-service.mock.ts +++ /dev/null @@ -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(); - private readonly builtInPromptOverrides = new Map(); - - 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 } : {}), - })), - }; - } -} diff --git a/packages/backend/server/src/__tests__/models/__snapshots__/copilot-session.spec.ts.md b/packages/backend/server/src/__tests__/models/__snapshots__/copilot-session.spec.ts.md index 3042b6e577..cdf0779383 100644 --- a/packages/backend/server/src/__tests__/models/__snapshots__/copilot-session.spec.ts.md +++ b/packages/backend/server/src/__tests__/models/__snapshots__/copilot-session.spec.ts.md @@ -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', }, }, diff --git a/packages/backend/server/src/__tests__/models/__snapshots__/copilot-session.spec.ts.snap b/packages/backend/server/src/__tests__/models/__snapshots__/copilot-session.spec.ts.snap index 59064e0446..e25208930e 100644 Binary files a/packages/backend/server/src/__tests__/models/__snapshots__/copilot-session.spec.ts.snap and b/packages/backend/server/src/__tests__/models/__snapshots__/copilot-session.spec.ts.snap differ diff --git a/packages/backend/server/src/__tests__/models/copilot-context.spec.ts b/packages/backend/server/src/__tests__/models/copilot-context.spec.ts index fbd041de24..151ff5795d 100644 --- a/packages/backend/server/src/__tests__/models/copilot-context.spec.ts +++ b/packages/backend/server/src/__tests__/models/copilot-context.spec.ts @@ -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', }); diff --git a/packages/backend/server/src/__tests__/models/copilot-session.spec.ts b/packages/backend/server/src/__tests__/models/copilot-session.spec.ts index 556f7f47c4..4563ede28d 100644 --- a/packages/backend/server/src/__tests__/models/copilot-session.spec.ts +++ b/packages/backend/server/src/__tests__/models/copilot-session.spec.ts @@ -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, 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; 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()]; diff --git a/packages/backend/server/src/__tests__/utils/copilot.ts b/packages/backend/server/src/__tests__/utils/copilot.ts index d06392d333..f7398f43f5 100644 --- a/packages/backend/server/src/__tests__/utils/copilot.ts +++ b/packages/backend/server/src/__tests__/utils/copilot.ts @@ -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 { diff --git a/packages/backend/server/src/__tests__/utils/testing-module.ts b/packages/backend/server/src/__tests__/utils/testing-module.ts index f8dc831026..50cafa3caf 100644 --- a/packages/backend/server/src/__tests__/utils/testing-module.ts +++ b/packages/backend/server/src/__tests__/utils/testing-module.ts @@ -127,6 +127,7 @@ export async function createTestingModule( }, }, copilot: { + enabled: true, storage: { provider: 'assetpack', bucket: 'copilot', diff --git a/packages/backend/server/src/core/backend-runtime/__tests__/provider.spec.ts b/packages/backend/server/src/core/backend-runtime/__tests__/provider.spec.ts index 844767692a..32e0cb7219 100644 --- a/packages/backend/server/src/core/backend-runtime/__tests__/provider.spec.ts +++ b/packages/backend/server/src/core/backend-runtime/__tests__/provider.spec.ts @@ -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); }); diff --git a/packages/backend/server/src/core/backend-runtime/provider.ts b/packages/backend/server/src/core/backend-runtime/provider.ts index 92b5e71e2a..3c80dc6c52 100644 --- a/packages/backend/server/src/core/backend-runtime/provider.ts +++ b/packages/backend/server/src/core/backend-runtime/provider.ts @@ -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; +class RuntimeEventStream implements AsyncIterableIterator { + private readonly values: T[] = []; + private readonly readers: Array<(result: IteratorResult) => 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> { + 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> { + 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 { return await this.runtime.health(); } @@ -298,6 +373,148 @@ export class BackendRuntimeProvider ); } + async listByokProfiles(workspaceId: string): Promise { + return await this.measured('listByokProfiles', runtime => + runtime.listByokProfiles(workspaceId) + ); + } + + async createByokProfile( + input: CreateByokProfileInput + ): Promise { + return await this.measured('createByokProfile', runtime => + runtime.createByokProfile(input) + ); + } + + async replaceByokProfile( + input: ReplaceByokProfileInput + ): Promise { + return await this.measured('replaceByokProfile', runtime => + runtime.replaceByokProfile(input) + ); + } + + async rotateByokCredential( + input: RotateByokCredentialInput + ): Promise { + return await this.measured('rotateByokCredential', runtime => + runtime.rotateByokCredential(input) + ); + } + + async probeByokProfile( + input: ProbeByokProfileInput + ): Promise { + return await this.measured('probeByokProfile', runtime => + runtime.probeByokProfile(input) + ); + } + + async probeByokDraft( + input: ProbeByokDraftInput + ): Promise { + 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 { + return await this.measured('reorderByokProfiles', runtime => + runtime.reorderByokProfiles(input) + ); + } + + async createByokLocalLease( + input: CreateByokLocalLeaseInput + ): Promise { + 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( + input: CopilotExecuteInput, + toolCallback: (request: string) => Promise, + options: { maxSteps: number; signal?: AbortSignal } + ): AsyncIterableIterator { + const stream = new RuntimeEventStream(); + 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) diff --git a/packages/backend/server/src/core/realtime/__tests__/registry.spec.ts b/packages/backend/server/src/core/realtime/__tests__/registry.spec.ts index ff8b6a8df7..3fca3edb83 100644 --- a/packages/backend/server/src/core/realtime/__tests__/registry.spec.ts +++ b/packages/backend/server/src/core/realtime/__tests__/registry.spec.ts @@ -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, { diff --git a/packages/backend/server/src/models/copilot-byok.ts b/packages/backend/server/src/models/copilot-byok.ts index 5b3462acbc..4088522356 100644 --- a/packages/backend/server/src/models/copilot-byok.ts +++ b/packages/backend/server/src/models/copilot-byok.ts @@ -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 }, diff --git a/packages/backend/server/src/models/copilot-session.ts b/packages/backend/server/src/models/copilot-session.ts index 5e3072c7b3..e04acd3407 100644 --- a/packages/backend/server/src/models/copilot-session.ts +++ b/packages/backend/server/src/models/copilot-session.ts @@ -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, '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(value: T): T { if (typeof value !== 'string') { return value; @@ -354,7 +337,7 @@ export class CopilotSessionModel extends BaseModel { private isCountedUserMessage( message: Pick ): boolean { - return message.role === AiPromptRole.user; + return message.role === AiSessionMessageRole.user; } getSessionType(session: Pick): SessionType { @@ -418,7 +401,6 @@ export class CopilotSessionModel extends BaseModel { reuseChat = false ): Promise { 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); diff --git a/packages/backend/server/src/models/copilot-transcript-task.ts b/packages/backend/server/src/models/copilot-transcript-task.ts index 9457090d46..f7b6e2c802 100644 --- a/packages/backend/server/src/models/copilot-transcript-task.ts +++ b/packages/backend/server/src/models/copilot-transcript-task.ts @@ -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 ) { @@ -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), diff --git a/packages/backend/server/src/native.ts b/packages/backend/server/src/native.ts index ae6040c760..d4e64f0dda 100644 --- a/packages/backend/server/src/native.ts +++ b/packages/backend/server/src/native.ts @@ -1,11 +1,10 @@ import serverNativeModule, { - type ActionEvent as NativeActionEventContract, - type ActionRuntimeInput as NativeActionRuntimeInputContract, type AssertSafeUrlRequest, type BackendRuntimeHealth, type BuiltInPromptRenderContract, type BuiltInPromptSessionContract, type BuiltInPromptSpec, + type BuiltInRouteOptions, type CanonicalChatRequestContract, type CanonicalStructuredRequestContract, type CapabilityAttachmentContract, @@ -29,23 +28,14 @@ import serverNativeModule, { type LlmRerankRequestContract, type LlmStructuredRequestContract, type ModelConditionsContract, - type ModelRegistryMatchResponse, - type ModelRegistryResolveResponse, type PortalResponse, type PromptMessageContract, - type PromptMetadataContract, - type PromptMetadataResult, - type PromptRenderContract, type PromptRenderResult, - type PromptSessionContract, type PromptSessionResult, type PromptStructuredResponseContract, - type PromptTokenCountContract, - type PromptTokenCountResult, type RemoteAttachmentFetchRequest, type RemoteAttachmentFetchResponse, type RemoteMimeTypeRequest, - type RequestedModelMatchResponse, type ResolvedEntitlement, type ResolveEntitlementInput, type RuntimeBlobCleanupExecuteResult, @@ -53,7 +43,6 @@ import serverNativeModule, { type RuntimeBlobCleanupResult, type RuntimeBlobCompleteResult, type RuntimeBlobMetadataBackfillResult, - type RuntimeByokLocalLeaseRecord, type RuntimeDocBlobRefsResult, type RuntimeDocCompactionResult, type RuntimeMagicLinkOtpConsumeResult, @@ -74,6 +63,39 @@ import serverNativeModule, { type Tokenizer, } from '@affine/server-native'; +export type { + BuiltInManagedTarget, + BuiltInManagedTargetTier, + BuiltInRouteOptions, + ByokCapabilityInput, + ByokCatalogModelOutput, + ByokCatalogOutput, + ByokCatalogProviderOutput, + ByokEndpointInput, + ByokLocalLeaseOutput, + ByokModelDeclarationInput, + ByokModelProbeCheckOutput, + ByokModelProbeOutput, + ByokProbeCheckInput, + ByokProbeResultOutput, + ByokProbeStatusOutput, + ByokProfileDefinitionInput, + ByokProfileOutput, + ByokValidationOutput, + CopilotAccessProjection, + CopilotExecuteInput, + CopilotRouteCheckInput, + CopilotTargetOverrideInput, + CreateByokLocalLeaseInput, + CreateByokLocalLeaseProviderInput, + CreateByokProfileInput, + ProbeByokDraftInput, + ProbeByokProfileInput, + ReorderByokProfilesInput, + ReplaceByokProfileInput, + RotateByokCredentialInput, +} from '@affine/server-native'; + export type { AssertSafeUrlRequest, BackendRuntimeHealth, @@ -105,7 +127,6 @@ export type { RuntimeBlobCleanupResult, RuntimeBlobCompleteResult, RuntimeBlobMetadataBackfillResult, - RuntimeByokLocalLeaseRecord, RuntimeDocBlobRefsResult, RuntimeDocCompactionResult, RuntimeMagicLinkOtpConsumeResult, @@ -140,27 +161,37 @@ export type ActionRunStatus = | 'failed' | 'aborted'; -export type NativeActionEvent = Omit< - NativeActionEventContract, - 'type' | 'status' -> & { +export type NativeActionEvent = { type: ActionEventType; + actionId: string; + actionVersion: string; status?: ActionRunStatus; + stepId?: string; + attachment?: unknown; + result?: unknown; + errorCode?: string; + errorMessage?: string; + trace?: unknown; }; -export type NativeActionRuntimeInput = Omit< - NativeActionRuntimeInputContract, - 'input' -> & { - input?: unknown; +export type CopilotActionRecipe = { + actionId: string; + actionVersion: string; + slot: string; + promptRef: string; + responseContract: { schema: Record; strict: boolean } | null; + outputProjection: string; }; -import type { - CopilotProviderModel, - ModelFullConditions, -} from './plugins/copilot/providers/types'; -import type { CopilotModelBackendKind } from './plugins/copilot/runtime/contracts'; -import { parseToolLoopStreamEvent } from './plugins/copilot/runtime/contracts/shared'; +export function getCopilotActionRecipe( + actionId: string, + actionVersion?: string +): CopilotActionRecipe { + return JSON.parse( + serverNativeModule.copilotActionRecipe(actionId, actionVersion) + ) as CopilotActionRecipe; +} + import type { ToolCallRequest, ToolCallResult, @@ -406,46 +437,6 @@ export const updateRootDocMetaTitle = serverNativeModule.updateRootDocMetaTitle; const nativeLlmModule = serverNativeModule; -export type LlmProtocol = - | 'openai_chat' - | 'openai_responses' - | 'openai_images' - | 'anthropic' - | 'gemini' - | 'fal_image'; - -type LlmAttachmentReferenceMode = 'remote' | 'inline'; - -type LlmAttachmentReferenceReason = - | 'non_url_source' - | 'unsupported_scheme' - | 'generic_remote_reference' - | 'gemini_api_file_uri' - | 'gemini_api_youtube_url' - | 'gemini_api_inline_http_url'; - -type LlmAttachmentReferencePlan = { - mode: LlmAttachmentReferenceMode; - reason: LlmAttachmentReferenceReason; -}; - -type LlmRequestIntentReasoning = { - enabled?: boolean; - effort?: 'low' | 'medium' | 'high'; - budget_tokens?: number; - include_reasoning?: boolean; -}; - -type LlmRequestIntent = { - include?: string[]; - reasoning?: LlmRequestIntentReasoning; -}; - -type LlmResolvedRequestIntent = { - include?: string[]; - reasoning?: Record; -}; - export type NativePromptMessageInput = Omit< PromptMessageContract, 'role' | 'attachments' | 'params' | 'responseFormat' @@ -511,54 +502,10 @@ export type NativePromptMessageInput = Omit< }; }; -export type LlmBackendConfig = { - base_url: string; - auth_token: string; - request_layer?: - | 'anthropic' - | 'chat_completions' - | 'chat_completions_no_v1' - | 'cloudflare_workers_ai' - | 'responses' - | 'openai_images' - | 'fal' - | 'vertex' - | 'vertex_anthropic' - | 'gemini_api' - | 'gemini_vertex'; - headers?: Record; - no_streaming?: boolean; - timeout_ms?: number; -}; - -export type LlmRoutedBackend = { - provider_id: string; - protocol: LlmProtocol; - model: string; - config: LlmBackendConfig; -}; - -export type LlmPreparedDispatchRoute = LlmRoutedBackend & { - request: LlmRequest; -}; - -export type LlmPreparedStructuredDispatchRoute = LlmRoutedBackend & { - request: LlmStructuredRequest; -}; - -export type LlmPreparedEmbeddingDispatchRoute = LlmRoutedBackend & { - request: LlmEmbeddingRequestContract; -}; - -export type LlmPreparedRerankDispatchRoute = LlmRoutedBackend & { - request: LlmRerankRequestContract; -}; - export type LlmImageRequest = LlmImageRequestContract; export type LlmImageRequestBuildInput = { model: string; - protocol: LlmProtocol; messages: PromptMessageContract[]; options?: { quality?: string; @@ -568,10 +515,6 @@ export type LlmImageRequestBuildInput = { }; }; -export type LlmPreparedImageDispatchRoute = LlmRoutedBackend & { - request: LlmImageRequest; -}; - export type LlmRequest = Omit< LlmRequestContract, | 'messages' @@ -648,16 +591,6 @@ export type LlmDispatchResponse = { reasoning_details?: unknown; }; -type LlmDispatchResult = { - provider_id: string; - response: LlmDispatchResponse; -}; - -type LlmRoutedDispatchResult = { - provider_id: string; - response: TResponse; -}; - export type LlmStructuredResponse = { id: string; model: string; @@ -668,59 +601,6 @@ export type LlmStructuredResponse = { reasoning_details?: unknown; }; -class StructuredDispatchError extends Error { - constructor( - readonly code: 'invalid_structured_output', - message: string, - override readonly cause?: unknown - ) { - super(message); - this.name = 'StructuredDispatchError'; - } -} - -const INVALID_STRUCTURED_OUTPUT_PREFIX = 'invalid_structured_output:'; - -export function isInvalidStructuredOutputError( - error: unknown -): error is { code: 'invalid_structured_output' } { - return ( - typeof error === 'object' && - error !== null && - 'code' in error && - (error as { code?: unknown }).code === 'invalid_structured_output' - ); -} - -function mapStructuredDispatchError(error: unknown): never { - const message = - error instanceof Error ? error.message : String(error ?? 'Unknown error'); - - if (message.startsWith(INVALID_STRUCTURED_OUTPUT_PREFIX)) { - throw new StructuredDispatchError( - 'invalid_structured_output', - message.slice(INVALID_STRUCTURED_OUTPUT_PREFIX.length).trim(), - error - ); - } - - throw error; -} - -type LlmEmbeddingResponse = { - model: string; - embeddings: number[][]; - usage?: { - prompt_tokens: number; - total_tokens: number; - }; -}; - -type LlmRerankResponse = { - model: string; - scores: number[]; -}; - export type LlmToolLoopStreamEvent = | { type: 'message_start'; id?: string; model?: string } | { type: 'provider_selected'; provider_id: string } @@ -767,90 +647,9 @@ export type LlmToolLoopStreamEvent = } | { type: 'error'; message: string; code?: string; raw?: string }; -type LlmStreamEvent = - | LlmToolLoopStreamEvent - | { - type: 'tool_call_delta'; - call_id: string; - name?: string; - arguments_delta: string; - }; export type LlmToolCallbackRequest = ToolCallRequest; export type LlmToolCallbackResponse = ToolCallResult; -const LLM_STREAM_END_MARKER = '__AFFINE_LLM_STREAM_END__'; - -async function callLlmToolCallback( - requestJson: string, - toolCallback: ( - request: LlmToolCallbackRequest - ) => LlmToolCallbackResponse | Promise -) { - const request = llmValidateContract( - 'toolCallbackRequest', - JSON.parse(requestJson) - ); - const response = await toolCallback(request); - return JSON.stringify( - llmValidateContract('toolCallbackResponse', response) - ); -} - -function parseLlmEventJson(eventJson: string): LlmStreamEvent { - return JSON.parse(eventJson) as LlmStreamEvent; -} - -function parseLlmToolLoopStreamEvent( - eventJson: string -): LlmToolLoopStreamEvent { - const event = parseLlmEventJson(eventJson); - if ( - event.type === 'provider_selected' && - typeof event.provider_id === 'string' - ) { - return event; - } - return parseToolLoopStreamEvent(event); -} - -export function llmMatchModelCapabilities( - models: CopilotProviderModel[], - cond: ModelFullConditions -): string | undefined { - if (!nativeLlmModule.llmMatchModelCapabilities) { - throw new Error('native llm capability matcher is not available'); - } - - const response = nativeLlmModule.llmMatchModelCapabilities({ - models, - cond, - }); - - return response.modelId ?? undefined; -} - -export function llmResolveModelRegistryVariant(input: { - backendKind?: CopilotModelBackendKind; - modelId: string; -}): ModelRegistryResolveResponse { - if (!nativeLlmModule.llmResolveModelRegistryVariant) { - throw new Error('native model registry resolver is not available'); - } - - return nativeLlmModule.llmResolveModelRegistryVariant(input); -} - -export function llmMatchModelRegistry(input: { - backendKind: CopilotModelBackendKind; - cond: ModelFullConditions; -}): ModelRegistryMatchResponse { - if (!nativeLlmModule.llmMatchModelRegistry) { - throw new Error('native model registry matcher is not available'); - } - - return nativeLlmModule.llmMatchModelRegistry(input); -} - export function llmInferPromptModelConditions( messages: NativePromptMessageInput[] ): ModelConditionsContract { @@ -861,204 +660,6 @@ export function llmInferPromptModelConditions( return nativeLlmModule.llmInferPromptModelConditions(messages); } -export function llmResolveRequestedModelMatch(input: { - providerIds: string[]; - optionalModels: string[]; - requestedModelId?: string; - defaultModel?: string; -}): RequestedModelMatchResponse { - if (!nativeLlmModule.llmResolveRequestedModelMatch) { - throw new Error('native requested model matcher is not available'); - } - - return nativeLlmModule.llmResolveRequestedModelMatch(input); -} - -async function llmDispatchPrepared( - routes: LlmPreparedDispatchRoute[] -): Promise { - if (!nativeLlmModule.llmDispatchPrepared) { - throw new Error('native prepared llm dispatch is not available'); - } - const response = nativeLlmModule.llmDispatchPrepared(JSON.stringify(routes)); - const responseText = await Promise.resolve(response); - return JSON.parse(responseText) as LlmDispatchResult; -} - -type LlmChatDispatchPlanInput = { - preparedRoutes: LlmPreparedDispatchRoute[]; -}; - -export async function llmDispatchPlan( - input: LlmChatDispatchPlanInput -): Promise<{ - provider_id: string; - response: LlmDispatchResponse; -}> { - return await llmDispatchPrepared(input.preparedRoutes); -} - -export async function llmStructuredDispatch( - protocol: LlmProtocol, - backendConfig: LlmBackendConfig, - request: LlmStructuredRequest -): Promise { - if (!nativeLlmModule.llmStructuredDispatch) { - throw new Error('native llm structured dispatch is not available'); - } - try { - const response = nativeLlmModule.llmStructuredDispatch( - protocol, - JSON.stringify(backendConfig), - JSON.stringify(request) - ); - const responseText = await Promise.resolve(response); - return JSON.parse(responseText) as LlmStructuredResponse; - } catch (error) { - mapStructuredDispatchError(error); - } -} - -async function llmStructuredDispatchPrepared( - routes: LlmPreparedStructuredDispatchRoute[] -): Promise> { - if (!nativeLlmModule.llmStructuredDispatchPrepared) { - throw new Error('native prepared structured dispatch is not available'); - } - try { - const response = nativeLlmModule.llmStructuredDispatchPrepared( - JSON.stringify(routes) - ); - const responseText = await Promise.resolve(response); - return JSON.parse( - responseText - ) as LlmRoutedDispatchResult; - } catch (error) { - mapStructuredDispatchError(error); - } -} - -type LlmStructuredDispatchPlanInput = { - preparedRoutes: LlmPreparedStructuredDispatchRoute[]; -}; - -export async function llmStructuredDispatchPlan( - input: LlmStructuredDispatchPlanInput -): Promise<{ - provider_id: string; - response: LlmStructuredResponse; -}> { - return await llmStructuredDispatchPrepared(input.preparedRoutes); -} - -export async function llmEmbeddingDispatch( - protocol: LlmProtocol, - backendConfig: LlmBackendConfig, - request: LlmEmbeddingRequestContract -): Promise<{ - model: string; - embeddings: number[][]; - usage?: { - prompt_tokens: number; - total_tokens: number; - }; -}> { - if (!nativeLlmModule.llmEmbeddingDispatch) { - throw new Error('native llm embedding dispatch is not available'); - } - const response = nativeLlmModule.llmEmbeddingDispatch( - protocol, - JSON.stringify(backendConfig), - JSON.stringify(request) - ); - const responseText = await Promise.resolve(response); - return JSON.parse(responseText) as LlmEmbeddingResponse; -} - -async function llmEmbeddingDispatchPrepared( - routes: LlmPreparedEmbeddingDispatchRoute[] -): Promise> { - if (!nativeLlmModule.llmEmbeddingDispatchPrepared) { - throw new Error('native prepared embedding dispatch is not available'); - } - const response = nativeLlmModule.llmEmbeddingDispatchPrepared( - JSON.stringify(routes) - ); - const responseText = await Promise.resolve(response); - return JSON.parse( - responseText - ) as LlmRoutedDispatchResult; -} - -type LlmEmbeddingDispatchPlanInput = { - preparedRoutes: LlmPreparedEmbeddingDispatchRoute[]; -}; - -export async function llmEmbeddingDispatchPlan( - input: LlmEmbeddingDispatchPlanInput -): Promise<{ - provider_id: string; - response: { - model: string; - embeddings: number[][]; - usage?: { - prompt_tokens: number; - total_tokens: number; - }; - }; -}> { - return await llmEmbeddingDispatchPrepared(input.preparedRoutes); -} - -export async function llmRerankDispatch( - protocol: LlmProtocol, - backendConfig: LlmBackendConfig, - request: LlmRerankRequestContract -): Promise<{ - model: string; - scores: number[]; -}> { - if (!nativeLlmModule.llmRerankDispatch) { - throw new Error('native llm rerank dispatch is not available'); - } - const response = nativeLlmModule.llmRerankDispatch( - protocol, - JSON.stringify(backendConfig), - JSON.stringify(request) - ); - const responseText = await Promise.resolve(response); - return JSON.parse(responseText) as LlmRerankResponse; -} - -async function llmRerankDispatchPrepared( - routes: LlmPreparedRerankDispatchRoute[] -): Promise> { - if (!nativeLlmModule.llmRerankDispatchPrepared) { - throw new Error('native prepared llm rerank dispatch is not available'); - } - const response = nativeLlmModule.llmRerankDispatchPrepared( - JSON.stringify(routes) - ); - const responseText = await Promise.resolve(response); - return JSON.parse(responseText) as LlmRoutedDispatchResult; -} - -type LlmRerankDispatchPlanInput = { - preparedRoutes: LlmPreparedRerankDispatchRoute[]; -}; - -export async function llmRerankDispatchPlan( - input: LlmRerankDispatchPlanInput -): Promise<{ - provider_id: string; - response: { - model: string; - scores: number[]; - }; -}> { - return await llmRerankDispatchPrepared(input.preparedRoutes); -} - export type LlmImageResponse = { images: Array<{ url?: string; @@ -1085,121 +686,6 @@ export function buildLlmImageRequestFromMessages( return nativeLlmModule.llmBuildImageRequestFromMessages(request); } -async function llmImageDispatchPrepared( - routes: LlmPreparedImageDispatchRoute[] -): Promise> { - if (!nativeLlmModule.llmImageDispatchPrepared) { - throw new Error('native prepared image dispatch is not available'); - } - const response = nativeLlmModule.llmImageDispatchPrepared( - JSON.stringify(routes) - ); - const responseText = await Promise.resolve(response); - return JSON.parse(responseText) as LlmRoutedDispatchResult; -} - -export async function llmImageDispatchPlan(input: { - preparedRoutes: LlmPreparedImageDispatchRoute[]; -}): Promise<{ - provider_id: string; - response: LlmImageResponse; -}> { - return await llmImageDispatchPrepared(input.preparedRoutes); -} - -export async function llmPlanAttachmentReference( - protocol: LlmProtocol, - backendConfig: LlmBackendConfig, - source: Record | string -): Promise<{ - mode: 'remote' | 'inline'; - reason: - | 'non_url_source' - | 'unsupported_scheme' - | 'generic_remote_reference' - | 'gemini_api_file_uri' - | 'gemini_api_youtube_url' - | 'gemini_api_inline_http_url'; -}> { - if (!nativeLlmModule.llmPlanAttachmentReference) { - throw new Error('native attachment reference planning is not available'); - } - const response = nativeLlmModule.llmPlanAttachmentReference( - protocol, - JSON.stringify(backendConfig), - JSON.stringify(source) - ); - const responseText = await Promise.resolve(response); - return JSON.parse(responseText) as LlmAttachmentReferencePlan; -} - -async function llmResolveRequestIntent( - protocol: LlmProtocol, - backendConfig: LlmBackendConfig, - intent: LlmRequestIntent -): Promise { - if (!nativeLlmModule.llmResolveRequestIntent) { - throw new Error('native request intent resolution is not available'); - } - const response = nativeLlmModule.llmResolveRequestIntent( - protocol, - JSON.stringify(backendConfig), - JSON.stringify(intent) - ); - const responseText = await Promise.resolve(response); - return JSON.parse(responseText) as LlmResolvedRequestIntent; -} - -export async function llmResolveRequestIntentOptions({ - protocol, - backendConfig, - include, - reasoning, -}: { - protocol: LlmProtocol; - backendConfig: LlmBackendConfig; - include?: string[]; - reasoning?: { - enabled?: boolean; - supported?: boolean; - effort?: 'low' | 'medium' | 'high'; - budgetTokens?: number; - includeReasoning?: boolean; - }; -}): Promise<{ - include?: string[]; - reasoning?: Record; -}> { - const intent: LlmRequestIntent = { - ...(include?.length ? { include } : {}), - ...(reasoning?.enabled && reasoning.supported !== false - ? { - reasoning: { - enabled: true, - effort: reasoning.effort, - budget_tokens: reasoning.budgetTokens, - include_reasoning: reasoning.includeReasoning, - }, - } - : {}), - }; - - if (!intent.include?.length && !intent.reasoning) { - return {}; - } - - return await llmResolveRequestIntent(protocol, backendConfig, intent); -} - -export function llmRenderPrompt( - request: PromptRenderContract -): PromptRenderResult { - if (!nativeLlmModule.llmRenderPrompt) { - throw new Error('native prompt render is not available'); - } - return nativeLlmModule.llmRenderPrompt(request); -} - export function llmRenderBuiltInPrompt( request: BuiltInPromptRenderContract ): PromptRenderResult { @@ -1210,15 +696,6 @@ export function llmRenderBuiltInPrompt( return nativeLlmModule.llmRenderBuiltInPrompt(request); } -export function llmRenderSessionPrompt( - request: PromptSessionContract -): PromptSessionResult { - if (!nativeLlmModule.llmRenderSessionPrompt) { - throw new Error('native session prompt render is not available'); - } - return nativeLlmModule.llmRenderSessionPrompt(request); -} - export function llmRenderBuiltInSessionPrompt( request: BuiltInPromptSessionContract ): PromptSessionResult { @@ -1229,24 +706,6 @@ export function llmRenderBuiltInSessionPrompt( return nativeLlmModule.llmRenderBuiltInSessionPrompt(request); } -export function llmCountPromptTokens( - request: PromptTokenCountContract -): PromptTokenCountResult { - if (!nativeLlmModule.llmCountPromptTokens) { - throw new Error('native prompt token counting is not available'); - } - return nativeLlmModule.llmCountPromptTokens(request); -} - -export function llmCollectPromptMetadata( - request: PromptMetadataContract -): PromptMetadataResult { - if (!nativeLlmModule.llmCollectPromptMetadata) { - throw new Error('native prompt metadata collection is not available'); - } - return nativeLlmModule.llmCollectPromptMetadata(request); -} - export function llmListBuiltInPromptSpecs(): BuiltInPromptSpec[] { if (!nativeLlmModule.llmListBuiltInPromptSpecs) { throw new Error('native built-in prompt specs are not available'); @@ -1263,6 +722,14 @@ export function llmGetBuiltInPromptSpec( return nativeLlmModule.llmGetBuiltInPromptSpec(name); } +export function llmGetBuiltInRouteOptions( + name: string +): BuiltInRouteOptions | null { + return nativeLlmModule.llmGetBuiltInRouteOptions(name); +} + +export const llmGetByokCatalog = nativeLlmModule.llmGetByokCatalog; + function stripLlmRequestMiddleware< T extends { middleware?: { request?: string[]; stream?: string[] } }, >(request: T): T { @@ -1401,10 +868,6 @@ export function llmCanonicalJsonSchemaHash( } export type LlmContractName = - | 'executionPlan' - | 'preparedRoutes' - | 'promptRenderContract' - | 'promptSessionContract' | 'toolCallbackRequest' | 'toolCallbackResponse' | 'toolLoopEvent' @@ -1433,444 +896,11 @@ export function llmValidateContract( return nativeLlmModule.llmValidateContract(name, value) as T; } -export function llmCompileExecutionPlan(value: unknown): T { - if (!nativeLlmModule.llmCompileExecutionPlan) { - throw new Error('native execution plan compiler is not available'); - } - - return nativeLlmModule.llmCompileExecutionPlan(value) as T; -} - -export function llmNormalizePreparedRoutes(value: unknown): T { - if (!nativeLlmModule.llmNormalizePreparedRoutes) { - throw new Error('native prepared route normalizer is not available'); - } - - return nativeLlmModule.llmNormalizePreparedRoutes(value) as T; -} - -class NativeStreamAdapter implements AsyncIterableIterator { - readonly #queue: T[] = []; - readonly #waiters: ((result: IteratorResult) => void)[] = []; - readonly #handle: { abort?: () => void } | undefined; - readonly #signal?: AbortSignal; - readonly #abortListener?: () => void; - #ended = false; - - constructor( - handle: { abort?: () => void } | undefined, - signal?: AbortSignal - ) { - this.#handle = handle; - this.#signal = signal; - - if (signal?.aborted) { - this.close(true); - return; - } - - if (signal) { - this.#abortListener = () => { - this.close(true); - }; - signal.addEventListener('abort', this.#abortListener, { once: true }); - } - } - - private close(abortHandle: boolean) { - if (this.#ended) { - return; - } - - this.#ended = true; - if (this.#signal && this.#abortListener) { - this.#signal.removeEventListener('abort', this.#abortListener); - } - if (abortHandle) { - this.#handle?.abort?.(); - } - - while (this.#waiters.length) { - const waiter = this.#waiters.shift(); - waiter?.({ value: undefined as T, done: true }); - } - } - - push(value: T | null) { - if (this.#ended) { - return; - } - - if (value === null) { - this.close(false); - return; - } - - const waiter = this.#waiters.shift(); - if (waiter) { - waiter({ value, done: false }); - return; - } - - this.#queue.push(value); - } - - [Symbol.asyncIterator]() { - return this; - } - - async next(): Promise> { - if (this.#queue.length > 0) { - const value = this.#queue.shift() as T; - return { value, done: false }; - } - - if (this.#ended) { - return { value: undefined as T, done: true }; - } - - return await new Promise(resolve => { - this.#waiters.push(resolve); - }); - } - - async return(): Promise> { - this.close(true); - - return { value: undefined as T, done: true }; - } -} - -export function runNativeActionRecipePreparedStream( - input: NativeActionRuntimeInput, - signal?: AbortSignal -): AsyncIterableIterator { - if (!nativeLlmModule.runNativeActionRecipePreparedStream) { - throw new Error('native action recipe stream runtime is not available'); - } - - let adapter: NativeStreamAdapter | undefined; - const buffer: (NativeActionEvent | null)[] = []; - let pushFn = (event: NativeActionEvent | null) => { - buffer.push(event); - }; - const handle = nativeLlmModule.runNativeActionRecipePreparedStream( - input as NativeActionRuntimeInputContract, - (error, eventJson) => { - if (error) { - pushFn({ - type: 'error', - actionId: input.recipeId, - actionVersion: input.recipeVersion ?? '', - errorCode: 'action_stream_callback_error', - errorMessage: error.message, - }); - return; - } - if (eventJson === LLM_STREAM_END_MARKER) { - pushFn(null); - return; - } - try { - pushFn(JSON.parse(eventJson) as NativeActionEvent); - } catch (error) { - pushFn({ - type: 'error', - actionId: input.recipeId, - actionVersion: input.recipeVersion ?? '', - errorCode: 'action_stream_event_parse_failed', - errorMessage: - error instanceof Error - ? error.message - : 'failed to parse native action stream event', - }); - } - } - ); - adapter = new NativeStreamAdapter(handle, signal); - pushFn = event => { - adapter.push(event); - }; - for (const event of buffer) { - adapter.push(event); - } - return adapter; -} - -function llmDispatchPreparedStream( - routes: LlmPreparedDispatchRoute[], - signal?: AbortSignal -): AsyncIterableIterator { - if (!nativeLlmModule.llmDispatchPreparedStream) { - throw new Error('native prepared llm stream dispatch is not available'); - } - - let adapter: NativeStreamAdapter | undefined; - const buffer: (LlmStreamEvent | null)[] = []; - let pushFn = (event: LlmStreamEvent | null) => { - buffer.push(event); - }; - const handle = nativeLlmModule.llmDispatchPreparedStream( - JSON.stringify(routes), - (error, eventJson) => { - if (error) { - pushFn({ type: 'error', message: error.message, raw: eventJson }); - return; - } - if (eventJson === LLM_STREAM_END_MARKER) { - pushFn(null); - return; - } - try { - pushFn(parseLlmEventJson(eventJson)); - } catch (error) { - pushFn({ - type: 'error', - message: - error instanceof Error - ? error.message - : 'failed to parse native prepared stream event', - raw: eventJson, - }); - } - } - ); - adapter = new NativeStreamAdapter(handle, signal); - pushFn = event => { - adapter.push(event); - }; - for (const event of buffer) { - adapter.push(event); - } - return adapter; -} - -type LlmChatStreamDispatchPlanInput = { - preparedRoutes: LlmPreparedDispatchRoute[]; - signal?: AbortSignal; -}; - -export function llmDispatchPlanStream( - input: LlmChatStreamDispatchPlanInput -): AsyncIterableIterator< - | LlmToolLoopStreamEvent - | { - type: 'tool_call_delta'; - call_id: string; - name?: string; - arguments_delta: string; - } -> { - return llmDispatchPreparedStream(input.preparedRoutes, input.signal); -} - -export function llmDispatchToolLoopStream( - protocol: LlmProtocol, - backendConfig: LlmBackendConfig, - request: LlmRequest, - toolCallback: ( - request: LlmToolCallbackRequest - ) => LlmToolCallbackResponse | Promise, - maxSteps: number, - signal?: AbortSignal -): AsyncIterableIterator { - if (!nativeLlmModule.llmDispatchToolLoopStream) { - throw new Error('native llm tool loop dispatch is not available'); - } - - let adapter: NativeStreamAdapter | undefined; - const buffer: (LlmToolLoopStreamEvent | null)[] = []; - let pushFn = (event: LlmToolLoopStreamEvent | null) => { - buffer.push(event); - }; - const handle = nativeLlmModule.llmDispatchToolLoopStream( - protocol, - JSON.stringify(backendConfig), - JSON.stringify(request), - maxSteps, - (error, eventJson) => { - if (error) { - pushFn({ type: 'error', message: error.message, raw: eventJson }); - return; - } - if (eventJson === LLM_STREAM_END_MARKER) { - pushFn(null); - return; - } - try { - pushFn(parseLlmToolLoopStreamEvent(eventJson)); - } catch (error) { - pushFn({ - type: 'error', - message: - error instanceof Error - ? error.message - : 'failed to parse native tool loop stream event', - raw: eventJson, - }); - } - }, - async (error, requestJson) => { - if (error) { - throw error; - } - return await callLlmToolCallback(requestJson, toolCallback); - } - ); - adapter = new NativeStreamAdapter(handle, signal); - pushFn = event => { - adapter.push(event); - }; - for (const event of buffer) { - adapter.push(event); - } - return adapter; -} - -export function llmDispatchToolLoopStreamRouted( - routes: LlmRoutedBackend[], - request: LlmRequest, - toolCallback: ( - request: LlmToolCallbackRequest - ) => LlmToolCallbackResponse | Promise, - maxSteps: number, - signal?: AbortSignal -): AsyncIterableIterator { - if (!nativeLlmModule.llmDispatchToolLoopStreamRouted) { - throw new Error('native routed llm tool loop dispatch is not available'); - } - - let adapter: NativeStreamAdapter | undefined; - const buffer: (LlmToolLoopStreamEvent | null)[] = []; - let pushFn = (event: LlmToolLoopStreamEvent | null) => { - buffer.push(event); - }; - const handle = nativeLlmModule.llmDispatchToolLoopStreamRouted( - JSON.stringify(routes), - JSON.stringify(request), - maxSteps, - (error, eventJson) => { - if (error) { - pushFn({ type: 'error', message: error.message, raw: eventJson }); - return; - } - if (eventJson === LLM_STREAM_END_MARKER) { - pushFn(null); - return; - } - try { - pushFn(parseLlmToolLoopStreamEvent(eventJson)); - } catch (error) { - pushFn({ - type: 'error', - message: - error instanceof Error - ? error.message - : 'failed to parse native routed tool loop stream event', - raw: eventJson, - }); - } - }, - async (error, requestJson) => { - if (error) { - throw error; - } - return await callLlmToolCallback(requestJson, toolCallback); - } - ); - - const originalAbort = handle?.abort?.bind(handle); - if (signal) { - if (signal.aborted) { - originalAbort?.(); - } else if (originalAbort) { - signal.addEventListener('abort', () => originalAbort(), { once: true }); - } - } - - adapter = new NativeStreamAdapter(handle, signal); - pushFn = event => { - adapter?.push(event); - }; - - for (const event of buffer) { - adapter.push(event); - } - return adapter; -} - -export function llmDispatchToolLoopStreamPrepared( - routes: LlmPreparedDispatchRoute[], - toolCallback: ( - request: LlmToolCallbackRequest - ) => LlmToolCallbackResponse | Promise, - maxSteps: number, - signal?: AbortSignal -): AsyncIterableIterator { - if (!nativeLlmModule.llmDispatchToolLoopStreamPrepared) { - throw new Error('native prepared llm tool loop dispatch is not available'); - } - - let adapter: NativeStreamAdapter | undefined; - const buffer: (LlmToolLoopStreamEvent | null)[] = []; - let pushFn = (event: LlmToolLoopStreamEvent | null) => { - buffer.push(event); - }; - const handle = nativeLlmModule.llmDispatchToolLoopStreamPrepared( - JSON.stringify(routes), - maxSteps, - (error, eventJson) => { - if (error) { - pushFn({ type: 'error', message: error.message, raw: eventJson }); - return; - } - if (eventJson === LLM_STREAM_END_MARKER) { - pushFn(null); - return; - } - try { - pushFn(parseLlmToolLoopStreamEvent(eventJson)); - } catch (error) { - pushFn({ - type: 'error', - message: - error instanceof Error - ? error.message - : 'failed to parse native prepared tool loop stream event', - raw: eventJson, - }); - } - }, - async (error, requestJson) => { - if (error) { - throw error; - } - return await callLlmToolCallback(requestJson, toolCallback); - } - ); - - adapter = new NativeStreamAdapter(handle, signal); - pushFn = event => { - adapter?.push(event); - }; - - for (const event of buffer) { - adapter.push(event); - } - return adapter; -} - export { type LlmEmbeddingRequestContract as LlmEmbeddingRequest, type LlmRerankRequestContract as LlmRerankRequest, type BuiltInPromptRenderContract as NativeBuiltInPromptRenderRequest, type BuiltInPromptSessionContract as NativeBuiltInPromptSessionRenderRequest, - type PromptTokenCountContract as NativePromptCountTokensRequest, - type PromptTokenCountResult as NativePromptCountTokensResponse, - type PromptMetadataContract as NativePromptMetadataRequest, - type PromptMetadataResult as NativePromptMetadataResponse, - type PromptRenderContract as NativePromptRenderRequest, type PromptRenderResult as NativePromptRenderResponse, - type PromptSessionContract as NativePromptSessionRenderRequest, type PromptSessionResult as NativePromptSessionRenderResponse, } from '@affine/server-native'; diff --git a/packages/backend/server/src/plugins/copilot/access/index.ts b/packages/backend/server/src/plugins/copilot/access/index.ts index fd86573c04..ab82188ee8 100644 --- a/packages/backend/server/src/plugins/copilot/access/index.ts +++ b/packages/backend/server/src/plugins/copilot/access/index.ts @@ -1,2 +1 @@ export * from './feature-coverage'; -export * from './policy'; diff --git a/packages/backend/server/src/plugins/copilot/access/policy.ts b/packages/backend/server/src/plugins/copilot/access/policy.ts deleted file mode 100644 index 49c81cad14..0000000000 --- a/packages/backend/server/src/plugins/copilot/access/policy.ts +++ /dev/null @@ -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 { - const [byokProfiles, quotaBackedRoutesAvailable] = await Promise.all([ - this.getByokProfiles(context), - this.canUseQuotaBackedRoutes(context), - ]); - - return { byokProfiles, quotaBackedRoutesAvailable }; - } - - async resolveTurnRouteAccess( - context: CopilotAccessContext - ): Promise { - 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); - } - } -} diff --git a/packages/backend/server/src/plugins/copilot/availability.ts b/packages/backend/server/src/plugins/copilot/availability.ts new file mode 100644 index 0000000000..604e009f44 --- /dev/null +++ b/packages/backend/server/src/plugins/copilot/availability.ts @@ -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.'); + } +} diff --git a/packages/backend/server/src/plugins/copilot/byok/index.ts b/packages/backend/server/src/plugins/copilot/byok/index.ts index 9c40c29da1..f14ec888a7 100644 --- a/packages/backend/server/src/plugins/copilot/byok/index.ts +++ b/packages/backend/server/src/plugins/copilot/byok/index.ts @@ -1,4 +1,3 @@ export { ByokEntitlementPolicy } from './policy'; export { WorkspaceByokResolver } from './resolver'; -export { type ByokProviderRequestContext, ByokService } from './service'; export * from './types'; diff --git a/packages/backend/server/src/plugins/copilot/byok/policy.ts b/packages/backend/server/src/plugins/copilot/byok/policy.ts index ef288a36fe..4bb99c7b01 100644 --- a/packages/backend/server/src/plugins/copilot/byok/policy.ts +++ b/packages/backend/server/src/plugins/copilot/byok/policy.ts @@ -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 = diff --git a/packages/backend/server/src/plugins/copilot/byok/probe.ts b/packages/backend/server/src/plugins/copilot/byok/probe.ts deleted file mode 100644 index e63a621ff3..0000000000 --- a/packages/backend/server/src/plugins/copilot/byok/probe.ts +++ /dev/null @@ -1,94 +0,0 @@ -import { BadRequestException } from '@nestjs/common'; - -import type { safeFetch } from '../../../base'; -import { ByokProvider } from './types'; - -const TEST_TIMEOUT_MS = 10_000; -export const PROVIDER_PROBE_MAX_BYTES = 1024 * 1024; - -type ProbeFetch = typeof safeFetch; - -export async function runProviderProbe( - probeFetch: ProbeFetch, - provider: ByokProvider, - apiKey: string, - endpoint: string | null, - allowPrivateEndpoint: boolean -) { - const request = buildProbeRequest(provider, apiKey, endpoint); - const response = await probeFetch( - request.url, - { - method: request.method, - headers: request.headers, - }, - { - timeoutMs: TEST_TIMEOUT_MS, - maxRedirects: 3, - maxBytes: PROVIDER_PROBE_MAX_BYTES, - allowedHeaders: Object.keys(request.headers), - allowHttp: endpoint?.startsWith('http:') ?? false, - allowPrivateTargetOrigin: allowPrivateEndpoint, - } - ); - if (!response.ok) { - throw new BadRequestException(providerProbeFailureMessage(response.status)); - } -} - -function buildProbeRequest( - provider: ByokProvider, - apiKey: string, - endpoint: string | null -): { - method: 'GET'; - url: string; - headers: Record; -} { - switch (provider) { - case ByokProvider.openai: - return { - method: 'GET', - url: `${endpoint ?? 'https://api.openai.com/v1'}/models`, - headers: { Authorization: `Bearer ${apiKey}` }, - }; - case ByokProvider.anthropic: - return { - method: 'GET', - url: `${endpoint ?? 'https://api.anthropic.com/v1'}/models`, - headers: { - 'x-api-key': apiKey, - 'anthropic-version': '2023-06-01', - }, - }; - case ByokProvider.gemini: - return { - method: 'GET', - url: `${endpoint ?? 'https://generativelanguage.googleapis.com/v1beta'}/models`, - headers: { 'x-goog-api-key': apiKey }, - }; - case ByokProvider.fal: - return { - method: 'GET', - url: 'https://api.fal.ai/v1/models?limit=10', - headers: { Authorization: `Key ${apiKey}` }, - }; - } -} - -function providerProbeFailureMessage(status: number) { - switch (status) { - case 401: - return 'Provider rejected the BYOK key.'; - case 403: - return 'Provider rejected the BYOK key permissions.'; - case 404: - return 'Provider probe endpoint was not found.'; - case 429: - return 'Provider rate limit exceeded while testing the key.'; - default: - return status >= 500 - ? 'Provider service is unavailable.' - : `Provider key test failed with HTTP ${status}.`; - } -} diff --git a/packages/backend/server/src/plugins/copilot/byok/resolver.ts b/packages/backend/server/src/plugins/copilot/byok/resolver.ts index 371256d43e..7b48161128 100644 --- a/packages/backend/server/src/plugins/copilot/byok/resolver.ts +++ b/packages/backend/server/src/plugins/copilot/byok/resolver.ts @@ -1,3 +1,4 @@ +import { BadRequestException } from '@nestjs/common'; import { Args, Field, @@ -11,18 +12,139 @@ import { } from '@nestjs/graphql'; import { SafeIntResolver } from 'graphql-scalars'; -import { Throttle } from '../../../base'; +import { Config, Throttle } from '../../../base'; import { CurrentUser } from '../../../core/auth'; +import { BackendRuntimeProvider } from '../../../core/backend-runtime'; import { PermissionAccess } from '../../../core/permission'; import { WorkspaceType } from '../../../core/workspaces'; +import { Models } from '../../../models'; +import { llmGetByokCatalog } from '../../../native'; +import { CopilotEnabled } from '../feature'; import { ByokEntitlementPolicy } from './policy'; -import { ByokKeyConfig, ByokLocalLeaseProvider, ByokService } from './service'; -import { ByokKeyStorage, ByokKeyTestStatus, ByokProvider } from './types'; +import { + BYOK_ALLOWED_PROVIDERS, + ByokProvider, + ByokProviderSource, +} from './types'; @ObjectType() -export class WorkspaceByokKeyConfigType implements ByokKeyConfig { +class WorkspaceByokCapabilityType { + @Field(() => [String]) + input!: string[]; + + @Field(() => [String]) + output!: string[]; + + @Field(() => [String]) + features!: string[]; + + @Field(() => [String]) + attachmentKinds!: string[]; + + @Field(() => [String]) + attachmentSources!: string[]; +} + +@ObjectType() +class WorkspaceByokModelDeclarationType { + @Field(() => String) + modelId!: string; + + @Field(() => Boolean) + enabled!: boolean; + + @Field(() => [WorkspaceByokCapabilityType]) + capabilities!: WorkspaceByokCapabilityType[]; +} + +@ObjectType() +class WorkspaceByokEndpointType { + @Field(() => String) + kind!: string; + + @Field(() => String, { nullable: true }) + url!: string | null; +} + +@ObjectType() +class WorkspaceByokProfileDefinitionType { + @Field(() => SafeIntResolver) + version!: number; + + @Field(() => WorkspaceByokEndpointType) + endpoint!: WorkspaceByokEndpointType; + + @Field(() => [WorkspaceByokModelDeclarationType]) + models!: WorkspaceByokModelDeclarationType[]; +} + +@ObjectType() +class WorkspaceByokProbeStatusType { + @Field(() => String) + kind!: string; + + @Field(() => Date, { nullable: true }) + testedAt!: Date | null; + + @Field(() => String, { nullable: true }) + errorKind!: string | null; +} + +@ObjectType() +class WorkspaceByokModelProbeCheckType { + @Field(() => String) + operation!: string; + + @Field(() => WorkspaceByokProbeStatusType) + status!: WorkspaceByokProbeStatusType; +} + +@ObjectType() +class WorkspaceByokModelProbeType { + @Field(() => String) + modelId!: string; + + @Field(() => [WorkspaceByokModelProbeCheckType]) + checks!: WorkspaceByokModelProbeCheckType[]; +} + +@ObjectType() +class WorkspaceByokValidationType { + @Field(() => String) + definitionFingerprint!: string; + + @Field(() => SafeIntResolver) + credentialGeneration!: number; + + @Field(() => WorkspaceByokProbeStatusType) + connection!: WorkspaceByokProbeStatusType; + + @Field(() => [WorkspaceByokModelProbeType]) + models!: WorkspaceByokModelProbeType[]; +} + +@ObjectType() +class WorkspaceByokProbeResultType { + @Field(() => String) + definitionFingerprint!: string; + + @Field(() => Boolean) + stale!: boolean; + + @Field(() => WorkspaceByokProbeStatusType) + connection!: WorkspaceByokProbeStatusType; + + @Field(() => [WorkspaceByokModelProbeType]) + models!: WorkspaceByokModelProbeType[]; +} + +@ObjectType() +export class WorkspaceByokProfileType { @Field(() => ID) - id!: string; + profileId!: string; + + @Field(() => String) + workspaceId!: string; @Field(() => ByokProvider) provider!: ByokProvider; @@ -33,59 +155,53 @@ export class WorkspaceByokKeyConfigType implements ByokKeyConfig { @Field(() => String, { nullable: true }) description!: string | null; - @Field(() => ByokKeyStorage) - storage!: ByokKeyStorage; - - @Field(() => Boolean) - configured!: boolean; + @Field(() => WorkspaceByokProfileDefinitionType) + definition!: WorkspaceByokProfileDefinitionType; @Field(() => Boolean) enabled!: boolean; - @Field(() => String, { nullable: true }) - endpoint!: string | null; - - @Field(() => Boolean) - endpointEditable!: boolean; - @Field(() => SafeIntResolver) sortOrder!: number; - @Field(() => [String]) - capabilities!: string[]; + @Field(() => SafeIntResolver) + revision!: number; - @Field(() => ByokKeyTestStatus) - testStatus!: ByokKeyTestStatus; - - @Field(() => String, { nullable: true }) - disabledReason!: string | null; - - @Field(() => Date, { nullable: true }) - lastTestedAt!: Date | null; - - @Field(() => String, { nullable: true }) - lastTestError!: string | null; - - @Field(() => Date, { nullable: true }) - lastUsedAt!: Date | null; - - @Field(() => Date, { nullable: true }) - lastErrorAt!: Date | null; - - @Field(() => String, { nullable: true }) - lastError!: string | null; + @Field(() => WorkspaceByokValidationType, { nullable: true }) + validation!: WorkspaceByokValidationType | null; } @ObjectType() -class WorkspaceByokCapabilityWarningType { +class WorkspaceByokCatalogModelType { @Field(() => String) - featureKind!: string; + modelId!: string; @Field(() => String) - reason!: string; + displayName!: string; - @Field(() => [ByokProvider]) - requiredProviders!: ByokProvider[]; + @Field(() => Boolean) + recommended!: boolean; + + @Field(() => [WorkspaceByokCapabilityType]) + capabilities!: WorkspaceByokCapabilityType[]; +} + +@ObjectType() +class WorkspaceByokCatalogProviderType { + @Field(() => ByokProvider) + provider!: ByokProvider; + + @Field(() => [WorkspaceByokCatalogModelType]) + models!: WorkspaceByokCatalogModelType[]; +} + +@ObjectType() +class WorkspaceByokCatalogType { + @Field(() => String) + version!: string; + + @Field(() => [WorkspaceByokCatalogProviderType]) + providers!: WorkspaceByokCatalogProviderType[]; } @ObjectType() @@ -102,29 +218,20 @@ class WorkspaceByokSettingsType { @Field(() => Boolean) localEntitled!: boolean; - @Field(() => [String]) - entitlementRequired!: string[]; - - @Field(() => [WorkspaceByokKeyConfigType]) - keys!: WorkspaceByokKeyConfigType[]; + @Field(() => [WorkspaceByokProfileType]) + profiles!: WorkspaceByokProfileType[]; @Field(() => [ByokProvider]) allowedProviders!: ByokProvider[]; - @Field(() => Boolean) - localStorageSupported!: boolean; - @Field(() => Boolean) customEndpointSupported!: boolean; @Field(() => Boolean) privateEndpointSupported!: boolean; - @Field(() => Boolean) - hasAiPlan!: boolean; - - @Field(() => [WorkspaceByokCapabilityWarningType]) - warnings!: WorkspaceByokCapabilityWarningType[]; + @Field(() => WorkspaceByokCatalogType) + catalog!: WorkspaceByokCatalogType; } @ObjectType() @@ -139,18 +246,6 @@ class WorkspaceByokUsagePointType { totalTokens!: number; } -@ObjectType() -class TestWorkspaceByokConfigResultType { - @Field(() => Boolean) - ok!: boolean; - - @Field(() => ByokKeyTestStatus) - status!: ByokKeyTestStatus; - - @Field(() => String, { nullable: true }) - message!: string | null; -} - @ObjectType() class CreateWorkspaceByokLocalLeaseResultType { @Field(() => String) @@ -161,10 +256,58 @@ class CreateWorkspaceByokLocalLeaseResultType { } @InputType() -class UpsertWorkspaceByokConfigInput { - @Field(() => ID, { nullable: true }) - id?: string; +class WorkspaceByokCapabilityInput { + @Field(() => [String]) + input!: string[]; + @Field(() => [String]) + output!: string[]; + + @Field(() => [String]) + features!: string[]; + + @Field(() => [String]) + attachmentKinds!: string[]; + + @Field(() => [String]) + attachmentSources!: string[]; +} + +@InputType() +class WorkspaceByokModelDeclarationInput { + @Field(() => String) + modelId!: string; + + @Field(() => Boolean) + enabled!: boolean; + + @Field(() => [WorkspaceByokCapabilityInput]) + capabilities!: WorkspaceByokCapabilityInput[]; +} + +@InputType() +class WorkspaceByokEndpointInput { + @Field(() => String) + kind!: string; + + @Field(() => String, { nullable: true }) + url!: string | null; +} + +@InputType() +class WorkspaceByokProfileDefinitionInput { + @Field(() => SafeIntResolver) + version!: number; + + @Field(() => WorkspaceByokEndpointInput) + endpoint!: WorkspaceByokEndpointInput; + + @Field(() => [WorkspaceByokModelDeclarationInput]) + models!: WorkspaceByokModelDeclarationInput[]; +} + +@InputType() +class CreateWorkspaceByokProfileInput { @Field(() => String) workspaceId!: string; @@ -175,59 +318,125 @@ class UpsertWorkspaceByokConfigInput { name!: string; @Field(() => String, { nullable: true }) - description?: string | null; + description!: string | null; - @Field(() => ByokKeyStorage) - storage!: ByokKeyStorage; + @Field(() => String) + credential!: string; - @Field(() => String, { nullable: true }) - apiKey?: string | null; + @Field(() => WorkspaceByokProfileDefinitionInput) + definition!: WorkspaceByokProfileDefinitionInput; - @Field(() => String, { nullable: true }) - endpoint?: string | null; - - @Field(() => SafeIntResolver, { nullable: true }) - sortOrder?: number | null; - - @Field(() => Boolean, { nullable: true }) - enabled?: boolean | null; + @Field(() => Boolean) + enabled!: boolean; } @InputType() -class TestWorkspaceByokConfigInput { +class ReplaceWorkspaceByokProfileInput { + @Field(() => String) + workspaceId!: string; + + @Field(() => ID) + profileId!: string; + + @Field(() => SafeIntResolver) + expectedRevision!: number; + + @Field(() => String) + name!: string; + + @Field(() => String, { nullable: true }) + description!: string | null; + + @Field(() => WorkspaceByokProfileDefinitionInput) + definition!: WorkspaceByokProfileDefinitionInput; + + @Field(() => String, { nullable: true }) + credential!: string | null; + + @Field(() => Boolean) + enabled!: boolean; +} + +@InputType() +class RotateWorkspaceByokCredentialInput { + @Field(() => String) + workspaceId!: string; + + @Field(() => ID) + profileId!: string; + + @Field(() => SafeIntResolver) + expectedRevision!: number; + + @Field(() => String) + credential!: string; +} + +@InputType() +class WorkspaceByokProbeCheckInput { + @Field(() => String) + modelId!: string; + + @Field(() => String) + operation!: string; +} + +@InputType() +class ProbeWorkspaceByokProfileInput { + @Field(() => String) + workspaceId!: string; + + @Field(() => ID) + profileId!: string; + + @Field(() => [WorkspaceByokProbeCheckInput]) + checks!: WorkspaceByokProbeCheckInput[]; +} + +@InputType() +class ProbeWorkspaceByokDraftInput { @Field(() => String) workspaceId!: string; @Field(() => ByokProvider) provider!: ByokProvider; - @Field(() => ByokKeyStorage) - storage!: ByokKeyStorage; - @Field(() => String, { nullable: true }) - apiKey?: string | null; - - @Field(() => String, { nullable: true }) - endpoint?: string | null; + credential!: string | null; @Field(() => ID, { nullable: true }) - configId?: string | null; + profileId!: string | null; + + @Field(() => SafeIntResolver, { nullable: true }) + expectedRevision!: number | null; + + @Field(() => WorkspaceByokProfileDefinitionInput) + definition!: WorkspaceByokProfileDefinitionInput; + + @Field(() => [WorkspaceByokProbeCheckInput]) + checks!: WorkspaceByokProbeCheckInput[]; } @InputType() -class ReorderWorkspaceByokConfigsInput { +class WorkspaceByokProfileOrderInput { + @Field(() => ID) + profileId!: string; + + @Field(() => SafeIntResolver) + expectedRevision!: number; +} + +@InputType() +class ReorderWorkspaceByokProfilesInput { @Field(() => String) workspaceId!: string; - @Field(() => ByokKeyStorage) - storage!: ByokKeyStorage; - - @Field(() => [ID]) - ids!: string[]; + @Field(() => [WorkspaceByokProfileOrderInput]) + profiles!: WorkspaceByokProfileOrderInput[]; } @InputType() -class CreateWorkspaceByokLocalLeaseProviderInput implements ByokLocalLeaseProvider { +class CreateWorkspaceByokLocalLeaseProviderInput { @Field(() => ByokProvider) provider!: ByokProvider; @@ -235,19 +444,16 @@ class CreateWorkspaceByokLocalLeaseProviderInput implements ByokLocalLeaseProvid name!: string; @Field(() => String, { nullable: true }) - description?: string | null; + description!: string | null; @Field(() => String) - apiKey!: string; + credential!: string; - @Field(() => String, { nullable: true }) - endpoint?: string | null; + @Field(() => WorkspaceByokProfileDefinitionInput) + definition!: WorkspaceByokProfileDefinitionInput; - @Field(() => SafeIntResolver, { nullable: true }) - sortOrder?: number | null; - - @Field(() => Boolean, { nullable: true }) - enabled?: boolean | null; + @Field(() => Boolean) + enabled!: boolean; } @InputType() @@ -259,12 +465,15 @@ class CreateWorkspaceByokLocalLeaseInput { providers!: CreateWorkspaceByokLocalLeaseProviderInput[]; } +@CopilotEnabled() @Resolver(() => WorkspaceType) export class WorkspaceByokResolver { constructor( private readonly ac: PermissionAccess, private readonly entitlement: ByokEntitlementPolicy, - private readonly byok: ByokService + private readonly runtime: BackendRuntimeProvider, + private readonly models: Models, + private readonly config: Config ) {} @ResolveField(() => WorkspaceByokSettingsType, { @@ -275,13 +484,35 @@ export class WorkspaceByokResolver { @CurrentUser() user: CurrentUser, @Parent() workspace: WorkspaceType ) { - await this.ac - .user(user.id) - .workspace(workspace.id) - .allowLocal() - .assert('Workspace.Settings.Read'); + await this.assertRead(user.id, workspace.id); await this.entitlement.assertManagementAccess(workspace.id, user.id); - return await this.byok.getSettings(workspace.id, user.id); + const [serverEntitled, localEntitled] = + await this.entitlement.hasEntitlement(workspace.id, user.id); + const profiles = serverEntitled + ? await this.runtime.listByokProfiles(workspace.id) + : []; + const customEndpointSupported = + this.config.copilot.byok.allowCustomEndpoint; + const catalog = llmGetByokCatalog(); + return { + workspaceId: workspace.id, + entitled: serverEntitled || localEntitled, + serverEntitled, + localEntitled, + profiles: profiles.map(profile => projectProfile(profile)), + allowedProviders: [...BYOK_ALLOWED_PROVIDERS], + customEndpointSupported, + privateEndpointSupported: + customEndpointSupported && + this.config.copilot.byok.allowPrivateEndpoint, + catalog: { + ...catalog, + providers: catalog.providers.map(provider => ({ + ...provider, + provider: provider.provider as ByokProvider, + })), + }, + }; } @ResolveField(() => [WorkspaceByokUsagePointType], { @@ -294,100 +525,131 @@ export class WorkspaceByokResolver { @Args('from', { type: () => Date }) from: Date, @Args('to', { type: () => Date }) to: Date ) { - await this.ac - .user(user.id) - .workspace(workspace.id) - .allowLocal() - .assert('Workspace.Settings.Read'); + await this.assertRead(user.id, workspace.id); await this.entitlement.assertManagementAccess(workspace.id, user.id); - return await this.byok.getUsage(workspace.id, from, to); + return await this.models.copilotUsage.aggregateByDay({ + workspaceId: workspace.id, + from, + to, + providerSources: [ByokProviderSource.Server, ByokProviderSource.Local], + }); } + @Mutation(() => WorkspaceByokProfileType) @Throttle('strict') - @Mutation(() => TestWorkspaceByokConfigResultType) - async testWorkspaceByokConfig( + async createWorkspaceByokProfile( @CurrentUser() user: CurrentUser, - @Args('input') input: TestWorkspaceByokConfigInput + @Args('input') input: CreateWorkspaceByokProfileInput ) { - await this.ac - .user(user.id) - .workspace(input.workspaceId) - .allowLocal() - .assert('Workspace.Settings.Update'); - await this.entitlement.assertManagementAccess(input.workspaceId, user.id); - if (input.storage === ByokKeyStorage.server) { + await this.assertUpdate(user.id, input.workspaceId); + await this.entitlement.assertServerEntitled(input.workspaceId); + requireExplicitDescription(input); + return projectProfile( + await this.runtime.createByokProfile({ + ...input, + description: input.description ?? undefined, + definition: nativeDefinition(input.definition), + actorUserId: user.id, + }) + ); + } + + @Mutation(() => WorkspaceByokProfileType) + @Throttle('strict') + async replaceWorkspaceByokProfile( + @CurrentUser() user: CurrentUser, + @Args('input') input: ReplaceWorkspaceByokProfileInput + ) { + await this.assertUpdate(user.id, input.workspaceId); + await this.entitlement.assertServerEntitled(input.workspaceId); + requireExplicitDescription(input); + return projectProfile( + await this.runtime.replaceByokProfile({ + ...input, + description: input.description ?? undefined, + credential: input.credential ?? undefined, + definition: nativeDefinition(input.definition), + actorUserId: user.id, + }) + ); + } + + @Mutation(() => WorkspaceByokProfileType) + @Throttle('strict') + async rotateWorkspaceByokCredential( + @CurrentUser() user: CurrentUser, + @Args('input') input: RotateWorkspaceByokCredentialInput + ) { + await this.assertUpdate(user.id, input.workspaceId); + await this.entitlement.assertServerEntitled(input.workspaceId); + return projectProfile( + await this.runtime.rotateByokCredential({ + ...input, + actorUserId: user.id, + }) + ); + } + + @Mutation(() => WorkspaceByokProbeResultType) + @Throttle('strict') + async probeWorkspaceByokProfile( + @CurrentUser() user: CurrentUser, + @Args('input') input: ProbeWorkspaceByokProfileInput + ) { + await this.assertUpdate(user.id, input.workspaceId); + await this.entitlement.assertServerEntitled(input.workspaceId); + return projectProbeResult(await this.runtime.probeByokProfile(input)); + } + + @Mutation(() => WorkspaceByokProbeResultType) + @Throttle('strict') + async probeWorkspaceByokDraft( + @CurrentUser() user: CurrentUser, + @Args('input') input: ProbeWorkspaceByokDraftInput + ) { + await this.assertUpdate(user.id, input.workspaceId); + if (input.profileId) { await this.entitlement.assertServerEntitled(input.workspaceId); } else { - await this.entitlement.assertLocalEntitled(input.workspaceId, user.id); + await this.entitlement.assertEntitled(input.workspaceId, user.id); } - return await this.byok.testConfig({ ...input, userId: user.id }); - } - - @Mutation(() => WorkspaceByokKeyConfigType) - @Throttle('strict') - async upsertWorkspaceByokConfig( - @CurrentUser() user: CurrentUser, - @Args('input') input: UpsertWorkspaceByokConfigInput - ) { - await this.ac - .user(user.id) - .workspace(input.workspaceId) - .allowLocal() - .assert('Workspace.Settings.Update'); - await this.entitlement.assertManagementAccess(input.workspaceId, user.id); - await this.entitlement.assertServerEntitled(input.workspaceId); - return await this.byok.upsertConfig({ ...input, userId: user.id }); - } - - @Mutation(() => [WorkspaceByokKeyConfigType]) - @Throttle('strict') - async reorderWorkspaceByokConfigs( - @CurrentUser() user: CurrentUser, - @Args('input') input: ReorderWorkspaceByokConfigsInput - ) { - await this.ac - .user(user.id) - .workspace(input.workspaceId) - .allowLocal() - .assert('Workspace.Settings.Update'); - await this.entitlement.assertManagementAccess(input.workspaceId, user.id); - await this.entitlement.assertServerEntitled(input.workspaceId); - return await this.byok.reorderConfigs({ ...input, userId: user.id }); + return projectProbeResult( + await this.runtime.probeByokDraft({ + ...input, + credential: input.credential ?? undefined, + profileId: input.profileId ?? undefined, + expectedRevision: input.expectedRevision ?? undefined, + definition: nativeDefinition(input.definition), + }) + ); } @Mutation(() => Boolean) @Throttle('strict') - async deleteWorkspaceByokConfig( + async deleteWorkspaceByokProfile( @CurrentUser() user: CurrentUser, - @Args('id', { type: () => ID }) id: string, + @Args('profileId', { type: () => ID }) profileId: string, @Args('workspaceId', { type: () => String }) workspaceId: string ) { - await this.ac - .user(user.id) - .workspace(workspaceId) - .allowLocal() - .assert('Workspace.Settings.Update'); - await this.entitlement.assertManagementAccess(workspaceId, user.id); + await this.assertUpdate(user.id, workspaceId); await this.entitlement.assertServerEntitled(workspaceId); - return await this.byok.deleteConfig(workspaceId, id, user.id); + return await this.runtime.deleteByokProfile(workspaceId, profileId); } - @Mutation(() => Boolean) + @Mutation(() => [WorkspaceByokProfileType]) @Throttle('strict') - async clearWorkspaceByokConfigs( + async reorderWorkspaceByokProfiles( @CurrentUser() user: CurrentUser, - @Args('workspaceId', { type: () => String }) workspaceId: string, - @Args('provider', { type: () => ByokProvider, nullable: true }) - provider?: ByokProvider | null + @Args('input') input: ReorderWorkspaceByokProfilesInput ) { - await this.ac - .user(user.id) - .workspace(workspaceId) - .allowLocal() - .assert('Workspace.Settings.Update'); - await this.entitlement.assertManagementAccess(workspaceId, user.id); - await this.entitlement.assertServerEntitled(workspaceId); - return await this.byok.clearConfigs(workspaceId, provider, user.id); + await this.assertUpdate(user.id, input.workspaceId); + await this.entitlement.assertServerEntitled(input.workspaceId); + return ( + await this.runtime.reorderByokProfiles({ + ...input, + actorUserId: user.id, + }) + ).map(profile => projectProfile(profile)); } @Mutation(() => CreateWorkspaceByokLocalLeaseResultType) @@ -403,6 +665,119 @@ export class WorkspaceByokResolver { .assert('Workspace.Copilot'); await this.entitlement.assertManagementAccess(input.workspaceId, user.id); await this.entitlement.assertLocalEntitled(input.workspaceId, user.id); - return await this.byok.createLocalLease({ ...input, userId: user.id }); + input.providers.forEach(requireExplicitDescription); + const result = await this.runtime.createByokLocalLease({ + ...input, + providers: input.providers.map(provider => ({ + ...provider, + description: provider.description ?? undefined, + definition: nativeDefinition(provider.definition), + })), + userId: user.id, + }); + return { + leaseId: result.leaseId, + expiresAt: new Date(result.expiresAtMs), + }; + } + + private async assertRead(userId: string, workspaceId: string) { + await this.ac + .user(userId) + .workspace(workspaceId) + .allowLocal() + .assert('Workspace.Settings.Read'); + } + + private async assertUpdate(userId: string, workspaceId: string) { + await this.ac + .user(userId) + .workspace(workspaceId) + .allowLocal() + .assert('Workspace.Settings.Update'); + await this.entitlement.assertManagementAccess(workspaceId, userId); } } + +function requireExplicitDescription(input: { description: string | null }) { + if (!Object.hasOwn(input, 'description')) { + throw new BadRequestException('description must be provided explicitly.'); + } +} + +function nativeDefinition(input: WorkspaceByokProfileDefinitionInput) { + return { + ...input, + endpoint: { + ...input.endpoint, + url: input.endpoint.url ?? undefined, + }, + }; +} + +function projectProbe(probe: { + kind: string; + testedAtMs?: number; + errorKind?: string; +}) { + return { + kind: probe.kind, + testedAt: probe.testedAtMs ? new Date(probe.testedAtMs) : null, + errorKind: probe.errorKind ?? null, + }; +} + +function projectProbeResult(result: { + definitionFingerprint: string; + stale: boolean; + connection: { + kind: string; + testedAtMs?: number; + errorKind?: string; + }; + models: Array<{ + modelId: string; + checks: Array<{ + operation: string; + status: { + kind: string; + testedAtMs?: number; + errorKind?: string; + }; + }>; + }>; +}) { + return { + ...result, + connection: projectProbe(result.connection), + models: result.models.map(model => ({ + ...model, + checks: model.checks.map(check => ({ + ...check, + status: projectProbe(check.status), + })), + })), + }; +} + +function projectProfile( + profile: Awaited> +) { + return { + ...profile, + provider: profile.provider as ByokProvider, + validation: profile.validation + ? { + ...profile.validation, + connection: projectProbe(profile.validation.connection), + models: profile.validation.models.map(model => ({ + ...model, + checks: model.checks.map(check => ({ + ...check, + status: projectProbe(check.status), + })), + })), + } + : null, + }; +} diff --git a/packages/backend/server/src/plugins/copilot/byok/service.ts b/packages/backend/server/src/plugins/copilot/byok/service.ts deleted file mode 100644 index d2f5b12c1f..0000000000 --- a/packages/backend/server/src/plugins/copilot/byok/service.ts +++ /dev/null @@ -1,833 +0,0 @@ -import { createHash, createHmac, randomUUID } from 'node:crypto'; - -import { BadRequestException, Injectable } from '@nestjs/common'; - -import { - BadRequest, - Cache, - Config, - CryptoHelper, - metrics, - safeFetch, -} from '../../../base'; -import { Models } from '../../../models'; -import type { CopilotProviderProfile } from '../config'; -import { ByokEntitlementPolicy } from './policy'; -import { runProviderProbe } from './probe'; -import { - BYOK_ALLOWED_PROVIDERS, - type ByokFeatureKind, - ByokKeyStorage, - ByokKeyTestStatus, - ByokProvider, - ByokProviderSource, - byokProviderToCopilotType, - isByokProvider, -} from './types'; - -const LOCAL_LEASE_TTL_MS = 10 * 60 * 1000; -const BYOK_PROFILE_PRIORITY_BASE = 10_000; -const SERVER_PROFILE_PRIORITY_OFFSET = 2_000; - -export type ByokProviderRequestContext = { - userId?: string; - workspaceId?: string; - byokLeaseId?: string; -}; - -export type ByokProfileSourceFilter = { - local?: boolean; - server?: boolean; -}; - -export type ByokKeyConfig = { - id: string; - provider: ByokProvider; - name: string; - description: string | null; - storage: ByokKeyStorage; - configured: boolean; - enabled: boolean; - endpoint: string | null; - endpointEditable: boolean; - sortOrder: number; - capabilities: string[]; - testStatus: ByokKeyTestStatus; - disabledReason: string | null; - lastTestedAt: Date | null; - lastTestError: string | null; - lastUsedAt: Date | null; - lastErrorAt: Date | null; - lastError: string | null; -}; - -export type ByokSettings = { - workspaceId: string; - entitled: boolean; - serverEntitled: boolean; - localEntitled: boolean; - entitlementRequired: string[]; - keys: ByokKeyConfig[]; - allowedProviders: ByokProvider[]; - localStorageSupported: boolean; - customEndpointSupported: boolean; - privateEndpointSupported: boolean; - hasAiPlan: boolean; - warnings: Array<{ - featureKind: string; - reason: string; - requiredProviders: ByokProvider[]; - }>; -}; - -export type ByokLocalLeaseProvider = { - provider: ByokProvider; - name: string; - description?: string | null; - apiKey: string; - endpoint?: string | null; - sortOrder?: number | null; - enabled?: boolean | null; -}; - -type LocalLeasePayload = { - workspaceId: string; - userId: string; - providers: Array< - Omit & { encryptedApiKey: string } - >; -}; - -type LocalLeaseActive = { - leaseId: string; - expiresAt: string; -}; - -type ByokProfileMeta = { - source: ByokProviderSource.Server | ByokProviderSource.Local; - keyId?: string; - provider: ByokProvider; -}; - -@Injectable() -export class ByokService { - private readonly probeFetch = safeFetch; - - constructor( - private readonly models: Models, - private readonly crypto: CryptoHelper, - private readonly cache: Cache, - private readonly entitlement: ByokEntitlementPolicy, - private readonly config: Config - ) {} - - get customEndpointSupported() { - return env.selfhosted && this.config.copilot.byok.allowCustomEndpoint; - } - - get privateEndpointSupported() { - return ( - this.customEndpointSupported && - this.config.copilot.byok.allowPrivateEndpoint - ); - } - - async getSettings( - workspaceId: string, - userId?: string - ): Promise { - if (!(await this.entitlement.hasManagementAccess(workspaceId, userId))) { - return { - workspaceId, - entitled: false, - serverEntitled: false, - localEntitled: false, - entitlementRequired: ['Workspace owner or admin'], - keys: [], - allowedProviders: [...BYOK_ALLOWED_PROVIDERS], - localStorageSupported: false, - customEndpointSupported: this.customEndpointSupported, - privateEndpointSupported: this.privateEndpointSupported, - hasAiPlan: await this.entitlement.hasAiPlan(userId), - warnings: [], - }; - } - - const [serverEntitled, localEntitled] = - await this.entitlement.hasEntitlement(workspaceId, userId); - const entitled = serverEntitled || localEntitled; - if (!entitled) { - return { - workspaceId, - entitled: false, - serverEntitled: false, - localEntitled: false, - entitlementRequired: ['Pro', 'Team', 'Believer'], - keys: [], - allowedProviders: [...BYOK_ALLOWED_PROVIDERS], - localStorageSupported: false, - customEndpointSupported: this.customEndpointSupported, - privateEndpointSupported: this.privateEndpointSupported, - hasAiPlan: await this.entitlement.hasAiPlan(userId), - warnings: [], - }; - } - - const rows = serverEntitled - ? await this.models.copilotWorkspaceByokConfig.list(workspaceId) - : []; - const keys = rows.map(row => this.toKeyConfig(row)); - - return { - workspaceId, - entitled: true, - serverEntitled, - localEntitled, - entitlementRequired: ['Pro', 'Team', 'Believer'], - keys, - allowedProviders: [...BYOK_ALLOWED_PROVIDERS], - localStorageSupported: false, - customEndpointSupported: this.customEndpointSupported, - privateEndpointSupported: this.privateEndpointSupported, - hasAiPlan: await this.entitlement.hasAiPlan(userId), - warnings: this.buildWarnings(keys), - }; - } - - async upsertConfig(input: { - id?: string | null; - workspaceId: string; - provider: ByokProvider; - name: string; - description?: string | null; - storage: ByokKeyStorage; - apiKey?: string | null; - endpoint?: string | null; - sortOrder?: number | null; - enabled?: boolean | null; - userId?: string; - }): Promise { - await this.entitlement.assertManagementAccess( - input.workspaceId, - input.userId - ); - await this.entitlement.assertServerEntitled(input.workspaceId); - this.assertProvider(input.provider); - if (input.storage !== ByokKeyStorage.server) { - throw new BadRequestException('Only server BYOK keys are persisted.'); - } - const existing = input.id - ? await this.models.copilotWorkspaceByokConfig.get(input.id) - : null; - if (input.id && (!existing || existing.workspaceId !== input.workspaceId)) { - throw new BadRequest('BYOK config not found.'); - } - const encryptedApiKey = input.apiKey - ? this.crypto.encrypt(input.apiKey) - : undefined; - - if (!input.id && !encryptedApiKey) { - throw new BadRequestException('apiKey is required.'); - } - - const description = - input.description !== undefined - ? input.description?.trim() || null - : (existing?.description ?? null); - const endpoint = - input.endpoint !== undefined - ? this.normalizeEndpoint(input.endpoint) - : (existing?.endpoint ?? null); - const sortOrder = input.sortOrder ?? existing?.sortOrder ?? 0; - const enabled = input.enabled ?? existing?.enabled ?? true; - - const row = await this.models.copilotWorkspaceByokConfig.upsert({ - id: input.id, - workspaceId: input.workspaceId, - provider: input.provider, - name: input.name.trim(), - description, - encryptedApiKey, - endpoint, - sortOrder, - enabled, - userId: input.userId, - }); - - return this.toKeyConfig(row); - } - - async reorderConfigs(input: { - workspaceId: string; - storage: ByokKeyStorage; - ids: string[]; - userId?: string; - }) { - await this.entitlement.assertManagementAccess( - input.workspaceId, - input.userId - ); - await this.entitlement.assertServerEntitled(input.workspaceId); - if (input.storage !== ByokKeyStorage.server) { - throw new BadRequestException('Only server BYOK keys are persisted.'); - } - await this.models.copilotWorkspaceByokConfig.reorder( - input.workspaceId, - input.ids, - input.userId - ); - return (await this.getSettings(input.workspaceId, input.userId)).keys; - } - - async deleteConfig(workspaceId: string, id: string, _userId?: string) { - await this.entitlement.assertManagementAccess(workspaceId, _userId); - await this.entitlement.assertServerEntitled(workspaceId); - await this.models.copilotWorkspaceByokConfig.delete(workspaceId, id); - return true; - } - - async clearConfigs( - workspaceId: string, - provider: ByokProvider | null | undefined, - _userId?: string - ) { - await this.entitlement.assertManagementAccess(workspaceId, _userId); - await this.entitlement.assertServerEntitled(workspaceId); - await this.models.copilotWorkspaceByokConfig.clear(workspaceId, provider); - return true; - } - - async testConfig(input: { - workspaceId: string; - provider: ByokProvider; - storage: ByokKeyStorage; - apiKey?: string | null; - endpoint?: string | null; - configId?: string | null; - userId?: string; - }) { - await this.entitlement.assertManagementAccess( - input.workspaceId, - input.userId - ); - if (input.storage === ByokKeyStorage.server) { - await this.entitlement.assertServerEntitled(input.workspaceId); - } else { - await this.entitlement.assertLocalEntitled( - input.workspaceId, - input.userId - ); - } - this.assertProvider(input.provider); - let apiKey = input.apiKey; - let endpoint = this.normalizeEndpoint(input.endpoint); - if (!apiKey && input.configId && input.storage === ByokKeyStorage.server) { - const config = await this.models.copilotWorkspaceByokConfig.get( - input.configId - ); - if ( - !config || - config.workspaceId !== input.workspaceId || - config.provider !== input.provider - ) { - throw new BadRequestException('BYOK config not found.'); - } - apiKey = this.crypto.decrypt(config.encryptedApiKey); - endpoint = - input.endpoint !== undefined - ? endpoint - : this.normalizeEndpoint(config.endpoint); - } - if (!apiKey) { - throw new BadRequestException('apiKey is required.'); - } - - try { - await runProviderProbe( - this.probeFetch, - input.provider, - apiKey, - endpoint, - this.privateEndpointSupported - ); - if (input.configId && input.storage === ByokKeyStorage.server) { - await this.models.copilotWorkspaceByokConfig.markValidated( - input.workspaceId, - input.configId, - input.userId - ); - } - metrics.ai.counter('byok_test_key').add(1, { - workspace: input.workspaceId, - provider: input.provider, - storage: input.storage, - result: 'passed', - }); - return { ok: true, status: ByokKeyTestStatus.passed, message: null }; - } catch (error) { - const message = this.sanitizeError(error); - if (input.configId && input.storage === ByokKeyStorage.server) { - await this.models.copilotWorkspaceByokConfig.markFailure( - input.workspaceId, - input.configId, - message - ); - } - metrics.ai.counter('byok_test_key').add(1, { - workspace: input.workspaceId, - provider: input.provider, - storage: input.storage, - result: 'failed', - }); - return { ok: false, status: ByokKeyTestStatus.failed, message }; - } - } - - async createLocalLease(input: { - workspaceId: string; - providers: ByokLocalLeaseProvider[]; - userId: string; - }) { - await this.entitlement.assertManagementAccess( - input.workspaceId, - input.userId - ); - await this.entitlement.assertLocalEntitled(input.workspaceId, input.userId); - const providers = input.providers.map(provider => { - this.assertProvider(provider.provider); - const endpoint = this.normalizeEndpoint(provider.endpoint); - return { ...provider, endpoint }; - }); - const activeCacheKey = this.localLeaseActiveCacheKey({ - ...input, - providers, - }); - const activeLease = await this.getActiveLocalLease(activeCacheKey); - if (activeLease) return activeLease; - - const leaseId = randomUUID(); - const expiresAt = new Date(Date.now() + LOCAL_LEASE_TTL_MS); - const payload: LocalLeasePayload = { - workspaceId: input.workspaceId, - userId: input.userId, - providers: providers.map(provider => ({ - provider: provider.provider, - name: provider.name, - description: provider.description, - encryptedApiKey: this.crypto.encrypt(provider.apiKey), - endpoint: provider.endpoint, - sortOrder: provider.sortOrder, - enabled: provider.enabled, - })), - }; - await this.cache.set(this.leaseCacheKey(leaseId), payload, { - ttl: LOCAL_LEASE_TTL_MS, - }); - const registered = await this.cache.setnx( - activeCacheKey, - { leaseId, expiresAt: expiresAt.toISOString() }, - { ttl: LOCAL_LEASE_TTL_MS } - ); - if (!registered) { - const current = await this.getActiveLocalLease(activeCacheKey); - if (current) { - await this.cache.delete(this.leaseCacheKey(leaseId)); - return current; - } - } - return { leaseId, expiresAt }; - } - - async getProfiles( - context: ByokProviderRequestContext = {}, - sources: ByokProfileSourceFilter = { local: true, server: true } - ): Promise { - if (!context.workspaceId) { - return []; - } - const [localEntitled, serverEntitled] = await Promise.all([ - this.entitlement.hasLocalEntitlement(context.workspaceId, context.userId), - this.entitlement.hasServerEntitlement(context.workspaceId), - ]); - const [localProfiles, serverProfiles] = await Promise.all([ - sources.local && localEntitled - ? this.getLocalProfiles(context) - : Promise.resolve([]), - sources.server && serverEntitled - ? this.getServerProfiles(context.workspaceId) - : Promise.resolve([]), - ]); - - return [...localProfiles, ...serverProfiles]; - } - - async recordUsage(input: { - workspaceId?: string; - userId?: string; - providerId?: string; - model?: string | null; - featureKind: ByokFeatureKind; - sessionId?: string; - taskId?: string; - actionId?: string; - billingUnitId?: string; - usage?: { - prompt_tokens?: number; - completion_tokens?: number; - total_tokens?: number; - cached_tokens?: number; - }; - }) { - if (!input.workspaceId || !input.providerId) return; - const meta = this.parseProfileMeta(input.providerId, input.workspaceId); - if (!meta) return; - - metrics.ai.counter('byok_usage').add(1, { - workspace: input.workspaceId, - provider: meta.provider, - source: meta.source, - feature: input.featureKind, - }); - await this.models.copilotUsage.create({ - workspaceId: input.workspaceId, - userId: input.userId, - provider: meta.provider, - providerSource: meta.source, - featureKind: input.featureKind, - model: input.model ?? null, - sessionId: input.sessionId, - taskId: input.taskId, - actionId: input.actionId, - billingUnitId: input.billingUnitId, - promptTokens: input.usage?.prompt_tokens ?? 0, - completionTokens: input.usage?.completion_tokens ?? 0, - totalTokens: input.usage?.total_tokens ?? 0, - cachedTokens: input.usage?.cached_tokens ?? 0, - }); - if (meta.source === ByokProviderSource.Server && meta.keyId) { - await this.models.copilotWorkspaceByokConfig.touchUsed( - input.workspaceId, - meta.keyId - ); - } - } - - async recordProviderFailure(input: { - workspaceId?: string; - providerId?: string; - featureKind: ByokFeatureKind; - error: unknown; - }) { - if (!input.workspaceId || !input.providerId) return; - const meta = this.parseProfileMeta(input.providerId, input.workspaceId); - if (!meta) return; - - const message = this.sanitizeError(input.error); - metrics.ai.counter('byok_route_failure').add(1, { - workspace: input.workspaceId, - provider: meta.provider, - source: meta.source, - feature: input.featureKind, - }); - if (meta.source === ByokProviderSource.Server && meta.keyId) { - await this.models.copilotWorkspaceByokConfig.markFailure( - input.workspaceId, - meta.keyId, - message - ); - } - } - - async getUsage(workspaceId: string, from: Date, to: Date) { - return await this.models.copilotUsage.aggregateByDay({ - workspaceId, - from, - to, - providerSources: [ByokProviderSource.Server, ByokProviderSource.Local], - }); - } - - private async getServerProfiles(workspaceId: string) { - const rows = - await this.models.copilotWorkspaceByokConfig.listEnabled(workspaceId); - - return rows - .filter(row => isByokProvider(row.provider)) - .map((row, index): CopilotProviderProfile => { - const provider = row.provider as ByokProvider; - return { - id: this.profileId(workspaceId, provider, row.id, 'server'), - type: byokProviderToCopilotType(provider), - priority: - BYOK_PROFILE_PRIORITY_BASE - SERVER_PROFILE_PRIORITY_OFFSET - index, - config: this.providerConfig( - provider, - row.encryptedApiKey, - row.endpoint - ), - } as CopilotProviderProfile; - }); - } - - private async getLocalProfiles(context: ByokProviderRequestContext) { - if (!context.byokLeaseId || !context.workspaceId || !context.userId) { - return []; - } - if ( - !(await this.entitlement.hasManagementAccess( - context.workspaceId, - context.userId - )) - ) { - return []; - } - const lease = await this.cache.get( - this.leaseCacheKey(context.byokLeaseId) - ); - if ( - !lease || - lease.workspaceId !== context.workspaceId || - lease.userId !== context.userId - ) { - return []; - } - return lease.providers - .filter(provider => provider.enabled !== false) - .map((provider, index): CopilotProviderProfile => { - return { - id: this.profileId( - context.workspaceId ?? lease.workspaceId, - provider.provider, - `${index}`, - 'local' - ), - type: byokProviderToCopilotType(provider.provider), - priority: BYOK_PROFILE_PRIORITY_BASE - index, - config: this.providerConfig( - provider.provider, - provider.encryptedApiKey, - provider.endpoint ?? null - ), - } as CopilotProviderProfile; - }); - } - - private providerConfig( - provider: ByokProvider, - encryptedApiKey: string, - endpoint: string | null - ) { - const apiKey = this.crypto.decrypt(encryptedApiKey); - switch (provider) { - case ByokProvider.openai: - case ByokProvider.gemini: - case ByokProvider.anthropic: - return { apiKey, ...(endpoint ? { baseURL: endpoint } : {}) }; - case ByokProvider.fal: - return { apiKey }; - } - } - - private profileId( - workspaceId: string, - provider: ByokProvider, - keyId: string, - storage: 'server' | 'local' - ) { - const hash = this.workspaceHash(workspaceId); - const sanitizedKeyId = keyId.replaceAll(/[^a-zA-Z0-9-_]/g, ''); - return storage === 'local' - ? `byok-${hash}-${provider}-local-${sanitizedKeyId}` - : `byok-${hash}-${provider}-${sanitizedKeyId}`; - } - - parseProfileMeta( - providerId: string, - workspaceId?: string - ): ByokProfileMeta | null { - const match = - /^byok-([a-f0-9]{12})-(openai|anthropic|gemini|fal)-(.+)$/.exec( - providerId - ); - if (!match) return null; - if (workspaceId && match[1] !== this.workspaceHash(workspaceId)) { - return null; - } - - const keyId = match[3]; - return { - provider: match[2] as ByokProvider, - source: keyId.startsWith('local-') - ? ByokProviderSource.Local - : ByokProviderSource.Server, - keyId: keyId.startsWith('local-') ? undefined : keyId, - }; - } - - private toKeyConfig(row: { - id: string; - provider: string; - name: string; - description: string | null; - endpoint: string | null; - sortOrder: number; - enabled: boolean; - disabledReason: string | null; - lastValidatedAt: Date | null; - lastValidationError: string | null; - lastUsedAt: Date | null; - lastErrorAt: Date | null; - lastError: string | null; - }): ByokKeyConfig { - const provider = row.provider as ByokProvider; - return { - id: row.id, - provider, - name: row.name, - description: row.description, - storage: ByokKeyStorage.server, - configured: true, - enabled: row.enabled, - endpoint: row.endpoint, - endpointEditable: this.customEndpointSupported, - sortOrder: row.sortOrder, - capabilities: this.capabilities(provider, 'server'), - testStatus: row.lastValidationError - ? ByokKeyTestStatus.failed - : row.lastValidatedAt - ? ByokKeyTestStatus.passed - : ByokKeyTestStatus.untested, - disabledReason: row.disabledReason, - lastTestedAt: row.lastValidatedAt, - lastTestError: row.lastValidationError, - lastUsedAt: row.lastUsedAt, - lastErrorAt: row.lastErrorAt, - lastError: row.lastError, - }; - } - - private capabilities(provider: ByokProvider, storage: 'server' | 'local') { - switch (provider) { - case ByokProvider.openai: - return ['Text', 'Image input', 'Actions', 'Image generate']; - case ByokProvider.anthropic: - return ['Text', 'Image input']; - case ByokProvider.gemini: - return storage === 'server' - ? [ - 'Text', - 'Image input', - 'Actions', - 'Image generate', - 'Transcript', - 'Indexing', - ] - : ['Text', 'Image input', 'Actions', 'Image generate']; - case ByokProvider.fal: - return ['Image generate']; - } - } - - private buildWarnings(keys: ByokKeyConfig[]) { - const activeServerGemini = keys.some( - key => - key.provider === ByokProvider.gemini && - key.storage === ByokKeyStorage.server && - key.enabled - ); - if (activeServerGemini) { - return []; - } - return [ - { - featureKind: 'transcript', - reason: - 'Transcript and workspace indexing require a server Gemini BYOK key or AFFiNE AI plan fallback.', - requiredProviders: [ByokProvider.gemini], - }, - { - featureKind: 'workspace_indexing', - reason: - 'Workspace indexing requires a server Gemini BYOK key or AFFiNE AI plan fallback.', - requiredProviders: [ByokProvider.gemini], - }, - ]; - } - - private normalizeEndpoint(endpoint?: string | null) { - if (!endpoint) return null; - if (!this.customEndpointSupported) { - throw new BadRequestException('Custom BYOK endpoint is not supported.'); - } - let parsed: URL; - try { - parsed = new URL(endpoint); - } catch { - throw new BadRequestException('Invalid BYOK endpoint.'); - } - if (!['https:', 'http:'].includes(parsed.protocol)) { - throw new BadRequestException('BYOK endpoint must use HTTP or HTTPS.'); - } - return parsed.toString().replace(/\/$/, ''); - } - - private assertProvider(provider: ByokProvider) { - if (!BYOK_ALLOWED_PROVIDERS.includes(provider)) { - throw new BadRequestException('Unsupported BYOK provider.'); - } - } - - private sanitizeError(error: unknown) { - if (error instanceof Error && error.name === 'AbortError') { - return 'Provider key test timed out.'; - } - if (error instanceof BadRequestException && error.message) { - return error.message.slice(0, 300); - } - return 'Provider request failed.'; - } - - private workspaceHash(workspaceId: string) { - return createHash('sha256').update(workspaceId).digest('hex').slice(0, 12); - } - - private leaseCacheKey(leaseId: string) { - return `copilot:byok:lease:${leaseId}`; - } - - private async getActiveLocalLease(activeCacheKey: string) { - const active = await this.cache.get(activeCacheKey); - if (!active) return null; - if (await this.cache.has(this.leaseCacheKey(active.leaseId))) { - return { leaseId: active.leaseId, expiresAt: new Date(active.expiresAt) }; - } - await this.cache.delete(activeCacheKey); - return null; - } - - private localLeaseActiveCacheKey(input: { - workspaceId: string; - userId: string; - providers: ByokLocalLeaseProvider[]; - }) { - const fingerprint = createHmac( - 'sha256', - this.crypto.keyPair.sha256.privateKey - ) - .update( - JSON.stringify( - input.providers.map(provider => ({ - provider: provider.provider, - name: provider.name, - description: provider.description ?? null, - apiKey: provider.apiKey, - endpoint: provider.endpoint ?? null, - sortOrder: provider.sortOrder ?? 0, - enabled: provider.enabled ?? true, - })) - ) - ) - .digest('hex'); - return `copilot:byok:lease:active:${input.workspaceId}:${input.userId}:${fingerprint}`; - } -} diff --git a/packages/backend/server/src/plugins/copilot/compat/history-projector.ts b/packages/backend/server/src/plugins/copilot/compat/history-projector.ts index 1e9d2fa7a8..34c1787117 100644 --- a/packages/backend/server/src/plugins/copilot/compat/history-projector.ts +++ b/packages/backend/server/src/plugins/copilot/compat/history-projector.ts @@ -1,5 +1,5 @@ import { Injectable } from '@nestjs/common'; -import { AiPromptRole } from '@prisma/client'; +import { AiSessionMessageRole } from '@prisma/client'; import type { Conversation, Turn } from '../core'; import { chatMessageFromTurn } from '../core'; @@ -16,7 +16,6 @@ export type CanonicalConversationHistory = { conversation: Conversation; turns: Turn[]; prompt: ResolvedPrompt; - tokenCost: number; }; export type CanonicalConversationMeta = Omit< @@ -35,7 +34,7 @@ export class CompatHistoryProjector { private projectSessionBase( history: CanonicalConversationMeta ): Omit { - const { conversation, prompt, tokenCost } = history; + const { conversation, prompt } = history; return { userId: conversation.userId, sessionId: conversation.id, @@ -45,10 +44,7 @@ export class CompatHistoryProjector { pinned: conversation.pinned, title: conversation.title, action: prompt.action || null, - model: prompt.model, - optionalModels: prompt.optionalModels || [], promptName: prompt.name, - tokens: tokenCost, createdAt: conversation.createdAt, updatedAt: conversation.updatedAt, }; @@ -84,7 +80,7 @@ export class CompatHistoryProjector { .concat(messages) .filter( message => - message.role !== AiPromptRole.user || + message.role !== AiSessionMessageRole.user || !!message.content.trim() || !!message.attachments?.length ) diff --git a/packages/backend/server/src/plugins/copilot/compat/history-prompt-preload-projector.ts b/packages/backend/server/src/plugins/copilot/compat/history-prompt-preload-projector.ts index 78147da0ff..839f9eedef 100644 --- a/packages/backend/server/src/plugins/copilot/compat/history-prompt-preload-projector.ts +++ b/packages/backend/server/src/plugins/copilot/compat/history-prompt-preload-projector.ts @@ -1,5 +1,5 @@ import { Injectable } from '@nestjs/common'; -import { AiPromptRole } from '@prisma/client'; +import { AiSessionMessageRole } from '@prisma/client'; import { PromptService } from '../prompt/service'; import type { ChatMessage } from '../types'; @@ -24,7 +24,9 @@ export class HistoryPromptPreloadProjector { history.turns[0] ? history.turns[0].metadata : {}, history.conversation.id ) - .filter(({ role }) => role !== AiPromptRole.system) as ChatMessage[]; + .filter( + ({ role }) => role !== AiSessionMessageRole.system + ) as ChatMessage[]; preload.forEach((message, index) => { message.createdAt = new Date( diff --git a/packages/backend/server/src/plugins/copilot/config.ts b/packages/backend/server/src/plugins/copilot/config.ts index 9bcce05aa8..91f678c738 100644 --- a/packages/backend/server/src/plugins/copilot/config.ts +++ b/packages/backend/server/src/plugins/copilot/config.ts @@ -5,32 +5,9 @@ import { StorageJSONSchema, StorageProviderConfig, } from '../../base'; -import { - AnthropicOfficialConfig, - AnthropicVertexConfig, -} from './providers/anthropic'; -import { CloudflareWorkersAIConfig } from './providers/cloudflare'; -import type { FalConfig } from './providers/fal'; -import { GeminiGenerativeConfig, GeminiVertexConfig } from './providers/gemini'; -import { OpenAIConfig } from './providers/openai'; -import { - CopilotProviderType, - ModelOutputType, - VertexSchema, -} from './providers/types'; +import { CopilotProviderType } from './providers/types'; -export type CopilotProviderConfigMap = { - [CopilotProviderType.OpenAI]: OpenAIConfig; - [CopilotProviderType.CloudflareWorkersAi]: CloudflareWorkersAIConfig; - [CopilotProviderType.FAL]: FalConfig; - [CopilotProviderType.Gemini]: GeminiGenerativeConfig; - [CopilotProviderType.GeminiVertex]: GeminiVertexConfig; - [CopilotProviderType.Anthropic]: AnthropicOfficialConfig; - [CopilotProviderType.AnthropicVertex]: AnthropicVertexConfig; -}; - -export type ProviderSpecificConfig = - CopilotProviderConfigMap[keyof CopilotProviderConfigMap]; +export type ProviderSpecificConfig = Record; export const RustRequestMiddlewareValues = [ 'normalize_messages', @@ -65,24 +42,13 @@ type CopilotProviderProfileCommon = { displayName?: string; priority?: number; enabled?: boolean; - models?: string[]; + models: string[]; middleware?: ProviderMiddlewareConfig; }; -type CopilotProviderProfileVariant = { - type: T; - config: CopilotProviderConfigMap[T]; -}; - -export type CopilotProviderProfile = CopilotProviderProfileCommon & - { - [Type in CopilotProviderType]: CopilotProviderProfileVariant; - }[CopilotProviderType]; - -export type CopilotProviderDefaults = Partial< - Record, string> -> & { - fallback?: string; +export type CopilotProviderProfile = CopilotProviderProfileCommon & { + type: CopilotProviderType; + config: ProviderSpecificConfig; }; const CopilotProviderProfileBaseShape = z.object({ @@ -90,7 +56,7 @@ const CopilotProviderProfileBaseShape = z.object({ displayName: z.string().optional(), priority: z.number().optional(), enabled: z.boolean().optional(), - models: z.array(z.string()).optional(), + models: z.array(z.string().min(1)).min(1), middleware: z .object({ rust: z @@ -106,79 +72,9 @@ const CopilotProviderProfileBaseShape = z.object({ .optional(), }); -const OpenAIConfigShape = z.object({ - apiKey: z.string(), - baseURL: z.string().optional(), - oldApiStyle: z.boolean().optional(), -}); - -const FalConfigShape = z.object({ - apiKey: z.string(), -}); - -const CloudflareWorkersAIConfigShape = z.object({ - apiToken: z.string(), - accountId: z.string().optional(), - baseURL: z.string().optional(), -}); - -const GeminiGenerativeConfigShape = z.object({ - apiKey: z.string(), - baseURL: z.string().optional(), -}); - -const VertexProviderConfigShape = z.object({ - location: z.string().optional(), - project: z.string().optional(), - baseURL: z.string().optional(), - googleAuthOptions: z.any().optional(), - fetch: z.any().optional(), -}); - -const AnthropicOfficialConfigShape = z.object({ - apiKey: z.string(), - baseURL: z.string().optional(), -}); - -const CopilotProviderProfileShape = z.discriminatedUnion('type', [ - CopilotProviderProfileBaseShape.extend({ - type: z.literal(CopilotProviderType.OpenAI), - config: OpenAIConfigShape, - }), - CopilotProviderProfileBaseShape.extend({ - type: z.literal(CopilotProviderType.FAL), - config: FalConfigShape, - }), - CopilotProviderProfileBaseShape.extend({ - type: z.literal(CopilotProviderType.CloudflareWorkersAi), - config: CloudflareWorkersAIConfigShape, - }), - CopilotProviderProfileBaseShape.extend({ - type: z.literal(CopilotProviderType.Gemini), - config: GeminiGenerativeConfigShape, - }), - CopilotProviderProfileBaseShape.extend({ - type: z.literal(CopilotProviderType.GeminiVertex), - config: VertexProviderConfigShape, - }), - CopilotProviderProfileBaseShape.extend({ - type: z.literal(CopilotProviderType.Anthropic), - config: AnthropicOfficialConfigShape, - }), - CopilotProviderProfileBaseShape.extend({ - type: z.literal(CopilotProviderType.AnthropicVertex), - config: VertexProviderConfigShape, - }), -]); - -const CopilotProviderDefaultsShape = z.object({ - [ModelOutputType.Text]: z.string().optional(), - [ModelOutputType.Object]: z.string().optional(), - [ModelOutputType.Embedding]: z.string().optional(), - [ModelOutputType.Image]: z.string().optional(), - [ModelOutputType.Rerank]: z.string().optional(), - [ModelOutputType.Structured]: z.string().optional(), - fallback: z.string().optional(), +const CopilotProviderProfileShape = CopilotProviderProfileBaseShape.extend({ + type: z.nativeEnum(CopilotProviderType), + config: z.record(z.string(), z.unknown()), }); declare global { @@ -202,14 +98,6 @@ declare global { storage: ConfigItem; providers: { profiles: ConfigItem; - defaults: ConfigItem; - openai: ConfigItem; - cloudflareWorkersAi: ConfigItem; - fal: ConfigItem; - gemini: ConfigItem; - geminiVertex: ConfigItem; - anthropic: ConfigItem; - anthropicVertex: ConfigItem; }; }; } @@ -245,56 +133,6 @@ defineModuleConfig('copilot', { default: [], shape: z.array(CopilotProviderProfileShape), }, - 'providers.defaults': { - desc: 'The default provider ids for model output types and global fallback.', - default: {}, - shape: CopilotProviderDefaultsShape, - }, - 'providers.openai': { - desc: 'The config for the openai provider.', - default: { - apiKey: '', - baseURL: 'https://api.openai.com/v1', - }, - link: 'https://github.com/openai/openai-node', - }, - 'providers.cloudflareWorkersAi': { - desc: 'The config for the Cloudflare Workers AI provider.', - default: { - apiToken: '', - accountId: '', - }, - }, - 'providers.fal': { - desc: 'The config for the fal provider.', - default: { - apiKey: '', - }, - }, - 'providers.gemini': { - desc: 'The config for the gemini provider.', - default: { - apiKey: '', - baseURL: 'https://generativelanguage.googleapis.com/v1beta', - }, - }, - 'providers.geminiVertex': { - desc: 'The config for the gemini provider in Google Vertex AI.', - default: {}, - schema: VertexSchema, - }, - 'providers.anthropic': { - desc: 'The config for the anthropic provider.', - default: { - apiKey: '', - baseURL: 'https://api.anthropic.com/v1', - }, - }, - 'providers.anthropicVertex': { - desc: 'The config for the anthropic provider in Google Vertex AI.', - default: {}, - schema: VertexSchema, - }, unsplash: { desc: 'The config for the unsplash key.', default: { diff --git a/packages/backend/server/src/plugins/copilot/context/realtime.ts b/packages/backend/server/src/plugins/copilot/context/realtime.ts index 6e9918e304..ab13b78f2c 100644 --- a/packages/backend/server/src/plugins/copilot/context/realtime.ts +++ b/packages/backend/server/src/plugins/copilot/context/realtime.ts @@ -1,7 +1,8 @@ import { Injectable, OnModuleInit } from '@nestjs/common'; import { z } from 'zod'; -import { OnEvent } from '../../../base'; +import { Config } from '../../../base/config'; +import { OnEvent } from '../../../base/event'; import { PermissionAccess } from '../../../core/permission'; import { RealtimePublisher, @@ -10,6 +11,7 @@ import { registerRealtimeLiveQuery, } from '../../../core/realtime'; import { Models } from '../../../models'; +import { assertCopilotEnabled } from '../availability'; export function workspaceEmbeddingRoom(workspaceId: string) { return realtimeWorkspaceEmbeddingProgressRoom(workspaceId); @@ -21,7 +23,8 @@ export class CopilotEmbeddingRealtimeProvider implements OnModuleInit { private readonly ac: PermissionAccess, private readonly models: Models, private readonly registry: RealtimeRegistry, - private readonly publisher: RealtimePublisher + private readonly publisher: RealtimePublisher, + private readonly config: Config ) {} onModuleInit() { @@ -118,6 +121,7 @@ export class CopilotEmbeddingRealtimeProvider implements OnModuleInit { } private async assertCopilot(userId: string, workspaceId: string) { + assertCopilotEnabled(this.config); await this.ac .user(userId) .workspace(workspaceId) diff --git a/packages/backend/server/src/plugins/copilot/context/resolver.ts b/packages/backend/server/src/plugins/copilot/context/resolver.ts index bef1b67c36..85fc7353f6 100644 --- a/packages/backend/server/src/plugins/copilot/context/resolver.ts +++ b/packages/backend/server/src/plugins/copilot/context/resolver.ts @@ -50,6 +50,7 @@ import { Models, } from '../../../models'; import { CopilotEmbeddingJob } from '../embedding/job'; +import { CopilotEnabled } from '../feature'; import { COPILOT_LOCKER, CopilotType } from '../resolver'; import { ChatSessionService } from '../session'; import { CopilotStorage } from '../storage'; @@ -286,6 +287,7 @@ class ContextMatchedDocChunk implements DocChunkSimilarity { } @Throttle() +@CopilotEnabled() @Resolver(() => CopilotType) export class CopilotContextRootResolver { constructor( @@ -435,6 +437,7 @@ export class CopilotContextRootResolver { } @Throttle() +@CopilotEnabled() @Resolver(() => CopilotContextType) export class CopilotContextResolver { constructor( diff --git a/packages/backend/server/src/plugins/copilot/context/service.ts b/packages/backend/server/src/plugins/copilot/context/service.ts index ee9d03af52..b74cdc3c52 100644 --- a/packages/backend/server/src/plugins/copilot/context/service.ts +++ b/packages/backend/server/src/plugins/copilot/context/service.ts @@ -1,3 +1,4 @@ +/* oxlint-disable import/no-cycle -- Context embedding reuses the shared capability runtime. */ import { Injectable, OnApplicationBootstrap } from '@nestjs/common'; import { diff --git a/packages/backend/server/src/plugins/copilot/controller.ts b/packages/backend/server/src/plugins/copilot/controller.ts index f1b61696eb..b6fd911270 100644 --- a/packages/backend/server/src/plugins/copilot/controller.ts +++ b/packages/backend/server/src/plugins/copilot/controller.ts @@ -36,6 +36,7 @@ import { UnsplashIsNotConfigured, } from '../../base'; import { CurrentUser, Public } from '../../core/auth'; +import { CopilotEnabled } from './feature'; import { ActionStreamHost, projectActionEventToChatEvent, @@ -52,6 +53,7 @@ export interface ChatEvent { const PING_INTERVAL = 5000; +@CopilotEnabled() @Controller('/api/copilot') export class CopilotController implements BeforeApplicationShutdown { private readonly logger = new Logger(CopilotController.name); diff --git a/packages/backend/server/src/plugins/copilot/conversation/store.ts b/packages/backend/server/src/plugins/copilot/conversation/store.ts index af3b92ec30..0d6a1b27f3 100644 --- a/packages/backend/server/src/plugins/copilot/conversation/store.ts +++ b/packages/backend/server/src/plugins/copilot/conversation/store.ts @@ -83,7 +83,6 @@ export class ConversationStore { conversation: Conversation; turns: Turn[]; promptName: string; - tokenCost: number; } | undefined > { @@ -96,7 +95,6 @@ export class ConversationStore { conversation: this.toConversation(session), turns: this.toTurns(session), promptName: session.promptName, - tokenCost: session.tokenCost, }; } @@ -104,7 +102,6 @@ export class ConversationStore { | { conversation: Conversation; promptName: string; - tokenCost: number; } | undefined > { @@ -124,7 +121,6 @@ export class ConversationStore { updatedAt: session.updatedAt, }, promptName: session.promptName, - tokenCost: session.tokenCost, }; } @@ -146,7 +142,6 @@ export class ConversationStore { turnFromChatMessage(message, session.id) ), promptName: session.promptName, - tokenCost: session.tokenCost, })); } @@ -168,14 +163,12 @@ export class ConversationStore { updatedAt: session.updatedAt, } satisfies Conversation, promptName: session.promptName, - tokenCost: session.tokenCost, })); } async appendTurns(input: { sessionId: string; userId: string; - prompt: { model: string }; turns: Turn[]; }) { return await this.models.copilotSession.updateMessages({ @@ -190,14 +183,12 @@ export class ConversationStore { async appendTurn(input: { sessionId: string; userId: string; - prompt: { model: string }; turn: Turn; compatSubmissionId?: string; }) { const message = await this.models.copilotSession.appendMessage({ sessionId: input.sessionId, userId: input.userId, - prompt: input.prompt, message: (() => { const { id: _id, ...message } = chatMessageFromTurn(input.turn); return { ...message, compatSubmissionId: input.compatSubmissionId }; diff --git a/packages/backend/server/src/plugins/copilot/cron.ts b/packages/backend/server/src/plugins/copilot/cron.ts index 17814227c2..411ca7d953 100644 --- a/packages/backend/server/src/plugins/copilot/cron.ts +++ b/packages/backend/server/src/plugins/copilot/cron.ts @@ -5,6 +5,7 @@ import { JOB_SIGNAL, JobQueue, OneDay, OnJob } from '../../base'; import { Models } from '../../models'; const CLEANUP_EMBEDDING_JOB_BATCH_SIZE = 100; +const BACKGROUND_COPILOT_JOB_PRIORITY = 100; declare global { interface Jobs { @@ -71,9 +72,11 @@ export class CopilotCronJobs { const sessions = await this.models.copilotSession.toBeGenerateTitle(); for (const session of sessions) { - await this.jobs.add('copilot.session.generateTitle', { - sessionId: session.id, - }); + await this.jobs.add( + 'copilot.session.generateTitle', + { sessionId: session.id }, + { priority: BACKGROUND_COPILOT_JOB_PRIORITY } + ); } this.logger.log( `Scheduled title generation for ${sessions.length} sessions` diff --git a/packages/backend/server/src/plugins/copilot/embedding/client.ts b/packages/backend/server/src/plugins/copilot/embedding/client.ts index 673b9c1d52..8df21b18e4 100644 --- a/packages/backend/server/src/plugins/copilot/embedding/client.ts +++ b/packages/backend/server/src/plugins/copilot/embedding/client.ts @@ -1,6 +1,7 @@ +/* oxlint-disable import/no-cycle -- Embedding delegates to the shared capability runtime. */ import { createHash } from 'node:crypto'; -import { Injectable, Logger } from '@nestjs/common'; +import { forwardRef, Inject, Injectable, Logger } from '@nestjs/common'; import { CopilotFailedToGenerateEmbedding } from '../../../base/error/errors.gen'; import { @@ -10,7 +11,6 @@ import { } from '../../../models'; import { type CopilotRerankRequest } from '../providers/types'; import { CapabilityRuntime } from '../runtime/capability-runtime'; -import { TaskPolicy } from '../runtime/task-policy'; import { type EmbeddingCallOptionsInput, EmbeddingClient, @@ -18,20 +18,20 @@ import { type ReRankResult, } from './types'; +type EmbeddingRuntime = Pick< + CapabilityRuntime, + 'embeddingConfigured' | 'embed' | 'rerank' +>; + class ProductionEmbeddingClient extends EmbeddingClient { private readonly logger = new Logger(ProductionEmbeddingClient.name); - constructor( - private readonly taskPolicy: TaskPolicy, - private readonly runtime: CapabilityRuntime - ) { + constructor(private readonly runtime: EmbeddingRuntime) { super(); } override async configured(): Promise { - const result = await this.runtime.embeddingConfigured( - this.taskPolicy.resolveEmbeddingModelId() - ); + const result = await this.runtime.embeddingConfigured('route-selected'); if (!result) { this.logger.warn( 'Copilot embedding client is not configured properly, please check your configuration.' @@ -45,7 +45,7 @@ class ProductionEmbeddingClient extends EmbeddingClient { options?: EmbeddingCallOptionsInput ): Promise { const normalizedOptions = normalizeEmbeddingCallOptions(options); - const modelId = this.taskPolicy.resolveEmbeddingModelId(); + const modelId = 'route-selected'; const embeddings = await this.runtime.embed(modelId, input, { dimensions: EMBEDDING_DIMENSIONS, signal: normalizedOptions.signal, @@ -94,17 +94,13 @@ class ProductionEmbeddingClient extends EmbeddingClient { })), }; - const ranks = await this.runtime.rerank( - this.taskPolicy.resolveRerankModelId(), - rerankRequest, - { - signal: normalizedOptions.signal, - user: normalizedOptions.userId, - workspace: normalizedOptions.workspaceId, - byokLeaseId: normalizedOptions.byokLeaseId, - featureKind: 'rerank', - } - ); + const ranks = await this.runtime.rerank('route-selected', rerankRequest, { + signal: normalizedOptions.signal, + user: normalizedOptions.userId, + workspace: normalizedOptions.workspaceId, + byokLeaseId: normalizedOptions.byokLeaseId, + featureKind: 'rerank', + }); try { return ranks.map((score, i) => { @@ -206,12 +202,12 @@ export class CopilotEmbeddingClientService { private client: EmbeddingClient | undefined; constructor( - private readonly taskPolicy: TaskPolicy, - private readonly runtime: CapabilityRuntime + @Inject(forwardRef(() => CapabilityRuntime)) + private readonly runtime: EmbeddingRuntime ) {} async refresh() { - const client = new ProductionEmbeddingClient(this.taskPolicy, this.runtime); + const client = new ProductionEmbeddingClient(this.runtime); await client.configured(); this.client = client; return this.client; diff --git a/packages/backend/server/src/plugins/copilot/feature.ts b/packages/backend/server/src/plugins/copilot/feature.ts new file mode 100644 index 0000000000..48a3413c00 --- /dev/null +++ b/packages/backend/server/src/plugins/copilot/feature.ts @@ -0,0 +1,54 @@ +import { CanActivate, Injectable, UseGuards } from '@nestjs/common'; + +import { Config } from '../../base/config'; +import { OnEvent } from '../../base/event'; +import { ServerFeature, ServerService } from '../../core/config'; +import { assertCopilotEnabled } from './availability'; + +@Injectable() +export class CopilotFeatureService { + constructor( + private readonly config: Config, + private readonly server: ServerService + ) {} + + get enabled() { + return this.config.copilot.enabled; + } + + @OnEvent('config.init') + onConfigInit() { + this.syncServerFeature(); + } + + @OnEvent('config.changed') + onConfigChanged(event: Events['config.changed']) { + if ('copilot' in event.updates) { + this.syncServerFeature(); + } + } + + assertEnabled() { + assertCopilotEnabled(this.config); + } + + private syncServerFeature() { + if (this.enabled) { + this.server.enableFeature(ServerFeature.Copilot); + } else { + this.server.disableFeature(ServerFeature.Copilot); + } + } +} + +@Injectable() +export class CopilotFeatureGuard implements CanActivate { + constructor(private readonly feature: CopilotFeatureService) {} + + canActivate() { + this.feature.assertEnabled(); + return true; + } +} + +export const CopilotEnabled = () => UseGuards(CopilotFeatureGuard); diff --git a/packages/backend/server/src/plugins/copilot/index.ts b/packages/backend/server/src/plugins/copilot/index.ts index 25398e2704..5897dcb90c 100644 --- a/packages/backend/server/src/plugins/copilot/index.ts +++ b/packages/backend/server/src/plugins/copilot/index.ts @@ -11,6 +11,7 @@ import { StorageModule } from '../../core/storage'; import { WorkspaceModule } from '../../core/workspaces'; import { IndexerModule } from '../indexer'; import { CopilotController } from './controller'; +import { CopilotFeatureGuard, CopilotFeatureService } from './feature'; import { WorkspaceMcpController } from './mcp/controller'; import { McpCredentialService } from './mcp/credential'; import { McpCredentialResolver } from './mcp/resolver'; @@ -34,20 +35,27 @@ const COPILOT_SHARED_IMPORTS = [ ]; @Module({ - imports: [...COPILOT_SHARED_IMPORTS], + imports: [ServerConfigModule], + providers: [CopilotFeatureService, CopilotFeatureGuard], + exports: [CopilotFeatureService, CopilotFeatureGuard], +}) +export class CopilotAvailabilityModule {} + +@Module({ + imports: [...COPILOT_SHARED_IMPORTS, CopilotAvailabilityModule], providers: [...COPILOT_KERNEL_PROVIDERS], - exports: [...COPILOT_KERNEL_PROVIDERS], + exports: [CopilotAvailabilityModule, ...COPILOT_KERNEL_PROVIDERS], }) export class CopilotKernelModule {} @Module({ - imports: [PermissionModule], + imports: [PermissionModule, CopilotAvailabilityModule], providers: [...COPILOT_TRANSCRIPT_REALTIME_PROVIDERS], }) export class CopilotRealtimeModule {} @Module({ - imports: [PermissionModule], + imports: [PermissionModule, CopilotAvailabilityModule], providers: [...COPILOT_CONTEXT_REALTIME_PROVIDERS], }) export class CopilotEmbeddingRealtimeModule {} diff --git a/packages/backend/server/src/plugins/copilot/mcp/controller.ts b/packages/backend/server/src/plugins/copilot/mcp/controller.ts index 3d2b4e0ff4..a82974fa8d 100644 --- a/packages/backend/server/src/plugins/copilot/mcp/controller.ts +++ b/packages/backend/server/src/plugins/copilot/mcp/controller.ts @@ -16,6 +16,7 @@ import type { Request, Response } from 'express'; import { ActionForbidden, Throttle } from '../../../base'; import { Public } from '../../../core/auth'; import { extractTokenFromHeader } from '../../../core/auth/input'; +import { CopilotEnabled } from '../feature'; import { McpCredentialService } from './credential'; import { WorkspaceMcpProvider, type WorkspaceMcpServer } from './provider'; @@ -46,6 +47,7 @@ const SUPPORTED_PROTOCOL_VERSIONS = new Set([ '2024-10-07', ]); +@CopilotEnabled() @Controller('/api/workspaces/:workspaceId/mcp') export class WorkspaceMcpController { private readonly logger = new Logger(WorkspaceMcpController.name); diff --git a/packages/backend/server/src/plugins/copilot/mcp/resolver.ts b/packages/backend/server/src/plugins/copilot/mcp/resolver.ts index 37d6263a64..0b2565db60 100644 --- a/packages/backend/server/src/plugins/copilot/mcp/resolver.ts +++ b/packages/backend/server/src/plugins/copilot/mcp/resolver.ts @@ -15,6 +15,7 @@ import { McpAccessMode } from '@prisma/client'; import { CurrentUser } from '../../../core/auth'; import { PermissionAccess } from '../../../core/permission'; +import { CopilotEnabled } from '../feature'; import { McpCredentialService } from './credential'; registerEnumType(McpAccessMode, { name: 'McpAccessMode' }); @@ -89,6 +90,7 @@ class CreateMcpCredentialInput { expirationDays!: number; } +@CopilotEnabled() @Resolver() export class McpCredentialResolver { constructor( diff --git a/packages/backend/server/src/plugins/copilot/module-providers.ts b/packages/backend/server/src/plugins/copilot/module-providers.ts index 8f946864c7..e9f9269817 100644 --- a/packages/backend/server/src/plugins/copilot/module-providers.ts +++ b/packages/backend/server/src/plugins/copilot/module-providers.ts @@ -1,9 +1,4 @@ -import { CopilotAccessPolicy } from './access'; -import { - ByokEntitlementPolicy, - ByokService, - WorkspaceByokResolver, -} from './byok'; +import { ByokEntitlementPolicy, WorkspaceByokResolver } from './byok'; import { HistoryAttachmentUrlProjector } from './compat/history-attachment-url-projector'; import { CompatHistoryProjector } from './compat/history-projector'; import { HistoryPromptPreloadProjector } from './compat/history-prompt-preload-projector'; @@ -25,30 +20,18 @@ import { } from './embedding'; import { WorkspaceMcpProvider } from './mcp/provider'; import { PromptService } from './prompt'; -import { - CopilotProviderFactory, - CopilotProviderLifecycleService, - CopilotProviderRegistryService, - CopilotProviders, -} from './providers'; import { CopilotResolver, UserCopilotResolver } from './resolver'; import { ActionRuntimeBridge } from './runtime/action-runtime-bridge'; import { CapabilityRuntime } from './runtime/capability-runtime'; -import { CopilotExecutionMetrics } from './runtime/execution-metrics'; -import { ExecutionPlanBuilder } from './runtime/execution-plan'; +import { CopilotRuntimeEventConsumer } from './runtime/copilot-runtime-event-consumer'; import { ActionStreamHost } from './runtime/hosts/action-stream-host'; import { AttachmentAdmissionHost } from './runtime/hosts/attachment-admission'; import { AttachmentMaterializer } from './runtime/hosts/attachment-materializer'; -import { CapabilityPolicyHost } from './runtime/hosts/capability-policy-host'; import { ConversationHost } from './runtime/hosts/conversation-host'; import { ImageResultHost } from './runtime/hosts/image-result-host'; import { ResponsePostprocessor } from './runtime/hosts/response-postprocessor'; -import { ToolExecutorHost } from './runtime/hosts/tool-executor-host'; import { TurnPersistence } from './runtime/hosts/turn-persistence'; -import { ModelSelectionPolicy } from './runtime/model-selection-policy'; -import { NativeExecutionEngine } from './runtime/native-execution-engine'; import { PromptRuntime } from './runtime/prompt-runtime'; -import { TaskPolicy } from './runtime/task-policy'; import { ToolRuntime } from './runtime/tool-runtime'; import { TurnOrchestrator } from './runtime/turn-orchestrator'; import { ChatSessionService } from './session'; @@ -65,21 +48,14 @@ import { CopilotWorkspaceService, } from './workspace'; -export const COPILOT_PROVIDER_PROVIDERS = [ - ...CopilotProviders, - CopilotProviderRegistryService, - CopilotProviderFactory, - CopilotProviderLifecycleService, -]; +export const COPILOT_PROVIDER_PROVIDERS: [] = []; export const COPILOT_RUNTIME_PROVIDERS = [ ByokEntitlementPolicy, - ByokService, ChatSessionService, ConversationStore, ConversationInboxService, ConversationPolicy, - CopilotAccessPolicy, HistoryAttachmentUrlProjector, CompatHistoryProjector, HistoryPromptPreloadProjector, @@ -88,18 +64,12 @@ export const COPILOT_RUNTIME_PROVIDERS = [ CopilotContextService, CopilotEmbeddingClientService, PromptService, - ModelSelectionPolicy, ActionRuntimeBridge, - CopilotExecutionMetrics, - ExecutionPlanBuilder, + CopilotRuntimeEventConsumer, PromptRuntime, - CapabilityPolicyHost, ConversationHost, CapabilityRuntime, - NativeExecutionEngine, - TaskPolicy, ToolRuntime, - ToolExecutorHost, AttachmentMaterializer, AttachmentAdmissionHost, ActionStreamHost, diff --git a/packages/backend/server/src/plugins/copilot/prompt/native-contract.ts b/packages/backend/server/src/plugins/copilot/prompt/native-contract.ts index 6ff8482cef..fb612588cc 100644 --- a/packages/backend/server/src/plugins/copilot/prompt/native-contract.ts +++ b/packages/backend/server/src/plugins/copilot/prompt/native-contract.ts @@ -1,36 +1,17 @@ import { - llmCollectPromptMetadata, - llmCountPromptTokens, llmGetBuiltInPromptSpec, llmListBuiltInPromptSpecs, llmRenderBuiltInPrompt, llmRenderBuiltInSessionPrompt, - llmRenderPrompt, - llmRenderSessionPrompt, type NativeBuiltInPromptRenderRequest as NativeBuiltInPromptRenderContract, type NativeBuiltInPromptSessionRenderRequest as NativeBuiltInPromptSessionContract, - type NativePromptCountTokensRequest as NativePromptTokenCountContract, - type NativePromptCountTokensResponse as NativePromptTokenCountResult, - type NativePromptMetadataRequest as NativePromptMetadataContract, - type NativePromptMetadataResponse as NativePromptMetadataResult, - type NativePromptRenderRequest as NativePromptRenderContract, type NativePromptRenderResponse as NativePromptRenderResult, - type NativePromptSessionRenderRequest as NativePromptSessionContract, type NativePromptSessionRenderResponse as NativePromptSessionResult, } from '../../../native'; import type { PromptMessage, PromptParams } from '../providers/types'; import { projectPromptMessageForNative } from '../runtime/contracts'; import type { PromptSpec } from './spec'; -export type NativePromptRenderRequest = Omit< - NativePromptRenderContract, - 'messages' | 'templateParams' | 'renderParams' -> & { - messages: PromptMessage[]; - templateParams: PromptParams; - renderParams: PromptParams; -}; - export type NativePromptRenderResponse = Omit< NativePromptRenderResult, 'messages' @@ -45,46 +26,6 @@ export type NativeBuiltInPromptRenderRequest = Omit< renderParams: PromptParams; }; -export type NativePromptCountTokensRequest = Omit< - NativePromptTokenCountContract, - 'messages' | 'model' -> & { - model?: string | null; - messages: Pick[]; -}; - -export type NativePromptCountTokensResponse = NativePromptTokenCountResult; - -export type NativePromptMetadataRequest = Omit< - NativePromptMetadataContract, - 'messages' -> & { - messages: PromptMessage[]; -}; - -export type NativePromptMetadataResponse = Omit< - NativePromptMetadataResult, - 'templateParams' -> & { - templateParams: PromptParams; -}; - -export type NativePromptSessionRenderRequest = Omit< - NativePromptSessionContract, - 'prompt' | 'turns' | 'renderParams' -> & { - prompt: Omit< - NativePromptSessionContract['prompt'], - 'templateParams' | 'messages' | 'model' - > & { - model?: string | null; - templateParams: PromptParams; - messages: PromptMessage[]; - }; - turns: PromptMessage[]; - renderParams: PromptParams; -}; - export type NativePromptSessionRenderResponse = Omit< NativePromptSessionResult, 'messages' @@ -100,8 +41,7 @@ export type NativeBuiltInPromptSessionRenderRequest = Omit< renderParams: PromptParams; }; -type NativePromptContractMessage = - NativePromptRenderContract['messages'][number]; +type NativePromptContractMessage = NativePromptRenderResult['messages'][number]; function toNativePromptMessage( message: PromptMessage @@ -123,22 +63,6 @@ function fromNativePromptMessage( }; } -export function renderPromptNative( - request: NativePromptRenderRequest -): NativePromptRenderResponse { - const normalizedMessages = request.messages.map(toNativePromptMessage); - const rendered = llmRenderPrompt({ - messages: normalizedMessages, - templateParams: request.templateParams, - renderParams: request.renderParams, - }); - - return { - ...rendered, - messages: rendered.messages.map(fromNativePromptMessage), - }; -} - export function renderBuiltInPromptNative( request: NativeBuiltInPromptRenderRequest ): NativePromptRenderResponse { @@ -153,25 +77,6 @@ export function renderBuiltInPromptNative( }; } -export function renderPromptSessionNative( - request: NativePromptSessionRenderRequest -): NativePromptSessionRenderResponse { - const rendered = llmRenderSessionPrompt({ - ...request, - prompt: { - ...request.prompt, - messages: request.prompt.messages.map(toNativePromptMessage), - model: request.prompt.model ?? undefined, - }, - turns: request.turns.map(toNativePromptMessage), - renderParams: request.renderParams, - }); - return { - ...rendered, - messages: rendered.messages.map(fromNativePromptMessage), - }; -} - export function renderBuiltInPromptSessionNative( request: NativeBuiltInPromptSessionRenderRequest ): NativePromptSessionRenderResponse { @@ -187,29 +92,10 @@ export function renderBuiltInPromptSessionNative( }; } -export function countPromptTokensNative( - request: NativePromptCountTokensRequest -): NativePromptCountTokensResponse { - return llmCountPromptTokens({ - ...request, - model: request.model ?? undefined, - }); -} - -export function collectPromptMetadataNative( - request: NativePromptMetadataRequest -): NativePromptMetadataResponse { - return llmCollectPromptMetadata({ - messages: request.messages.map(toNativePromptMessage), - }); -} - export function listBuiltInPromptSpecsNative(): PromptSpec[] { return llmListBuiltInPromptSpecs().map(spec => ({ name: spec.name, action: spec.action, - model: spec.model, - optionalModels: spec.optionalModels, config: spec.config, params: spec.params ? Object.fromEntries( @@ -238,8 +124,6 @@ export function getBuiltInPromptSpecNative(name: string): PromptSpec | null { return { name: spec.name, action: spec.action, - model: spec.model, - optionalModels: spec.optionalModels, config: spec.config, params: spec.params ? Object.fromEntries( diff --git a/packages/backend/server/src/plugins/copilot/prompt/service.ts b/packages/backend/server/src/plugins/copilot/prompt/service.ts index 0e4a2e46b2..490629b93d 100644 --- a/packages/backend/server/src/plugins/copilot/prompt/service.ts +++ b/packages/backend/server/src/plugins/copilot/prompt/service.ts @@ -2,15 +2,11 @@ import { Injectable, Logger } from '@nestjs/common'; import type { PromptMessage, PromptParams } from '../providers/types'; import { - collectPromptMetadataNative, - countPromptTokensNative, getBuiltInPromptSpecNative, renderBuiltInPromptNative, renderBuiltInPromptSessionNative, - renderPromptNative, - renderPromptSessionNative, } from './native-contract'; -import type { Prompt, PromptSpec, ResolvedPrompt } from './spec'; +import type { PromptSpec, ResolvedPrompt } from './spec'; @Injectable() export class PromptService { @@ -20,11 +16,6 @@ export class PromptService { } async get(name: string): Promise { - const compatPrompt = this.lookupCompatPrompt(name); - if (compatPrompt) { - return this.describeCompatPrompt(this.clonePrompt(compatPrompt)); - } - const builtInPromptSpec = this.lookupBuiltInPromptSpec(name); if (!builtInPromptSpec) return null; @@ -36,17 +27,10 @@ export class PromptService { params: PromptParams, sessionId?: string ): PromptMessage[] { - const rendered = - prompt.source === 'built_in' - ? renderBuiltInPromptNative({ - name: prompt.name, - renderParams: params, - }) - : renderPromptNative({ - messages: this.requireCompatMessages(prompt), - templateParams: prompt.params, - renderParams: params, - }); + const rendered = renderBuiltInPromptNative({ + name: prompt.name, + renderParams: params, + }); this.logWarnings(rendered.warnings, sessionId); return rendered.messages; @@ -56,38 +40,18 @@ export class PromptService { prompt: ResolvedPrompt, turns: PromptMessage[], params: PromptParams, - maxTokenSize = prompt.config?.maxTokens || 128 * 1024, sessionId?: string ): PromptMessage[] { - const rendered = - prompt.source === 'built_in' - ? renderBuiltInPromptSessionNative({ - name: prompt.name, - turns, - renderParams: params, - maxTokenSize, - }) - : renderPromptSessionNative({ - prompt: { - action: prompt.action, - model: prompt.model, - promptTokens: this.countCompatPromptTokens(prompt), - templateParams: prompt.params, - messages: this.requireCompatMessages(prompt), - }, - turns, - renderParams: params, - maxTokenSize, - }); + const rendered = renderBuiltInPromptSessionNative({ + name: prompt.name, + turns, + renderParams: params, + }); this.logWarnings(rendered.warnings, sessionId); return rendered.messages; } - protected lookupCompatPrompt(_name: string): Prompt | null { - return null; - } - protected lookupBuiltInPromptSpec(name: string): PromptSpec | null { const spec = getBuiltInPromptSpecNative(name); return spec ? this.clonePromptSpec(spec) : null; @@ -104,23 +68,9 @@ export class PromptService { })); } - protected clonePrompt(prompt: Prompt): Prompt { - return { - ...prompt, - optionalModels: prompt.optionalModels - ? [...prompt.optionalModels] - : undefined, - config: prompt.config ? structuredClone(prompt.config) : undefined, - messages: this.cloneMessages(prompt.messages), - }; - } - protected clonePromptSpec(spec: PromptSpec): PromptSpec { return { ...spec, - optionalModels: spec.optionalModels - ? [...spec.optionalModels] - : undefined, config: spec.config ? structuredClone(spec.config) : undefined, params: spec.params ? structuredClone(spec.params) : undefined, messages: spec.messages.map(message => ({ ...message })), @@ -132,27 +82,9 @@ export class PromptService { return { name: spec.name, action: spec.action, - model: spec.model, - optionalModels: spec.optionalModels ?? [], config: spec.config ? structuredClone(spec.config) : undefined, paramKeys: Object.keys(params), params, - source: 'built_in', - }; - } - - private describeCompatPrompt(prompt: Prompt): ResolvedPrompt { - const metadata = collectPromptMetadataNative({ messages: prompt.messages }); - return { - name: prompt.name, - action: prompt.action, - model: prompt.model, - optionalModels: prompt.optionalModels ?? [], - config: prompt.config ? structuredClone(prompt.config) : undefined, - paramKeys: metadata.paramKeys, - params: metadata.templateParams, - source: 'compat', - messages: prompt.messages, }; } @@ -178,23 +110,6 @@ export class PromptService { ); } - private countCompatPromptTokens(prompt: ResolvedPrompt): number { - return countPromptTokensNative({ - model: prompt.model, - messages: this.requireCompatMessages(prompt).map(message => ({ - content: message.content, - })), - }).tokens; - } - - private requireCompatMessages(prompt: ResolvedPrompt): PromptMessage[] { - if (prompt.source === 'compat' && prompt.messages) { - return this.cloneMessages(prompt.messages); - } - - throw new Error(`Prompt ${prompt.name} does not expose compat messages`); - } - private logWarnings(warnings: string[], sessionId?: string) { if (!sessionId) { return; diff --git a/packages/backend/server/src/plugins/copilot/prompt/spec.ts b/packages/backend/server/src/plugins/copilot/prompt/spec.ts index 6da12069b2..0e003a5916 100644 --- a/packages/backend/server/src/plugins/copilot/prompt/spec.ts +++ b/packages/backend/server/src/plugins/copilot/prompt/spec.ts @@ -1,28 +1,11 @@ -import type { - PromptConfig, - PromptMessage, - PromptParams, -} from '../providers/types'; - -export type Prompt = { - name: string; - model: string; - optionalModels?: string[]; - action?: string; - messages: PromptMessage[]; - config?: PromptConfig; -}; +import type { PromptConfig, PromptParams } from '../providers/types'; export type ResolvedPrompt = { name: string; - model: string; - optionalModels: string[]; action?: string; config?: PromptConfig; paramKeys: string[]; params: PromptParams; - source: 'built_in' | 'compat'; - messages?: PromptMessage[]; }; type PromptParamSpec = { @@ -38,8 +21,6 @@ type PromptSpecMessage = { export type PromptSpec = { name: string; action?: string; - model: string; - optionalModels?: string[]; config?: PromptConfig; params?: Record; messages: PromptSpecMessage[]; diff --git a/packages/backend/server/src/plugins/copilot/providers/anthropic/anthropic.ts b/packages/backend/server/src/plugins/copilot/providers/anthropic/anthropic.ts deleted file mode 100644 index e2fde83f2b..0000000000 --- a/packages/backend/server/src/plugins/copilot/providers/anthropic/anthropic.ts +++ /dev/null @@ -1,100 +0,0 @@ -import { CopilotProviderSideError, UserFriendlyError } from '../../../../base'; -import { - type LlmBackendConfig, - llmResolveRequestIntentOptions, -} from '../../../../native'; -import { CopilotProvider } from '../provider'; -import { hasProviderModelBehaviorFlag } from '../provider-model-runtime'; -import { - type CopilotProviderExecution, - type ProviderDriverSpec, -} from '../provider-runtime-contract'; -import { CopilotProviderType } from '../types'; -import { - getGoogleAuth, - getVertexAnthropicBaseUrl, - type VertexAnthropicProviderConfig, -} from '../utils'; - -export abstract class AnthropicProvider extends CopilotProvider { - protected resolveModelBackendKind() { - return this.type === CopilotProviderType.AnthropicVertex - ? ('anthropic_vertex' as const) - : ('anthropic' as const); - } - - override getDriverSpec(): ProviderDriverSpec { - return { - createBackendConfig: execution => this.createNativeConfig(execution), - mapError: error => this.handleError(error), - chat: { - resolveRequestOptions: async context => { - const requestIntent = await llmResolveRequestIntentOptions({ - protocol: context.protocol, - backendConfig: context.backendConfig, - reasoning: { - enabled: context.options.reasoning, - supported: hasProviderModelBehaviorFlag( - context.model, - 'reasoning_budget_12000' - ), - budgetTokens: hasProviderModelBehaviorFlag( - context.model, - 'reasoning_budget_12000' - ) - ? 12000 - : undefined, - }, - }); - - return { - attachmentCapability: this.getAttachCapability( - context.model, - context.outputType - ), - reasoning: requestIntent.reasoning, - }; - }, - }, - structured: false, - embedding: false, - rerank: false, - }; - } - - private handleError(e: any) { - if (e instanceof UserFriendlyError) { - return e; - } - return new CopilotProviderSideError({ - provider: this.type, - kind: 'unexpected_response', - message: e?.message || 'Unexpected anthropic response', - }); - } - - private async createNativeConfig( - execution?: CopilotProviderExecution - ): Promise { - const config = this.getConfig(execution); - if (this.type === CopilotProviderType.AnthropicVertex) { - const vertexConfig = config as VertexAnthropicProviderConfig; - const auth = await getGoogleAuth(vertexConfig, 'anthropic'); - const { Authorization: authHeader } = auth.headers(); - const token = authHeader.replace(/^Bearer\s+/i, ''); - const baseUrl = getVertexAnthropicBaseUrl(vertexConfig) || auth.baseUrl; - return { - base_url: baseUrl || '', - auth_token: token, - headers: { Authorization: authHeader }, - }; - } - - const officialConfig = config as { apiKey: string; baseURL?: string }; - const baseUrl = officialConfig.baseURL || 'https://api.anthropic.com/v1'; - return { - base_url: baseUrl.replace(/\/v1\/?$/, ''), - auth_token: officialConfig.apiKey, - }; - } -} diff --git a/packages/backend/server/src/plugins/copilot/providers/anthropic/index.ts b/packages/backend/server/src/plugins/copilot/providers/anthropic/index.ts deleted file mode 100644 index f37327b774..0000000000 --- a/packages/backend/server/src/plugins/copilot/providers/anthropic/index.ts +++ /dev/null @@ -1,2 +0,0 @@ -export * from './official'; -export * from './vertex'; diff --git a/packages/backend/server/src/plugins/copilot/providers/anthropic/official.ts b/packages/backend/server/src/plugins/copilot/providers/anthropic/official.ts deleted file mode 100644 index 7855524c31..0000000000 --- a/packages/backend/server/src/plugins/copilot/providers/anthropic/official.ts +++ /dev/null @@ -1,16 +0,0 @@ -import type { CopilotProviderExecution } from '../provider-runtime-contract'; -import { CopilotProviderType } from '../types'; -import { AnthropicProvider } from './anthropic'; - -export type AnthropicOfficialConfig = { - apiKey: string; - baseURL?: string; -}; - -export class AnthropicOfficialProvider extends AnthropicProvider { - override readonly type = CopilotProviderType.Anthropic; - - override configured(execution?: CopilotProviderExecution): boolean { - return !!this.getConfig(execution).apiKey; - } -} diff --git a/packages/backend/server/src/plugins/copilot/providers/anthropic/vertex.ts b/packages/backend/server/src/plugins/copilot/providers/anthropic/vertex.ts deleted file mode 100644 index b866b32b04..0000000000 --- a/packages/backend/server/src/plugins/copilot/providers/anthropic/vertex.ts +++ /dev/null @@ -1,16 +0,0 @@ -import type { CopilotProviderExecution } from '../provider-runtime-contract'; -import { CopilotProviderType } from '../types'; -import { getVertexAnthropicBaseUrl, type VertexProviderConfig } from '../utils'; -import { AnthropicProvider } from './anthropic'; - -export type AnthropicVertexConfig = VertexProviderConfig; - -export class AnthropicVertexProvider extends AnthropicProvider { - override readonly type = CopilotProviderType.AnthropicVertex; - - override configured(execution?: CopilotProviderExecution): boolean { - const config = this.getConfig(execution); - if (!config.location || !config.googleAuthOptions) return false; - return !!config.project || !!getVertexAnthropicBaseUrl(config); - } -} diff --git a/packages/backend/server/src/plugins/copilot/providers/attachments.ts b/packages/backend/server/src/plugins/copilot/providers/attachments.ts index 2e4026ef64..f8fee3e931 100644 --- a/packages/backend/server/src/plugins/copilot/providers/attachments.ts +++ b/packages/backend/server/src/plugins/copilot/providers/attachments.ts @@ -1,42 +1,4 @@ -import type { - ModelAttachmentCapability, - PromptAttachment, - PromptMessage, -} from './types'; - -export const IMAGE_ATTACHMENT_CAPABILITY: ModelAttachmentCapability = { - kinds: ['image'], - sourceKinds: ['url', 'data'], - allowRemoteUrls: true, -}; - -export const GEMINI_ATTACHMENT_CAPABILITY: ModelAttachmentCapability = { - kinds: ['image', 'audio', 'file'], - sourceKinds: ['url', 'data', 'bytes', 'file_handle'], - allowRemoteUrls: true, -}; - -export function promptAttachmentHasSource( - attachment: PromptAttachment -): boolean { - if (typeof attachment === 'string') { - return !!attachment.trim(); - } - - if ('attachment' in attachment) { - return !!attachment.attachment; - } - - switch (attachment.kind) { - case 'url': - return !!attachment.url; - case 'data': - case 'bytes': - return !!attachment.data; - case 'file_handle': - return !!attachment.fileHandle; - } -} +import type { PromptAttachment, PromptMessage } from './types'; export function applyPromptAttachmentMimeTypeHintForNative( attachment: PromptAttachment, diff --git a/packages/backend/server/src/plugins/copilot/providers/cloudflare.ts b/packages/backend/server/src/plugins/copilot/providers/cloudflare.ts deleted file mode 100644 index e0f8593a2e..0000000000 --- a/packages/backend/server/src/plugins/copilot/providers/cloudflare.ts +++ /dev/null @@ -1,65 +0,0 @@ -import { CopilotProviderSideError, UserFriendlyError } from '../../../base'; -import { type LlmBackendConfig } from '../../../native'; -import { CopilotProvider } from './provider'; -import { - type CopilotProviderExecution, - type ProviderDriverSpec, -} from './provider-runtime-contract'; -import { CopilotProviderType } from './types'; - -export type CloudflareWorkersAIConfig = { - apiToken: string; - accountId?: string; - baseURL?: string; -}; - -export class CloudflareWorkersAIProvider extends CopilotProvider { - override readonly type = CopilotProviderType.CloudflareWorkersAi; - - protected resolveModelBackendKind() { - return 'cloudflare_workers_ai' as const; - } - - override configured(execution?: CopilotProviderExecution): boolean { - const config = this.getConfig(execution); - return !!config.apiToken && (!!config.accountId || !!config.baseURL); - } - private handleError(e: any) { - if (e instanceof UserFriendlyError) { - return e; - } - return new CopilotProviderSideError({ - provider: this.type, - kind: 'unexpected_response', - message: e?.message || 'Unexpected cloudflare workers ai response', - }); - } - - private createNativeConfig( - execution?: CopilotProviderExecution - ): LlmBackendConfig { - const config = this.getConfig(execution); - return { - base_url: this.resolveBaseUrl(execution), - auth_token: config.apiToken, - }; - } - - private resolveBaseUrl(execution?: CopilotProviderExecution) { - const config = this.getConfig(execution); - if (config.baseURL) { - return config.baseURL.replace(/\/v1\/?$/, '').replace(/\/$/, ''); - } - const accountId = config.accountId ?? ''; - return `https://api.cloudflare.com/client/v4/accounts/${accountId}/ai`; - } - - override getDriverSpec(): ProviderDriverSpec { - return { - createBackendConfig: execution => this.createNativeConfig(execution), - mapError: error => this.handleError(error), - structured: false, - embedding: false, - }; - } -} diff --git a/packages/backend/server/src/plugins/copilot/providers/factory.ts b/packages/backend/server/src/plugins/copilot/providers/factory.ts deleted file mode 100644 index 4e69a63615..0000000000 --- a/packages/backend/server/src/plugins/copilot/providers/factory.ts +++ /dev/null @@ -1,527 +0,0 @@ -import { Injectable, Logger } from '@nestjs/common'; - -import { CopilotQuotaExceeded } from '../../../base'; -import { ServerFeature, ServerService } from '../../../core'; -import { type CopilotAccessContext, CopilotAccessPolicy } from '../access'; -import type { RequiredStructuredOutputContract } from '../runtime/contracts'; -import { getProviderRuntimeHost } from '../runtime/provider-runtime-context'; -import type { CopilotProvider } from './provider'; -import { - buildProviderRegistry, - type CopilotProviderRegistry, - type NormalizedCopilotProviderProfile, - resolveModel, - stripProviderPrefix, -} from './provider-registry'; -import type { - CopilotProviderExecution, - PreparedNativeEmbeddingExecution, - PreparedNativeExecution, - PreparedNativeImageExecution, - PreparedNativeRerankExecution, - PreparedNativeStructuredExecution, -} from './provider-runtime-contract'; -import { CopilotProviderRegistryService } from './registry-service'; -import { - type CopilotChatOptions, - type CopilotEmbeddingOptions, - type CopilotImageOptions, - CopilotProviderType, - type CopilotRerankRequest, - type CopilotStructuredOptions, - ModelFullConditions, - ModelOutputType, - type PromptMessage, -} from './types'; - -export type ResolvedCopilotProvider = { - providerId: string; - provider: CopilotProvider; - execution: CopilotProviderExecution; - profile: NormalizedCopilotProviderProfile; - rawModelId?: string; - modelId?: string; - explicitProviderId?: string; - prepared?: PreparedNativeExecution; - preparedStructured?: PreparedNativeStructuredExecution; - preparedEmbedding?: PreparedNativeEmbeddingExecution; - preparedRerank?: PreparedNativeRerankExecution; - preparedImage?: PreparedNativeImageExecution; -}; - -type RoutePreparationResult = Partial< - Pick< - ResolvedCopilotProvider, - | 'prepared' - | 'preparedStructured' - | 'preparedEmbedding' - | 'preparedRerank' - | 'preparedImage' - | 'modelId' - > ->; - -type EffectiveProviderRegistry = { - byokRegistry: CopilotProviderRegistry; - quotaBackedRegistry: CopilotProviderRegistry; - quotaBackedRoutesAvailable: boolean; -}; - -@Injectable() -export class CopilotProviderFactory { - constructor( - private readonly server: ServerService, - private readonly registries: CopilotProviderRegistryService, - private readonly access: CopilotAccessPolicy - ) {} - - private readonly logger = new Logger(CopilotProviderFactory.name); - - readonly #providers = new Map(); - readonly #providerIdsByType = new Map>(); - - private getRegistry() { - return this.registries.getRegistry(); - } - - private getProviderByProfile( - providerId: string, - profile: NormalizedCopilotProviderProfile - ) { - return ( - this.#providers.get(providerId) ?? - Array.from(this.#providerIdsByType.get(profile.type) ?? []) - .map(id => this.#providers.get(id)) - .find((provider): provider is CopilotProvider => !!provider) - ); - } - - private providerAvailable( - providerId: string, - profile: NormalizedCopilotProviderProfile - ) { - return !!this.getProviderByProfile(providerId, profile); - } - - private getAvailableProviderIds(registry: CopilotProviderRegistry) { - return Array.from(registry.profiles.entries()) - .filter(([providerId, profile]) => - this.providerAvailable(providerId, profile) - ) - .map(([providerId]) => providerId); - } - - private getPreferredProviderIds( - registry: CopilotProviderRegistry, - type?: CopilotProviderType - ) { - if (!type) return undefined; - return registry.byType.get(type)?.filter(providerId => { - const profile = registry.profiles.get(providerId); - return profile ? this.providerAvailable(providerId, profile) : false; - }); - } - - private normalizeCond( - registry: CopilotProviderRegistry, - providerId: string, - cond: ModelFullConditions - ): ModelFullConditions { - const modelId = stripProviderPrefix(registry, providerId, cond.modelId); - return { ...cond, modelId }; - } - - private async getEffectiveRegistry( - context: CopilotAccessContext = {} - ): Promise { - const quotaBackedRegistry = this.getRegistry(); - const routeAccess = await this.access.resolveRouteAccess(context); - - return { - byokRegistry: buildProviderRegistry({ - profiles: routeAccess.byokProfiles, - defaults: {}, - }), - quotaBackedRegistry, - quotaBackedRoutesAvailable: routeAccess.quotaBackedRoutesAvailable, - }; - } - - private getRequestContext( - options?: - | CopilotChatOptions - | CopilotStructuredOptions - | CopilotImageOptions - ): CopilotAccessContext { - return { - userId: options?.user, - workspaceId: options?.workspace, - byokLeaseId: options?.byokLeaseId, - featureKind: options?.featureKind, - quotaBackedRoutesAllowed: options?.quotaBackedRoutesAllowed, - }; - } - - private filterPreparedRoutes(routes: Array) { - return routes.filter( - (route): route is ResolvedCopilotProvider => route !== null - ); - } - - private async prepareResolvedRoutes( - routes: ResolvedCopilotProvider[], - prepare: ( - route: ResolvedCopilotProvider - ) => Promise - ) { - const preparedRoutes = await Promise.all( - routes.map(async route => { - const prepared = await prepare(route); - return prepared ? { ...route, ...prepared } : null; - }) - ); - - return this.filterPreparedRoutes(preparedRoutes); - } - - async resolveProvider( - cond: ModelFullConditions, - filter: { - prefer?: CopilotProviderType; - } = {}, - context: CopilotAccessContext = {} - ): Promise { - return (await this.resolveRoutes(cond, filter, context))[0] ?? null; - } - - async resolveRoutes( - cond: ModelFullConditions, - filter: { - prefer?: CopilotProviderType; - } = {}, - context: CopilotAccessContext = {} - ): Promise { - this.logger.debug( - `Resolving copilot provider for output type: ${cond.outputType}` - ); - const { byokRegistry, quotaBackedRegistry, quotaBackedRoutesAvailable } = - await this.getEffectiveRegistry(context); - const byokRoutes = await this.resolveRoutesFromRegistry( - byokRegistry, - cond, - filter - ); - const resolved = byokRoutes.length - ? byokRoutes - : quotaBackedRoutesAvailable - ? await this.resolveRoutesFromRegistry( - quotaBackedRegistry, - cond, - filter - ) - : []; - for (const route of resolved) { - this.logger.debug( - `Copilot provider candidate found: ${route.provider.type} (${route.providerId})` - ); - } - - if ( - !resolved.length && - !quotaBackedRoutesAvailable && - context.quotaBackedRoutesAllowed !== false - ) { - const quotaBackedRoutes = await this.resolveRoutesFromRegistry( - quotaBackedRegistry, - cond, - filter - ); - if (quotaBackedRoutes.length) { - throw new CopilotQuotaExceeded(); - } - } - - return resolved; - } - - private async resolveRoutesFromRegistry( - registry: CopilotProviderRegistry, - cond: ModelFullConditions, - filter: { - prefer?: CopilotProviderType; - } = {} - ): Promise { - const route = resolveModel({ - registry, - modelId: cond.modelId, - outputType: cond.outputType, - availableProviderIds: this.getAvailableProviderIds(registry), - preferredProviderIds: this.getPreferredProviderIds( - registry, - filter.prefer - ), - }); - - const resolved: ResolvedCopilotProvider[] = []; - for (const providerId of route.candidateProviderIds) { - const profile = registry.profiles.get(providerId); - const provider = profile - ? this.getProviderByProfile(providerId, profile) - : undefined; - if (!provider || !profile) continue; - - const normalizedCond = this.normalizeCond(registry, providerId, cond); - if ( - normalizedCond.modelId && - profile.models?.length && - !profile.models.includes(normalizedCond.modelId) - ) { - continue; - } - - const execution = { providerId, profile }; - const matched = await provider.match(normalizedCond, execution); - if (!matched) continue; - - resolved.push({ - providerId, - provider, - execution, - profile, - rawModelId: route.rawModelId, - modelId: normalizedCond.modelId, - explicitProviderId: route.explicitProviderId, - }); - } - - return resolved; - } - - async prepareRoutes( - kind: 'text' | 'streamText' | 'streamObject', - cond: ModelFullConditions, - messages: PromptMessage[], - options: CopilotChatOptions = {}, - filter: { - prefer?: CopilotProviderType; - } = {} - ): Promise { - const routes = await this.resolveRoutes( - cond, - filter, - this.getRequestContext(options) - ); - return await this.prepareResolvedRoutes(routes, async route => { - const prepared = await getProviderRuntimeHost( - route.provider - ).prepare.chat( - kind, - { ...cond, modelId: route.modelId }, - messages, - options, - route.execution - ); - const normalizedPrepared = prepared?.route ? prepared : undefined; - if (!normalizedPrepared) { - return null; - } - - return { - modelId: normalizedPrepared.route.model, - prepared: normalizedPrepared, - }; - }); - } - - async prepareStructuredRoutes( - cond: ModelFullConditions, - messages: PromptMessage[], - options: CopilotStructuredOptions = {}, - filter: { - prefer?: CopilotProviderType; - } = {}, - responseContract?: RequiredStructuredOutputContract - ): Promise { - const routes = await this.resolveRoutes( - cond, - filter, - this.getRequestContext(options) - ); - return await this.prepareResolvedRoutes(routes, async route => { - const preparedStructured = - (await getProviderRuntimeHost(route.provider).prepare.structured( - { ...cond, modelId: route.modelId }, - messages, - options, - responseContract, - route.execution - )) ?? undefined; - if (!preparedStructured) { - return null; - } - - return { - modelId: preparedStructured.route.model, - preparedStructured, - }; - }); - } - - async prepareEmbeddingRoutes( - modelId: string, - input: string | string[], - options: CopilotEmbeddingOptions = {} - ): Promise { - const routes = await this.resolveRoutes( - { modelId, outputType: ModelOutputType.Embedding }, - {}, - { - ...this.getRequestContext(options), - featureKind: options?.featureKind ?? 'embedding', - } - ); - return await this.prepareResolvedRoutes(routes, async route => { - const preparedEmbedding = - (await getProviderRuntimeHost(route.provider).prepare.embedding( - { modelId: route.modelId }, - input, - options, - route.execution - )) ?? undefined; - if (!preparedEmbedding) { - return null; - } - - return { - modelId: preparedEmbedding.route.model, - preparedEmbedding, - }; - }); - } - - async prepareRerankRoutes( - modelId: string, - request: CopilotRerankRequest, - options: CopilotChatOptions = {} - ): Promise { - const routes = await this.resolveRoutes( - { - modelId, - outputType: ModelOutputType.Rerank, - }, - {}, - { ...this.getRequestContext(options), featureKind: 'rerank' } - ); - return await this.prepareResolvedRoutes(routes, async route => { - const preparedRerank = - (await getProviderRuntimeHost(route.provider).prepare.rerank( - { modelId: route.modelId }, - request, - options, - route.execution - )) ?? undefined; - if (!preparedRerank) { - return null; - } - - return { - modelId: preparedRerank.route.model, - preparedRerank, - }; - }); - } - - async prepareImageRoutes( - cond: ModelFullConditions, - messages: PromptMessage[], - options: CopilotImageOptions = {}, - filter: { - prefer?: CopilotProviderType; - } = {} - ): Promise { - const routes = await this.resolveRoutes(cond, filter, { - ...this.getRequestContext(options), - featureKind: options?.featureKind ?? 'image', - }); - return await this.prepareResolvedRoutes(routes, async route => { - const preparedImage = - (await getProviderRuntimeHost(route.provider).prepare.image( - { ...cond, modelId: route.modelId }, - messages, - options, - route.execution - )) ?? undefined; - if (!preparedImage) { - return null; - } - - return { - modelId: preparedImage.route.model, - preparedImage, - }; - }); - } - - async getProvider( - cond: ModelFullConditions, - filter: { - prefer?: CopilotProviderType; - } = {} - ): Promise { - return (await this.resolveProvider(cond, filter))?.provider ?? null; - } - - async getProviderByModel( - modelId: string, - filter: { - prefer?: CopilotProviderType; - } = {} - ): Promise { - this.logger.debug(`Resolving copilot provider for model: ${modelId}`); - return this.getProvider({ modelId }, filter); - } - - register(providerId: string, provider: CopilotProvider) { - const existed = this.#providers.get(providerId); - if (existed?.type && existed.type !== provider.type) { - const ids = this.#providerIdsByType.get(existed.type); - ids?.delete(providerId); - if (!ids?.size) { - this.#providerIdsByType.delete(existed.type); - } - } - - this.#providers.set(providerId, provider); - - const ids = this.#providerIdsByType.get(provider.type) ?? new Set(); - ids.add(providerId); - this.#providerIdsByType.set(provider.type, ids); - - this.logger.log( - `Copilot provider [${provider.type}] registered as [${providerId}].` - ); - this.server.enableFeature(ServerFeature.Copilot); - } - - unregister(providerId: string, provider: CopilotProvider) { - const existed = this.#providers.get(providerId); - if (!existed || existed !== provider) { - return; - } - - this.#providers.delete(providerId); - - const ids = this.#providerIdsByType.get(provider.type); - ids?.delete(providerId); - if (!ids?.size) { - this.#providerIdsByType.delete(provider.type); - } - - this.logger.log( - `Copilot provider [${provider.type}] unregistered from [${providerId}].` - ); - if (this.#providers.size === 0) { - this.server.disableFeature(ServerFeature.Copilot); - } - } -} diff --git a/packages/backend/server/src/plugins/copilot/providers/fal.ts b/packages/backend/server/src/plugins/copilot/providers/fal.ts deleted file mode 100644 index ef49f610fe..0000000000 --- a/packages/backend/server/src/plugins/copilot/providers/fal.ts +++ /dev/null @@ -1,59 +0,0 @@ -import { Injectable } from '@nestjs/common'; - -import { CopilotProviderSideError, UserFriendlyError } from '../../../base'; -import { CopilotProvider } from './provider'; -import type { - CopilotProviderExecution, - ProviderDriverSpec, -} from './provider-runtime-contract'; -import { CopilotProviderType } from './types'; - -export type FalConfig = { - apiKey: string; -}; - -@Injectable() -export class FalProvider extends CopilotProvider { - override type = CopilotProviderType.FAL; - - protected resolveModelBackendKind() { - return 'fal' as const; - } - - override configured(execution?: CopilotProviderExecution): boolean { - return !!this.getConfig(execution).apiKey; - } - - private createNativeConfig(execution?: CopilotProviderExecution) { - return { - base_url: 'https://fal.run', - auth_token: this.getConfig(execution).apiKey, - }; - } - - override getDriverSpec(): ProviderDriverSpec { - return { - createBackendConfig: execution => this.createNativeConfig(execution), - mapError: error => this.handleError(error), - chat: false, - structured: false, - embedding: false, - rerank: false, - image: {}, - }; - } - - private handleError(e: any) { - if (e instanceof UserFriendlyError) { - // pass through user friendly errors - return e; - } else { - const error = new CopilotProviderSideError({ - provider: this.type, - kind: 'unexpected_response', - message: e?.message || 'Unexpected fal response', - }); - return error; - } - } -} diff --git a/packages/backend/server/src/plugins/copilot/providers/gemini/gemini.ts b/packages/backend/server/src/plugins/copilot/providers/gemini/gemini.ts deleted file mode 100644 index 99552f824b..0000000000 --- a/packages/backend/server/src/plugins/copilot/providers/gemini/gemini.ts +++ /dev/null @@ -1,250 +0,0 @@ -import { setTimeout as delay } from 'node:timers/promises'; - -import { Inject } from '@nestjs/common'; -import { ZodError } from 'zod'; - -import { - CopilotProviderSideError, - OneMB, - UserFriendlyError, -} from '../../../../base'; -import { - isInvalidStructuredOutputError, - type LlmBackendConfig, - llmResolveRequestIntentOptions, -} from '../../../../native'; -import { - admittedAttachmentToPromptAttachment, - AttachmentAdmissionHost, -} from '../../runtime/hosts/attachment-admission'; -import { - planAdmittedAttachmentMaterialization, - planHostUrlAttachmentMaterialization, -} from '../../runtime/hosts/attachment-materialization-planner'; -import { AttachmentMaterializer } from '../../runtime/hosts/attachment-materializer'; -import { CopilotProvider } from '../provider'; -import { hasProviderModelBehaviorFlag } from '../provider-model-runtime'; -import { - type CopilotProviderExecution, - type ProviderDriverSpec, -} from '../provider-runtime-contract'; -import type { PromptAttachment, PromptMessage } from '../types'; -import { promptAttachmentMimeType, promptAttachmentToUrl } from '../utils'; - -export const DEFAULT_DIMENSIONS = 256; -const GEMINI_REMOTE_ATTACHMENT_MAX_BYTES = 64 * OneMB; -const TRUSTED_ATTACHMENT_HOST_SUFFIXES = ['cdn.affine.pro']; -const GEMINI_RETRY_INITIAL_DELAY_MS = 2_000; - -function normalizeMimeType(mediaType?: string) { - return mediaType?.split(';', 1)[0]?.trim() || 'application/octet-stream'; -} - -export abstract class GeminiProvider extends CopilotProvider { - @Inject() protected readonly attachmentMaterializer!: AttachmentMaterializer; - @Inject() - protected readonly attachmentAdmissionHost?: AttachmentAdmissionHost; - - protected resolveModelBackendKind() { - return this.type === 'geminiVertex' - ? ('gemini_vertex' as const) - : ('gemini_api' as const); - } - - protected abstract createNativeConfig( - execution?: CopilotProviderExecution - ): Promise; - - private handleError(e: any) { - if (e instanceof UserFriendlyError) { - return e; - } else { - return new CopilotProviderSideError({ - provider: this.type, - kind: 'unexpected_response', - message: e?.message || 'Unexpected google response', - }); - } - } - - private getAttachmentAdmissionHost() { - return ( - this.attachmentAdmissionHost ?? - new AttachmentAdmissionHost(this.attachmentMaterializer) - ); - } - - protected async prepareMessages( - messages: PromptMessage[], - backendConfig: LlmBackendConfig, - options?: { - signal?: AbortSignal; - user?: string; - workspace?: string; - session?: string; - } - ): Promise { - const prepared: PromptMessage[] = []; - - for (const message of messages) { - options?.signal?.throwIfAborted(); - if (!Array.isArray(message.attachments) || !message.attachments.length) { - prepared.push(message); - continue; - } - - const attachments: PromptAttachment[] = []; - let changed = false; - for (const attachment of message.attachments) { - options?.signal?.throwIfAborted(); - const rawUrl = promptAttachmentToUrl(attachment); - if (!rawUrl || rawUrl.startsWith('data:')) { - attachments.push(attachment); - continue; - } - - try { - new URL(rawUrl); - } catch { - attachments.push(attachment); - continue; - } - - const declaredMimeType = promptAttachmentMimeType( - attachment, - typeof message.params?.mimetype === 'string' - ? message.params.mimetype - : undefined - ); - const referencePlan = await planHostUrlAttachmentMaterialization( - 'gemini', - backendConfig, - { - attachmentId: rawUrl, - url: rawUrl, - expectedMime: declaredMimeType - ? normalizeMimeType(declaredMimeType) - : undefined, - maxSize: GEMINI_REMOTE_ATTACHMENT_MAX_BYTES, - } - ); - if (referencePlan.mode === 'remote_reference') { - attachments.push(attachment); - continue; - } - - const admitted = - await this.getAttachmentAdmissionHost().admitPromptAttachment( - attachment, - { - userId: options?.user ?? 'provider-runtime', - workspaceId: options?.workspace ?? 'provider-runtime', - sessionId: options?.session, - signal: options?.signal, - maxBytes: referencePlan.request.maxSize, - trustedHostSuffixes: TRUSTED_ATTACHMENT_HOST_SUFFIXES, - } - ); - const materialization = planAdmittedAttachmentMaterialization(admitted); - attachments.push( - materialization.mode === 'inline' - ? materialization.attachment - : admittedAttachmentToPromptAttachment(admitted) - ); - changed = true; - } - - prepared.push(changed ? { ...message, attachments } : message); - } - - return prepared; - } - - protected async waitForStructuredRetry( - delayMs: number, - signal?: AbortSignal - ) { - await delay(delayMs, undefined, signal ? { signal } : undefined); - } - - override getDriverSpec(): ProviderDriverSpec { - return { - createBackendConfig: execution => this.createNativeConfig(execution), - mapError: error => this.handleError(error), - chat: { - prepareMessages: async context => - await this.prepareMessages( - context.input.messages, - context.backendConfig, - context.options - ), - resolveRequestOptions: async context => { - const requestIntent = await llmResolveRequestIntentOptions({ - protocol: context.protocol, - backendConfig: context.backendConfig, - reasoning: { - enabled: context.options.reasoning, - supported: - hasProviderModelBehaviorFlag( - context.model, - 'reasoning_medium' - ) || - hasProviderModelBehaviorFlag(context.model, 'reasoning_high'), - effort: hasProviderModelBehaviorFlag( - context.model, - 'reasoning_high' - ) - ? 'high' - : 'medium', - includeReasoning: - hasProviderModelBehaviorFlag( - context.model, - 'reasoning_medium' - ) || - hasProviderModelBehaviorFlag(context.model, 'reasoning_high'), - }, - }); - - return { - attachmentCapability: this.getAttachCapability( - context.model, - context.outputType - ), - include: requestIntent.include, - reasoning: requestIntent.reasoning, - }; - }, - }, - structured: { - prepareMessages: (inputMessages, backendConfig, structuredOptions) => - this.prepareMessages(inputMessages, backendConfig, structuredOptions), - shouldRetry: async ({ error, attempt, options: structuredOptions }) => { - const isParsingError = - isInvalidStructuredOutputError(error) || error instanceof ZodError; - const retryableError = - isParsingError || !(error instanceof UserFriendlyError); - const maxRetries = Math.max(structuredOptions.maxRetries ?? 3, 0); - if (!retryableError || attempt >= maxRetries) { - return false; - } - if (!isParsingError) { - await this.waitForStructuredRetry( - GEMINI_RETRY_INITIAL_DELAY_MS * 2 ** attempt, - structuredOptions.signal - ); - } - return true; - }, - }, - embedding: { - defaultDimensions: DEFAULT_DIMENSIONS, - taskType: 'RETRIEVAL_DOCUMENT', - }, - rerank: false, - image: { - prepareMessages: (inputMessages, backendConfig, imageOptions) => - this.prepareMessages(inputMessages, backendConfig, imageOptions), - }, - }; - } -} diff --git a/packages/backend/server/src/plugins/copilot/providers/gemini/generative.ts b/packages/backend/server/src/plugins/copilot/providers/gemini/generative.ts deleted file mode 100644 index 76f1d157e9..0000000000 --- a/packages/backend/server/src/plugins/copilot/providers/gemini/generative.ts +++ /dev/null @@ -1,28 +0,0 @@ -import type { LlmBackendConfig } from '../../../../native'; -import type { CopilotProviderExecution } from '../provider-runtime-contract'; -import { CopilotProviderType } from '../types'; -import { GeminiProvider } from './gemini'; - -export type GeminiGenerativeConfig = { - apiKey: string; - baseURL?: string; -}; - -export class GeminiGenerativeProvider extends GeminiProvider { - override readonly type = CopilotProviderType.Gemini; - override configured(execution?: CopilotProviderExecution): boolean { - return !!this.getConfig(execution).apiKey; - } - - protected override async createNativeConfig( - execution?: CopilotProviderExecution - ): Promise { - const config = this.getConfig(execution); - return { - base_url: ( - config.baseURL || 'https://generativelanguage.googleapis.com/v1beta' - ).replace(/\/$/, ''), - auth_token: config.apiKey, - }; - } -} diff --git a/packages/backend/server/src/plugins/copilot/providers/gemini/index.ts b/packages/backend/server/src/plugins/copilot/providers/gemini/index.ts deleted file mode 100644 index fbacc7fdc2..0000000000 --- a/packages/backend/server/src/plugins/copilot/providers/gemini/index.ts +++ /dev/null @@ -1,2 +0,0 @@ -export * from './generative'; -export * from './vertex'; diff --git a/packages/backend/server/src/plugins/copilot/providers/gemini/vertex.ts b/packages/backend/server/src/plugins/copilot/providers/gemini/vertex.ts deleted file mode 100644 index 70bf196dc4..0000000000 --- a/packages/backend/server/src/plugins/copilot/providers/gemini/vertex.ts +++ /dev/null @@ -1,34 +0,0 @@ -import type { LlmBackendConfig } from '../../../../native'; -import type { CopilotProviderExecution } from '../provider-runtime-contract'; -import { CopilotProviderType } from '../types'; -import { - getGoogleAuth, - getVertexGoogleBaseUrl, - type VertexProviderConfig, -} from '../utils'; -import { GeminiProvider } from './gemini'; - -export type GeminiVertexConfig = VertexProviderConfig; - -export class GeminiVertexProvider extends GeminiProvider { - override readonly type = CopilotProviderType.GeminiVertex; - override configured(execution?: CopilotProviderExecution): boolean { - const config = this.getConfig(execution); - return !!getVertexGoogleBaseUrl(config) && !!config.googleAuthOptions; - } - protected async resolveVertexAuth(execution?: CopilotProviderExecution) { - return await getGoogleAuth(this.getConfig(execution), 'google'); - } - - protected override async createNativeConfig( - execution?: CopilotProviderExecution - ): Promise { - const auth = await this.resolveVertexAuth(execution); - const { Authorization: authHeader } = auth.headers(); - - return { - base_url: auth.baseUrl || '', - auth_token: authHeader.replace(/^Bearer\s+/i, ''), - }; - } -} diff --git a/packages/backend/server/src/plugins/copilot/providers/index.ts b/packages/backend/server/src/plugins/copilot/providers/index.ts deleted file mode 100644 index 44a762bc7f..0000000000 --- a/packages/backend/server/src/plugins/copilot/providers/index.ts +++ /dev/null @@ -1,14 +0,0 @@ -export { - AnthropicOfficialProvider, - AnthropicVertexProvider, -} from './anthropic'; -export { CloudflareWorkersAIProvider } from './cloudflare'; -export { CopilotProviderFactory } from './factory'; -export { FalProvider } from './fal'; -export { GeminiGenerativeProvider, GeminiVertexProvider } from './gemini'; -export { CopilotProviderLifecycleService } from './lifecycle-service'; -export { OpenAIProvider } from './openai'; -export type { CopilotProvider } from './provider'; -export { CopilotProviders } from './provider-tokens'; -export { CopilotProviderRegistryService } from './registry-service'; -export * from './types'; diff --git a/packages/backend/server/src/plugins/copilot/providers/lifecycle-service.ts b/packages/backend/server/src/plugins/copilot/providers/lifecycle-service.ts deleted file mode 100644 index 7835a74ada..0000000000 --- a/packages/backend/server/src/plugins/copilot/providers/lifecycle-service.ts +++ /dev/null @@ -1,90 +0,0 @@ -import { Injectable, Type } from '@nestjs/common'; -import { ModuleRef } from '@nestjs/core'; - -import { OnEvent } from '../../../base'; -import { CopilotProviderFactory } from './factory'; -import type { CopilotProvider } from './provider'; -import type { CopilotProviderExecution } from './provider-runtime-contract'; -import { CopilotProviders } from './provider-tokens'; -import { CopilotProviderRegistryService } from './registry-service'; - -@Injectable() -export class CopilotProviderLifecycleService { - private readonly registeredByProvider = new WeakMap< - CopilotProvider, - Set - >(); - - constructor( - private readonly moduleRef: ModuleRef, - private readonly factory: CopilotProviderFactory, - private readonly registries: CopilotProviderRegistryService - ) {} - - private getProviders(): CopilotProvider[] { - return CopilotProviders.flatMap(token => { - const provider = this.moduleRef.get(token as Type, { - strict: false, - }); - return provider ? [provider] : []; - }); - } - - private getRegisteredProviderIds(provider: CopilotProvider) { - const current = this.registeredByProvider.get(provider); - if (current) { - return current; - } - - const next = new Set(); - this.registeredByProvider.set(provider, next); - return next; - } - - private async syncProvider(provider: CopilotProvider) { - const registry = this.registries.getRegistry(); - const configuredIds = new Set(); - - for (const providerId of registry.byType.get(provider.type) ?? []) { - const profile = registry.profiles.get(providerId); - if (!profile) { - continue; - } - - const execution: CopilotProviderExecution = { providerId, profile }; - if (!provider.configured(execution)) { - this.factory.unregister(providerId, provider); - continue; - } - - configuredIds.add(providerId); - this.factory.register(providerId, provider); - } - - const previous = this.getRegisteredProviderIds(provider); - for (const providerId of previous) { - if (!configuredIds.has(providerId)) { - this.factory.unregister(providerId, provider); - } - } - this.registeredByProvider.set(provider, configuredIds); - } - - async syncProviders() { - for (const provider of this.getProviders()) { - await this.syncProvider(provider); - } - } - - @OnEvent('config.init') - async onConfigInit() { - await this.syncProviders(); - } - - @OnEvent('config.changed') - async onConfigChanged(event: Events['config.changed']) { - if ('copilot' in event.updates) { - await this.syncProviders(); - } - } -} diff --git a/packages/backend/server/src/plugins/copilot/providers/openai.ts b/packages/backend/server/src/plugins/copilot/providers/openai.ts deleted file mode 100644 index fc501b0466..0000000000 --- a/packages/backend/server/src/plugins/copilot/providers/openai.ts +++ /dev/null @@ -1,172 +0,0 @@ -import { Inject } from '@nestjs/common'; - -import { - CopilotProviderSideError, - OneMB, - UserFriendlyError, -} from '../../../base'; -import { - type LlmBackendConfig, - llmResolveRequestIntentOptions, -} from '../../../native'; -import { - admittedAttachmentToPromptAttachment, - AttachmentAdmissionHost, -} from '../runtime/hosts/attachment-admission'; -import { AttachmentMaterializer } from '../runtime/hosts/attachment-materializer'; -import { CopilotProvider } from './provider'; -import { hasProviderModelBehaviorFlag } from './provider-model-runtime'; -import type { - CopilotProviderExecution, - ProviderDriverSpec, -} from './provider-runtime-contract'; -import { - CopilotProviderType, - type PromptAttachment, - type PromptMessage, -} from './types'; -import { promptAttachmentToUrl } from './utils'; - -export const DEFAULT_DIMENSIONS = 256; - -export type OpenAIConfig = { - apiKey: string; - baseURL?: string; - oldApiStyle?: boolean; -}; - -export class OpenAIProvider extends CopilotProvider { - readonly type = CopilotProviderType.OpenAI; - @Inject() protected readonly attachmentMaterializer!: AttachmentMaterializer; - @Inject() - protected readonly attachmentAdmissionHost?: AttachmentAdmissionHost; - - protected resolveModelBackendKind(execution?: CopilotProviderExecution) { - return this.getConfig(execution).oldApiStyle - ? ('openai_chat' as const) - : ('openai_responses' as const); - } - - override configured(execution?: CopilotProviderExecution): boolean { - return !!this.getConfig(execution).apiKey; - } - - private handleError(e: any) { - if (e instanceof UserFriendlyError) { - return e; - } - return new CopilotProviderSideError({ - provider: this.type, - kind: 'unexpected_response', - message: e?.message || 'Unexpected openai response', - }); - } - - protected createNativeConfig( - execution?: CopilotProviderExecution - ): LlmBackendConfig { - const config = this.getConfig(execution); - const baseUrl = config.baseURL || 'https://api.openai.com/v1'; - return { - base_url: baseUrl.replace(/\/v1\/?$/, ''), - auth_token: config.apiKey, - }; - } - - private getAttachmentAdmissionHost() { - return ( - this.attachmentAdmissionHost ?? - new AttachmentAdmissionHost(this.attachmentMaterializer) - ); - } - - private async prepareImageMessages( - messages: PromptMessage[], - options: { - signal?: AbortSignal; - user?: string; - workspace?: string; - session?: string; - } - ) { - const prepared: PromptMessage[] = []; - - for (const message of messages) { - options.signal?.throwIfAborted(); - if (!Array.isArray(message.attachments) || !message.attachments.length) { - prepared.push(message); - continue; - } - - let changed = false; - const attachments: PromptAttachment[] = []; - for (const attachment of message.attachments) { - options.signal?.throwIfAborted(); - const url = promptAttachmentToUrl(attachment); - if (!url || url.startsWith('data:')) { - attachments.push(attachment); - continue; - } - - const admitted = - await this.getAttachmentAdmissionHost().admitPromptAttachment( - attachment, - { - userId: options.user ?? 'provider-runtime', - workspaceId: options.workspace ?? 'provider-runtime', - sessionId: options.session, - signal: options.signal, - maxBytes: 50 * OneMB, - } - ); - attachments.push(admittedAttachmentToPromptAttachment(admitted)); - changed = true; - } - - prepared.push(changed ? { ...message, attachments } : message); - } - - return prepared; - } - - override getDriverSpec(): ProviderDriverSpec { - return { - createBackendConfig: execution => this.createNativeConfig(execution), - mapError: error => this.handleError(error), - chat: { - resolveRequestOptions: async context => { - const requestIntent = await llmResolveRequestIntentOptions({ - protocol: context.protocol, - backendConfig: context.backendConfig, - include: context.options.webSearch ? ['citations'] : undefined, - reasoning: { - enabled: context.options.reasoning, - supported: hasProviderModelBehaviorFlag( - context.model, - 'reasoning_supported' - ), - }, - }); - - return { - attachmentCapability: this.getAttachCapability( - context.model, - context.outputType - ), - include: requestIntent.include, - reasoning: requestIntent.reasoning, - }; - }, - }, - structured: {}, - embedding: { - defaultDimensions: DEFAULT_DIMENSIONS, - taskType: 'RETRIEVAL_DOCUMENT', - }, - image: { - prepareMessages: async (messages, _backendConfig, options) => - await this.prepareImageMessages(messages, options), - }, - }; - } -} diff --git a/packages/backend/server/src/plugins/copilot/providers/provider-middleware.ts b/packages/backend/server/src/plugins/copilot/providers/provider-middleware.ts deleted file mode 100644 index 4d9a88e388..0000000000 --- a/packages/backend/server/src/plugins/copilot/providers/provider-middleware.ts +++ /dev/null @@ -1,73 +0,0 @@ -import type { ProviderMiddlewareConfig } from '../config'; -import { CopilotProviderType } from './types'; - -const DEFAULT_NODE_TEXT_MIDDLEWARE: NonNullable< - NonNullable['text'] -> = ['citation_footnote', 'callout']; - -const DEFAULT_MIDDLEWARE_BY_TYPE: Record< - CopilotProviderType, - ProviderMiddlewareConfig -> = { - [CopilotProviderType.OpenAI]: { - node: { text: DEFAULT_NODE_TEXT_MIDDLEWARE }, - }, - [CopilotProviderType.CloudflareWorkersAi]: { - node: { text: DEFAULT_NODE_TEXT_MIDDLEWARE }, - }, - [CopilotProviderType.Anthropic]: { - node: { text: DEFAULT_NODE_TEXT_MIDDLEWARE }, - }, - [CopilotProviderType.AnthropicVertex]: { - node: { text: DEFAULT_NODE_TEXT_MIDDLEWARE }, - }, - [CopilotProviderType.Gemini]: { - node: { text: DEFAULT_NODE_TEXT_MIDDLEWARE }, - }, - [CopilotProviderType.GeminiVertex]: { - node: { text: DEFAULT_NODE_TEXT_MIDDLEWARE }, - }, - [CopilotProviderType.FAL]: {}, -}; - -function unique(items: T[]) { - return [...new Set(items)]; -} - -function mergeArray(base: T[] | undefined, override: T[] | undefined) { - if (!base?.length && !override?.length) { - return undefined; - } - return unique([...(base ?? []), ...(override ?? [])]); -} - -function compactMiddlewareSection>( - section: T -): T | undefined { - return Object.values(section).some(value => value !== undefined) - ? section - : undefined; -} - -export function mergeProviderMiddleware( - defaults: ProviderMiddlewareConfig, - override?: ProviderMiddlewareConfig -): ProviderMiddlewareConfig { - return { - rust: compactMiddlewareSection({ - request: mergeArray(defaults.rust?.request, override?.rust?.request), - stream: mergeArray(defaults.rust?.stream, override?.rust?.stream), - }), - node: compactMiddlewareSection({ - text: mergeArray(defaults.node?.text, override?.node?.text), - }), - }; -} - -export function resolveProviderMiddleware( - type: CopilotProviderType, - override?: ProviderMiddlewareConfig -): ProviderMiddlewareConfig { - const defaults = DEFAULT_MIDDLEWARE_BY_TYPE[type] ?? {}; - return mergeProviderMiddleware(defaults, override); -} diff --git a/packages/backend/server/src/plugins/copilot/providers/provider-model-runtime.ts b/packages/backend/server/src/plugins/copilot/providers/provider-model-runtime.ts deleted file mode 100644 index 1cfa10cfcd..0000000000 --- a/packages/backend/server/src/plugins/copilot/providers/provider-model-runtime.ts +++ /dev/null @@ -1,385 +0,0 @@ -import { z } from 'zod'; - -import { CopilotPromptInvalid } from '../../../base'; -import { - type LlmBackendConfig, - llmInferPromptModelConditions, - llmMatchModelCapabilities, - llmMatchModelRegistry, - type LlmProtocol, - llmResolveModelRegistryVariant, -} from '../../../native'; -import { applyPromptAttachmentMimeTypeHintForNative } from './attachments'; -import { - type CopilotChatOptions, - type CopilotImageOptions, - type CopilotModelBackendKind, - type CopilotProviderModel, - type CopilotProviderType, - type CopilotStructuredOptions, - EmbeddingMessage, - type ModelAttachmentCapability, - type ModelCapability, - type ModelFullConditions, - ModelInputType, - ModelOutputType, - type PromptAttachmentKind, - type PromptAttachmentSourceKind, - type PromptMessage, - PromptMessageSchema, -} from './types'; - -// Owner: backend host model-selection glue. -// Capability matching and catalog lookup are delegated to native/adapter; this -// file keeps provider prefix/default/prefer behavior and Node prompt checks. -export type ProviderModelRuntimeContext = { - type: CopilotProviderType; - backendKind: CopilotModelBackendKind; -}; - -export type ResolvedProviderModel = CopilotProviderModel & { - backendKind: CopilotModelBackendKind; - canonicalKey: string; - protocol?: LlmProtocol; - requestLayer?: LlmBackendConfig['request_layer']; - routeOverrides?: Partial< - Record< - ModelOutputType, - { - protocol?: LlmProtocol; - requestLayer?: LlmBackendConfig['request_layer']; - } - > - >; - behaviorFlags?: string[]; -}; - -function unique(values: Iterable) { - return Array.from(new Set(values)); -} - -function resolveAttachmentCapability( - cap: ModelCapability, - outputType?: ModelOutputType -): ModelAttachmentCapability | undefined { - if (outputType === ModelOutputType.Structured) { - return cap.structuredAttachments ?? cap.attachments; - } - return cap.attachments; -} - -function toProviderModel( - variant: NonNullable< - ReturnType['variant'] - > -): ResolvedProviderModel { - return { - id: variant.rawModelId, - name: variant.displayName, - backendKind: variant.backendKind, - canonicalKey: variant.canonicalKey, - protocol: variant.protocol, - requestLayer: variant.requestLayer, - routeOverrides: variant.routeOverrides, - behaviorFlags: variant.behaviorFlags, - capabilities: variant.capabilities.map(capability => ({ - input: capability.input as ModelInputType[], - output: capability.output as ModelOutputType[], - attachments: capability.attachments - ? { - kinds: capability.attachments.kinds as PromptAttachmentKind[], - sourceKinds: capability.attachments.sourceKinds as - | ModelAttachmentCapability['sourceKinds'] - | undefined, - allowRemoteUrls: capability.attachments.allowRemoteUrls, - } - : undefined, - structuredAttachments: capability.structuredAttachments - ? { - kinds: capability.structuredAttachments - .kinds as PromptAttachmentKind[], - sourceKinds: capability.structuredAttachments.sourceKinds as - | ModelAttachmentCapability['sourceKinds'] - | undefined, - allowRemoteUrls: capability.structuredAttachments.allowRemoteUrls, - } - : undefined, - defaultForOutputType: capability.defaultForOutputType, - })), - }; -} - -export type ProviderModelSelection = { - kind: 'configured'; - model: ResolvedProviderModel; -}; - -export function resolveProviderModelSelection( - context: ProviderModelRuntimeContext, - cond: ModelFullConditions -): ProviderModelSelection | undefined { - if (cond.modelId) { - const resolved = llmResolveModelRegistryVariant({ - backendKind: context.backendKind, - modelId: cond.modelId, - }).variant; - if (!resolved) { - return; - } - - const model = toProviderModel(resolved); - const matchedModelId = llmMatchModelCapabilities([model], { - ...cond, - modelId: model.id, - }); - if (!matchedModelId) { - return; - } - - return { - kind: 'configured', - model, - }; - } - - const resolved = llmMatchModelRegistry({ - backendKind: context.backendKind, - cond, - }).variant; - if (!resolved) { - return; - } - - return { - kind: 'configured', - model: toProviderModel(resolved), - }; -} - -function isMultimodal(model: CopilotProviderModel) { - return model.capabilities.some(c => - [ModelInputType.Image, ModelInputType.Audio, ModelInputType.File].some(t => - c.input.includes(t) - ) - ); -} - -function handleZodError(ret: z.SafeParseReturnType) { - if (ret.success) return; - const issues = ret.error.issues.map(i => { - const path = - 'root' + - (i.path.length - ? `.${i.path.map(seg => (typeof seg === 'number' ? `[${seg}]` : `.${seg}`)).join('')}` - : ''); - return `${i.message}${path}`; - }); - throw new CopilotPromptInvalid(issues.join('; ')); -} - -export async function inferModelConditionsFromMessages( - messages?: PromptMessage[], - withAttachment = true -): Promise> { - if (!messages?.length || !withAttachment) return {}; - const projectedMessages = messages.map(message => ({ - role: message.role, - content: message.content, - ...(Array.isArray(message.attachments) && message.attachments.length - ? { - attachments: message.attachments.map(attachment => - applyPromptAttachmentMimeTypeHintForNative(attachment, message) - ), - } - : {}), - })); - const inferredCond = llmInferPromptModelConditions(projectedMessages); - - return { - ...(inferredCond.attachmentKinds?.length - ? { attachmentKinds: unique(inferredCond.attachmentKinds) } - : {}), - ...(inferredCond.attachmentSourceKinds?.length - ? { - attachmentSourceKinds: unique( - inferredCond.attachmentSourceKinds - ) as PromptAttachmentSourceKind[], - } - : {}), - ...(inferredCond.inputTypes?.length - ? { inputTypes: unique(inferredCond.inputTypes) as ModelInputType[] } - : {}), - ...(inferredCond.hasRemoteAttachments - ? { hasRemoteAttachments: true } - : {}), - }; -} - -export function mergeModelConditions( - cond: ModelFullConditions, - inferredCond: Partial -): ModelFullConditions { - return { - ...inferredCond, - ...cond, - inputTypes: unique([ - ...(inferredCond.inputTypes ?? []), - ...(cond.inputTypes ?? []), - ]), - attachmentKinds: unique([ - ...(inferredCond.attachmentKinds ?? []), - ...(cond.attachmentKinds ?? []), - ]), - attachmentSourceKinds: unique([ - ...(inferredCond.attachmentSourceKinds ?? []), - ...(cond.attachmentSourceKinds ?? []), - ]), - hasRemoteAttachments: - cond.hasRemoteAttachments ?? inferredCond.hasRemoteAttachments, - }; -} - -export function getAttachCapability( - model: CopilotProviderModel, - outputType: ModelOutputType -): ModelAttachmentCapability | undefined { - const capability = - model.capabilities.find(cap => cap.output.includes(outputType)) ?? - model.capabilities[0]; - if (!capability) { - return; - } - return resolveAttachmentCapability(capability, outputType); -} - -export function matchProviderModel( - context: ProviderModelRuntimeContext, - cond: ModelFullConditions -): boolean { - return !!resolveProviderModelSelection(context, cond); -} - -export function resolveProviderModel( - context: ProviderModelRuntimeContext, - modelId: string -): ResolvedProviderModel | undefined { - return resolveProviderModelSelection(context, { - modelId, - })?.model; -} - -export function hasProviderModelBehaviorFlag( - model: CopilotProviderModel, - flag: string -) { - const behaviorFlags = (model as ResolvedProviderModel).behaviorFlags; - return Array.isArray(behaviorFlags) && behaviorFlags.includes(flag); -} - -export function resolveProviderModelRoute( - model: CopilotProviderModel, - outputType: ModelOutputType -) { - const resolved = model as ResolvedProviderModel; - const override = resolved.routeOverrides?.[outputType]; - - return { - protocol: override?.protocol ?? resolved.protocol, - requestLayer: override?.requestLayer ?? resolved.requestLayer, - }; -} - -export function requireProviderModelSelection( - context: ProviderModelRuntimeContext, - cond: ModelFullConditions -): ResolvedProviderModel { - const selection = resolveProviderModelSelection(context, cond); - if (selection) return selection.model; - - const { modelId, outputType, inputTypes } = cond; - throw new CopilotPromptInvalid( - modelId - ? `Model ${modelId} does not support ${outputType ?? ''} output with ${inputTypes ?? ''} input` - : outputType - ? `No model supports ${outputType} output with ${inputTypes ?? ''} input for provider ${context.type}` - : 'Output type is required when modelId is not provided' - ); -} - -export async function checkProviderParams( - context: ProviderModelRuntimeContext, - { - cond, - messages, - embeddings, - options = {}, - withAttachment = true, - }: { - cond: ModelFullConditions; - messages?: PromptMessage[]; - embeddings?: string[]; - options?: - | CopilotChatOptions - | CopilotStructuredOptions - | CopilotImageOptions; - withAttachment?: boolean; - execution?: unknown; - } -): Promise { - if (messages) { - const { requireContent = true, requireAttachment = false } = options; - - const MessageSchema = z - .array( - PromptMessageSchema.extend({ - content: requireContent - ? z.string().trim().min(1) - : z.string().optional().nullable(), - }) - .passthrough() - .catchall(z.union([z.string(), z.number(), z.date(), z.null()])) - ) - .optional(); - - handleZodError(MessageSchema.safeParse(messages)); - - const inferredCond = await inferModelConditionsFromMessages( - messages, - withAttachment - ); - const mergedCond = mergeModelConditions(cond, inferredCond); - const model = requireProviderModelSelection(context, mergedCond); - const multimodal = isMultimodal(model); - - if ( - multimodal && - requireAttachment && - !messages.some( - message => - message.role === 'user' && - Array.isArray(message.attachments) && - message.attachments.length > 0 - ) - ) { - throw new CopilotPromptInvalid('attachments required in multimodal mode'); - } - - if (embeddings) { - handleZodError(EmbeddingMessage.safeParse(embeddings)); - } - - return mergedCond; - } - - const inferredCond = await inferModelConditionsFromMessages( - messages, - withAttachment - ); - const mergedCond = mergeModelConditions(cond, inferredCond); - - if (embeddings) { - handleZodError(EmbeddingMessage.safeParse(embeddings)); - } - - return mergedCond; -} diff --git a/packages/backend/server/src/plugins/copilot/providers/provider-native-runtime.ts b/packages/backend/server/src/plugins/copilot/providers/provider-native-runtime.ts deleted file mode 100644 index c095dc5ea9..0000000000 --- a/packages/backend/server/src/plugins/copilot/providers/provider-native-runtime.ts +++ /dev/null @@ -1,337 +0,0 @@ -import type { - LlmBackendConfig, - LlmEmbeddingRequest, - LlmProtocol, - LlmRerankRequest, - LlmStructuredRequest, -} from '../../../native'; -import { - buildLlmImageRequestFromMessages, - llmEmbeddingDispatch, - llmRerankDispatch, - llmStructuredDispatch, -} from '../../../native'; -import type { NodeTextMiddleware, ProviderMiddlewareConfig } from '../config'; -import { - buildToolContracts, - projectPromptMessageForNative, -} from '../runtime/contracts'; -import { buildNativeRequest } from '../runtime/native-request-runtime'; -import type { ToolLoopBackend } from '../runtime/tool/bridge'; -import type { NativeProviderAdapter } from '../runtime/tool/native-adapter'; -import type { CopilotToolSet } from '../tools'; -import type { - CopilotProviderExecution, - PreparedNativeEmbeddingExecution, - PreparedNativeExecution, - PreparedNativeImageExecution, - PreparedNativeRequestOptions, - PreparedNativeRerankExecution, - PreparedNativeStructuredExecution, -} from './provider-runtime-contract'; -import type { - CopilotChatOptions, - CopilotImageOptions, - PromptMessage, -} from './types'; - -export type CreateToolAdapterOptions = { - maxSteps?: number; - nodeTextMiddleware?: NodeTextMiddleware[]; -}; - -export type CreateNativeAdapter = ( - backend: ToolLoopBackend, - tools: CopilotToolSet, - nodeTextMiddleware?: NodeTextMiddleware[], - options?: CreateToolAdapterOptions -) => NativeProviderAdapter; - -export type CreatePreparedExecutionRuntimeInput = { - resolveProviderId: (execution?: CopilotProviderExecution) => string; - getTools: ( - options: CopilotChatOptions, - model: string - ) => Promise; - getActiveProviderMiddleware: ( - execution?: CopilotProviderExecution - ) => ProviderMiddlewareConfig; - createNativeAdapter: CreateNativeAdapter; - maxSteps: number; -}; -export type PreparedExecutionRuntime = ReturnType< - typeof createPreparedExecutionRuntime ->; - -export function createPreparedExecutionRuntime( - input: CreatePreparedExecutionRuntimeInput -) { - return { - buildPreparedNativeExecution: async ( - prepared: PreparedNativeRequestOptions - ) => - await buildPreparedNativeExecution( - input.resolveProviderId(prepared.execution), - input.getTools, - input.getActiveProviderMiddleware, - input.maxSteps, - prepared - ), - createPreparedExecutionAdapter: (prepared: PreparedNativeExecution) => - createPreparedExecutionAdapter( - input.createNativeAdapter, - input.maxSteps, - prepared - ), - buildPreparedNativeStructuredExecution: ( - protocol: LlmProtocol, - backendConfig: LlmBackendConfig, - model: string, - request: LlmStructuredRequest, - execution?: CopilotProviderExecution - ) => - buildPreparedNativeStructuredExecution( - input.resolveProviderId(execution), - protocol, - backendConfig, - model, - request - ), - buildPreparedNativeEmbeddingExecution: ( - protocol: LlmProtocol, - backendConfig: LlmBackendConfig, - model: string, - request: LlmEmbeddingRequest, - execution?: CopilotProviderExecution - ) => - buildPreparedNativeEmbeddingExecution( - input.resolveProviderId(execution), - protocol, - backendConfig, - model, - request - ), - buildPreparedNativeRerankExecution: ( - protocol: LlmProtocol, - backendConfig: LlmBackendConfig, - model: string, - request: LlmRerankRequest, - execution?: CopilotProviderExecution - ) => - buildPreparedNativeRerankExecution( - input.resolveProviderId(execution), - protocol, - backendConfig, - model, - request - ), - buildPreparedNativeImageExecution: ( - protocol: LlmProtocol, - backendConfig: LlmBackendConfig, - model: string, - messages: PromptMessage[], - options: CopilotImageOptions = {}, - execution?: CopilotProviderExecution - ) => - buildPreparedNativeImageExecution( - input.resolveProviderId(execution), - protocol, - backendConfig, - model, - messages, - options - ), - }; -} - -export function createPreparedExecutionAdapter( - createNativeAdapter: CreateNativeAdapter, - maxSteps: number, - prepared: PreparedNativeExecution -) { - return createNativeAdapter( - { - protocol: prepared.route.protocol, - backendConfig: prepared.route.backendConfig, - }, - prepared.tools, - prepared.postprocess?.nodeTextMiddleware, - { - maxSteps, - nodeTextMiddleware: prepared.postprocess?.nodeTextMiddleware, - } - ); -} - -export function createNativeStructuredDispatch( - backendConfig: LlmBackendConfig, - protocol: LlmProtocol -) { - return (request: LlmStructuredRequest) => - llmStructuredDispatch(protocol, backendConfig, request); -} - -export function createNativeEmbeddingDispatch( - backendConfig: LlmBackendConfig, - protocol: LlmProtocol -) { - return (request: LlmEmbeddingRequest) => - llmEmbeddingDispatch(protocol, backendConfig, request); -} - -export function createNativeRerankDispatch( - backendConfig: LlmBackendConfig, - protocol: LlmProtocol -) { - return (request: LlmRerankRequest) => - llmRerankDispatch(protocol, backendConfig, request); -} - -function buildPreparedRoute( - providerId: string, - protocol: LlmProtocol, - backendConfig: LlmBackendConfig, - model: string -): PreparedNativeExecution['route'] { - return { - providerId, - protocol, - requestLayer: backendConfig.request_layer, - model, - backendConfig, - }; -} - -export async function buildPreparedNativeExecution( - providerId: string, - getTools: ( - options: CopilotChatOptions, - model: string - ) => Promise, - getActiveProviderMiddleware: ( - execution?: CopilotProviderExecution - ) => ProviderMiddlewareConfig, - maxSteps: number, - { - protocol, - backendConfig, - model, - messages, - options = {}, - execution, - withAttachment = true, - attachmentCapability, - include, - reasoning, - tools, - middleware, - }: PreparedNativeRequestOptions -): Promise { - const resolvedTools = tools ?? (await getTools(options, model)); - const resolvedMiddleware = - middleware ?? getActiveProviderMiddleware(execution); - const { request } = await buildNativeRequest({ - model, - messages, - options, - toolContracts: buildToolContracts(resolvedTools), - withAttachment, - attachmentCapability, - include, - reasoning, - middleware: resolvedMiddleware, - }); - - return { - route: buildPreparedRoute(providerId, protocol, backendConfig, model), - request, - tools: resolvedTools, - maxSteps, - postprocess: { - nodeTextMiddleware: resolvedMiddleware.node?.text, - }, - }; -} - -type BuildPreparedNativeDispatchExecution = < - TRequest extends - | LlmStructuredRequest - | LlmEmbeddingRequest - | LlmRerankRequest, ->( - providerId: string, - protocol: LlmProtocol, - backendConfig: LlmBackendConfig, - model: string, - request: TRequest -) => { - route: PreparedNativeExecution['route']; - request: TRequest; -}; - -const buildPreparedNativeDispatchExecution: BuildPreparedNativeDispatchExecution = - (providerId, protocol, backendConfig, model, request) => { - return { - route: buildPreparedRoute(providerId, protocol, backendConfig, model), - request, - }; - }; - -export const buildPreparedNativeStructuredExecution = - buildPreparedNativeDispatchExecution as ( - providerId: string, - protocol: LlmProtocol, - backendConfig: LlmBackendConfig, - model: string, - request: LlmStructuredRequest - ) => PreparedNativeStructuredExecution; - -export const buildPreparedNativeEmbeddingExecution = - buildPreparedNativeDispatchExecution as ( - providerId: string, - protocol: LlmProtocol, - backendConfig: LlmBackendConfig, - model: string, - request: LlmEmbeddingRequest - ) => PreparedNativeEmbeddingExecution; - -export const buildPreparedNativeRerankExecution = - buildPreparedNativeDispatchExecution as ( - providerId: string, - protocol: LlmProtocol, - backendConfig: LlmBackendConfig, - model: string, - request: LlmRerankRequest - ) => PreparedNativeRerankExecution; - -export function buildPreparedNativeImageExecution( - providerId: string, - protocol: LlmProtocol, - backendConfig: LlmBackendConfig, - model: string, - messages: PromptMessage[], - options: CopilotImageOptions = {} -): PreparedNativeImageExecution { - const nativeMessages = messages.map( - message => projectPromptMessageForNative(message).message - ); - - return { - route: buildPreparedRoute(providerId, protocol, backendConfig, model), - request: buildLlmImageRequestFromMessages({ - model, - protocol, - messages: nativeMessages, - options: projectImageRequestOptions(options), - }), - }; -} - -function projectImageRequestOptions(options: CopilotImageOptions = {}) { - return { - quality: options.quality, - seed: options.seed, - modelName: options.modelName, - loras: options.loras, - }; -} diff --git a/packages/backend/server/src/plugins/copilot/providers/provider-registry.ts b/packages/backend/server/src/plugins/copilot/providers/provider-registry.ts deleted file mode 100644 index 477facfe4f..0000000000 --- a/packages/backend/server/src/plugins/copilot/providers/provider-registry.ts +++ /dev/null @@ -1,287 +0,0 @@ -import type { - CopilotProviderConfigMap, - CopilotProviderDefaults, - CopilotProviderProfile, - ProviderMiddlewareConfig, -} from '../config'; -import { resolveProviderMiddleware } from './provider-middleware'; -import { CopilotProviderType, ModelOutputType } from './types'; - -const PROVIDER_ID_PATTERN = /^[a-zA-Z0-9-_]+$/; - -const LEGACY_PROVIDER_ORDER: CopilotProviderType[] = [ - CopilotProviderType.OpenAI, - CopilotProviderType.CloudflareWorkersAi, - CopilotProviderType.FAL, - CopilotProviderType.Gemini, - CopilotProviderType.GeminiVertex, - CopilotProviderType.Anthropic, - CopilotProviderType.AnthropicVertex, -]; - -const LEGACY_PROVIDER_PRIORITY = LEGACY_PROVIDER_ORDER.reduce( - (acc, type, index) => { - acc[type] = LEGACY_PROVIDER_ORDER.length - index; - return acc; - }, - {} as Record -); - -type LegacyProvidersConfig = Partial< - Record ->; - -export type CopilotProvidersConfigInput = LegacyProvidersConfig & { - profiles?: CopilotProviderProfile[] | null; - defaults?: CopilotProviderDefaults | null; -}; - -export type NormalizedCopilotProviderProfile = Omit< - CopilotProviderProfile, - 'enabled' | 'priority' | 'middleware' -> & { - enabled: boolean; - priority: number; - middleware: ProviderMiddlewareConfig; -}; - -export type CopilotProviderRegistry = { - profiles: Map; - defaults: CopilotProviderDefaults; - order: string[]; - byType: Map; -}; - -export type ResolveModelResult = { - rawModelId?: string; - modelId?: string; - explicitProviderId?: string; - candidateProviderIds: string[]; -}; - -type ResolveModelOptions = { - registry: CopilotProviderRegistry; - modelId?: string; - outputType?: ModelOutputType; - availableProviderIds?: Iterable; - preferredProviderIds?: Iterable; -}; - -function unique(list: T[]): T[] { - return [...new Set(list)]; -} - -function asArray(iter?: Iterable): T[] { - return iter ? Array.from(iter) : []; -} - -function parseModelPrefix( - registry: CopilotProviderRegistry, - modelId: string -): { providerId: string; modelId?: string } | null { - const index = modelId.indexOf('/'); - if (index <= 0) { - return null; - } - - const providerId = modelId.slice(0, index); - if (!registry.profiles.has(providerId)) { - return null; - } - - const model = modelId.slice(index + 1); - return { providerId, modelId: model || undefined }; -} - -function normalizeProfile( - profile: CopilotProviderProfile -): NormalizedCopilotProviderProfile { - return { - ...profile, - enabled: profile.enabled !== false, - priority: profile.priority ?? 0, - middleware: resolveProviderMiddleware(profile.type, profile.middleware), - }; -} - -function toLegacyProfiles( - config: CopilotProvidersConfigInput -): CopilotProviderProfile[] { - const legacyProfiles: CopilotProviderProfile[] = []; - for (const type of LEGACY_PROVIDER_ORDER) { - const legacyConfig = config[type]; - if (!legacyConfig) { - continue; - } - legacyProfiles.push({ - id: `${type}-default`, - type, - priority: LEGACY_PROVIDER_PRIORITY[type], - config: legacyConfig, - } as CopilotProviderProfile); - } - return legacyProfiles; -} - -function mergeProfiles( - explicitProfiles: CopilotProviderProfile[], - legacyProfiles: CopilotProviderProfile[] -): CopilotProviderProfile[] { - const profiles = new Map(); - - for (const profile of explicitProfiles) { - if (!PROVIDER_ID_PATTERN.test(profile.id)) { - throw new Error(`Invalid copilot provider profile id: ${profile.id}`); - } - if (profiles.has(profile.id)) { - throw new Error(`Duplicated copilot provider profile id: ${profile.id}`); - } - profiles.set(profile.id, profile); - } - - for (const profile of legacyProfiles) { - if (!profiles.has(profile.id)) { - profiles.set(profile.id, profile); - } - } - - return Array.from(profiles.values()); -} - -function sortProfiles(profiles: NormalizedCopilotProviderProfile[]) { - return profiles.toSorted((a, b) => { - if (a.priority !== b.priority) { - return b.priority - a.priority; - } - return a.id.localeCompare(b.id); - }); -} - -function assertDefaults( - defaults: CopilotProviderDefaults, - profiles: Map -) { - for (const providerId of Object.values(defaults)) { - if (!providerId) { - continue; - } - if (!profiles.has(providerId)) { - throw new Error( - `Copilot provider defaults references unknown providerId: ${providerId}` - ); - } - } -} - -export function buildProviderRegistry( - config: CopilotProvidersConfigInput -): CopilotProviderRegistry { - const explicitProfiles = config.profiles ?? []; - const legacyProfiles = toLegacyProfiles(config); - const mergedProfiles = mergeProfiles(explicitProfiles, legacyProfiles) - .map(normalizeProfile) - .filter(profile => profile.enabled); - const sortedProfiles = sortProfiles(mergedProfiles); - - const profiles = new Map( - sortedProfiles.map(profile => [profile.id, profile] as const) - ); - const defaults = config.defaults ?? {}; - assertDefaults(defaults, profiles); - - const order = sortedProfiles.map(profile => profile.id); - const byType = new Map(); - for (const profile of sortedProfiles) { - const ids = byType.get(profile.type) ?? []; - ids.push(profile.id); - byType.set(profile.type, ids); - } - - return { profiles, defaults, order, byType }; -} - -export function resolveModel({ - registry, - modelId, - outputType, - availableProviderIds, - preferredProviderIds, -}: ResolveModelOptions): ResolveModelResult { - const available = new Set(asArray(availableProviderIds)); - const preferred = new Set(asArray(preferredProviderIds)); - const hasAvailableFilter = available.size > 0; - const hasPreferredFilter = preferred.size > 0; - - const isAllowed = (providerId: string) => { - const profile = registry.profiles.get(providerId); - if (!profile?.enabled) { - return false; - } - if (hasAvailableFilter && !available.has(providerId)) { - return false; - } - if (hasPreferredFilter && !preferred.has(providerId)) { - return false; - } - return true; - }; - - const prefixed = modelId ? parseModelPrefix(registry, modelId) : null; - if (prefixed) { - return { - rawModelId: modelId, - modelId: prefixed.modelId, - explicitProviderId: prefixed.providerId, - candidateProviderIds: isAllowed(prefixed.providerId) - ? [prefixed.providerId] - : [], - }; - } - - if (modelId) { - return { - rawModelId: modelId, - modelId, - candidateProviderIds: registry.order.filter(providerId => - isAllowed(providerId) - ), - }; - } - - const defaultProviderId = - outputType && outputType !== ModelOutputType.Rerank - ? registry.defaults[outputType] - : undefined; - - const fallbackOrder = [ - ...(defaultProviderId ? [defaultProviderId] : []), - registry.defaults.fallback, - ...registry.order, - ].filter((id): id is string => !!id); - - return { - rawModelId: modelId, - modelId, - candidateProviderIds: unique( - fallbackOrder.filter(providerId => isAllowed(providerId)) - ), - }; -} - -export function stripProviderPrefix( - registry: CopilotProviderRegistry, - providerId: string, - modelId?: string -) { - if (!modelId) { - return modelId; - } - const prefixed = parseModelPrefix(registry, modelId); - if (!prefixed) { - return modelId; - } - if (prefixed.providerId !== providerId) { - return modelId; - } - return prefixed.modelId; -} diff --git a/packages/backend/server/src/plugins/copilot/providers/provider-runtime-contract.ts b/packages/backend/server/src/plugins/copilot/providers/provider-runtime-contract.ts deleted file mode 100644 index 1c74d50e57..0000000000 --- a/packages/backend/server/src/plugins/copilot/providers/provider-runtime-contract.ts +++ /dev/null @@ -1,456 +0,0 @@ -import type { - LlmBackendConfig, - LlmEmbeddingRequest, - LlmImageRequest, - LlmProtocol, - LlmRequest, - LlmRerankRequest, - LlmStructuredRequest, -} from '../../../native'; -import type { NodeTextMiddleware, ProviderMiddlewareConfig } from '../config'; -import type { CopilotToolSet } from '../tools'; -import { - type ProviderModelRuntimeContext, - resolveProviderModelRoute, -} from './provider-model-runtime'; -import type { NormalizedCopilotProviderProfile } from './provider-registry'; -import { - CopilotChatOptions, - CopilotImageOptions, - CopilotProviderModel, - CopilotStructuredOptions, - ModelAttachmentCapability, - ModelConditions, - ModelFullConditions, - ModelOutputType, - PromptMessage, -} from './types'; - -export type NativeExecutionRoute = { - protocol: LlmProtocol; - requestLayer?: LlmBackendConfig['request_layer']; - model: string; - backendConfig: LlmBackendConfig; -}; - -export type CopilotProviderExecution = { - providerId: string; - profile: NormalizedCopilotProviderProfile; -}; - -export type PreparedNativeExecution = { - route: NativeExecutionRoute & { - providerId: string; - }; - request: LlmRequest; - tools: CopilotToolSet; - maxSteps?: number; - postprocess?: { - nodeTextMiddleware?: NodeTextMiddleware[]; - }; -}; - -export type PreparedNativeStructuredExecution = { - route: NativeExecutionRoute & { - providerId: string; - }; - request: LlmStructuredRequest; -}; - -export type PreparedNativeEmbeddingExecution = { - route: NativeExecutionRoute & { - providerId: string; - }; - request: LlmEmbeddingRequest; -}; - -export type PreparedNativeRerankExecution = { - route: NativeExecutionRoute & { - providerId: string; - }; - request: LlmRerankRequest; -}; - -export type PreparedNativeImageExecution = { - route: NativeExecutionRoute & { - providerId: string; - }; - request: LlmImageRequest; -}; - -export type PreparedNativeRequestOptions = { - protocol: LlmProtocol; - backendConfig: LlmBackendConfig; - model: string; - messages: PromptMessage[]; - options?: CopilotChatOptions; - execution?: CopilotProviderExecution; - withAttachment?: boolean; - attachmentCapability?: ModelAttachmentCapability; - include?: string[]; - reasoning?: Record; - tools?: CopilotToolSet; - middleware?: ProviderMiddlewareConfig; -}; - -type ProviderChatDriverPrepareResult = Omit< - PreparedNativeRequestOptions, - 'execution' | 'options' ->; - -type Awaitable = T | Promise; - -type NativeBackendConfigResolver = ( - execution?: CopilotProviderExecution -) => Awaitable; - -export type StructuredProviderDriver = { - createBackendConfig: NativeBackendConfigResolver; - prepareMessages?: ( - messages: PromptMessage[], - backendConfig: LlmBackendConfig, - options: NonNullable - ) => Promise; - shouldRetry?: (context: { - error: unknown; - attempt: number; - options: NonNullable; - }) => Awaitable; - mapError: (error: unknown) => unknown; -}; - -export type EmbeddingProviderDriver = { - createBackendConfig: NativeBackendConfigResolver; - defaultDimensions?: number; - taskType?: string; - mapError: (error: unknown) => unknown; -}; - -export type RerankProviderDriver = { - createBackendConfig: NativeBackendConfigResolver; - mapError: (error: unknown) => unknown; -}; - -export type ImageProviderDriver = { - createBackendConfig: NativeBackendConfigResolver; - prepareMessages?: ( - messages: PromptMessage[], - backendConfig: LlmBackendConfig, - options: NonNullable - ) => Promise; - mapError: (error: unknown) => unknown; -}; - -export type ProviderMetricLabels = Record< - string, - string | number | boolean | undefined ->; - -export type ProviderExecutionDrivers = { - chat?: ProviderChatDriver; - structured?: StructuredProviderDriver; - embedding?: EmbeddingProviderDriver; - rerank?: RerankProviderDriver; - image?: ImageProviderDriver; -}; - -export type ProviderDriverSpec = NativeProviderDriverBase & { - chat?: NativeChatDriverOverrides | false; - structured?: NativeStructuredDriverOverrides | false; - embedding?: NativeEmbeddingDriverOverrides | false; - rerank?: NativeRerankDriverOverrides | false; - image?: NativeImageDriverOverrides | false; -}; - -export type ProviderRuntimeHostSeed = { - model: ProviderModelRuntimeContext; - resolveExecutionDrivers: () => ProviderExecutionDrivers | undefined; - selectModel: NativeChatDriverBase['selectModel']; - checkParams: NativeChatDriverBase['checkParams']; - getAttachCapability: ( - model: CopilotProviderModel, - outputType: ModelOutputType - ) => ModelAttachmentCapability | undefined; - getActiveProviderMiddleware: ( - execution?: CopilotProviderExecution - ) => ProviderMiddlewareConfig; - getTools: ( - options: CopilotChatOptions, - model: string - ) => Promise; - metricLabels: ( - model: string, - labels?: ProviderMetricLabels, - execution?: CopilotProviderExecution - ) => ProviderMetricLabels; -}; - -export type ProviderChatDriverPrepareInput = { - kind: 'text' | 'streamText' | 'streamObject'; - cond: ModelConditions; - messages: PromptMessage[]; - options: CopilotChatOptions; - execution?: CopilotProviderExecution; -}; - -export type ProviderChatDriver = { - prepare: (input: { - kind: ProviderChatDriverPrepareInput['kind']; - cond: ProviderChatDriverPrepareInput['cond']; - messages: ProviderChatDriverPrepareInput['messages']; - options: ProviderChatDriverPrepareInput['options']; - execution?: ProviderChatDriverPrepareInput['execution']; - }) => Promise; - mapError: (error: unknown) => unknown; -}; - -type NativeProviderDriverBase = Pick< - StructuredProviderDriver, - 'createBackendConfig' | 'mapError' ->; - -type ChatToolingResult = Pick< - ProviderChatDriverPrepareResult, - 'tools' | 'middleware' ->; - -type NativeChatDriverBase = NativeProviderDriverBase & { - checkParams: (input: { - cond: ModelFullConditions; - messages?: PromptMessage[]; - embeddings?: string[]; - options?: - | CopilotChatOptions - | CopilotStructuredOptions - | CopilotImageOptions; - withAttachment?: boolean; - execution?: CopilotProviderExecution; - }) => Promise; - selectModel: ( - cond: ModelFullConditions, - execution?: CopilotProviderExecution - ) => CopilotProviderModel; - getTools?: ( - options: CopilotChatOptions, - model: string - ) => Promise; - getActiveProviderMiddleware?: ( - execution?: CopilotProviderExecution - ) => ProviderMiddlewareConfig; -}; - -type NativeStructuredDriverOverrides = Partial; -type NativeEmbeddingDriverOverrides = Partial; -type NativeRerankDriverOverrides = Partial; -type NativeImageDriverOverrides = Partial; - -type NativeChatDriverContext = { - input: ProviderChatDriverPrepareInput; - outputType: ModelOutputType; - normalizedCond: ModelFullConditions; - model: CopilotProviderModel; - backendConfig: LlmBackendConfig; - protocol: LlmProtocol; - messages: PromptMessage[]; - options: NonNullable; - execution?: CopilotProviderExecution; -}; - -type NativeChatDriverOverrides = { - resolveOutputType?: ( - kind: ProviderChatDriverPrepareInput['kind'] - ) => ModelOutputType | null; - withAttachment?: boolean; - prepareMessages?: ( - context: Omit - ) => Awaitable; - resolveTooling?: ( - context: NativeChatDriverContext - ) => Awaitable; - resolveRequestOptions?: ( - context: NativeChatDriverContext - ) => Awaitable< - Partial< - Pick< - ProviderChatDriverPrepareResult, - 'withAttachment' | 'attachmentCapability' | 'include' | 'reasoning' - > - > - >; -}; - -export function createNativeProviderDriverFactory( - base: NativeProviderDriverBase -) { - return { - structured( - overrides: NativeStructuredDriverOverrides = {} - ): StructuredProviderDriver { - return { - createBackendConfig: - overrides.createBackendConfig ?? base.createBackendConfig, - mapError: overrides.mapError ?? base.mapError, - ...(overrides.prepareMessages - ? { prepareMessages: overrides.prepareMessages } - : {}), - ...(overrides.shouldRetry - ? { shouldRetry: overrides.shouldRetry } - : {}), - }; - }, - embedding( - overrides: NativeEmbeddingDriverOverrides = {} - ): EmbeddingProviderDriver { - return { - createBackendConfig: - overrides.createBackendConfig ?? base.createBackendConfig, - mapError: overrides.mapError ?? base.mapError, - ...(overrides.defaultDimensions !== undefined - ? { defaultDimensions: overrides.defaultDimensions } - : {}), - ...(overrides.taskType ? { taskType: overrides.taskType } : {}), - }; - }, - rerank(overrides: NativeRerankDriverOverrides = {}): RerankProviderDriver { - return { - createBackendConfig: - overrides.createBackendConfig ?? base.createBackendConfig, - mapError: overrides.mapError ?? base.mapError, - }; - }, - image(overrides: NativeImageDriverOverrides = {}): ImageProviderDriver { - return { - createBackendConfig: - overrides.createBackendConfig ?? base.createBackendConfig, - mapError: overrides.mapError ?? base.mapError, - ...(overrides.prepareMessages - ? { prepareMessages: overrides.prepareMessages } - : {}), - }; - }, - }; -} - -function compileProviderChatDriver( - spec: NativeProviderDriverBase & NativeChatDriverOverrides, - base: NativeChatDriverBase -): ProviderChatDriver { - return { - prepare: async (input: ProviderChatDriverPrepareInput) => { - const options: NonNullable = input.options ?? {}; - const resolvedOutputType = spec.resolveOutputType?.(input.kind); - const outputType = - resolvedOutputType === undefined - ? input.kind === 'streamObject' - ? ModelOutputType.Object - : ModelOutputType.Text - : resolvedOutputType; - if (!outputType) { - return null; - } - - const normalizedCond = await base.checkParams({ - messages: input.messages, - cond: { - ...input.cond, - outputType, - }, - options, - execution: input.execution, - ...(spec.withAttachment !== undefined - ? { withAttachment: spec.withAttachment } - : {}), - }); - const model = base.selectModel(normalizedCond, input.execution); - const backendConfig = await spec.createBackendConfig(input.execution); - const route = resolveProviderModelRoute(model, outputType); - if (!route.protocol) { - throw new Error(`Missing native protocol for model ${model.id}`); - } - const partialContext = { - input, - outputType, - normalizedCond, - model, - backendConfig: - route.requestLayer === backendConfig.request_layer - ? backendConfig - : { ...backendConfig, request_layer: route.requestLayer }, - protocol: route.protocol, - options, - execution: input.execution, - }; - const messages = spec.prepareMessages - ? await spec.prepareMessages(partialContext) - : input.messages; - const context = { - ...partialContext, - messages, - }; - const tooling = spec.resolveTooling - ? await spec.resolveTooling(context) - : { - ...(base.getTools - ? { tools: await base.getTools(options, model.id) } - : {}), - ...(base.getActiveProviderMiddleware - ? { - middleware: base.getActiveProviderMiddleware(input.execution), - } - : {}), - }; - const requestOptions = spec.resolveRequestOptions - ? await spec.resolveRequestOptions(context) - : {}; - - return { - protocol: context.protocol, - backendConfig: context.backendConfig, - model: model.id, - messages, - ...(spec.withAttachment === false ? { withAttachment: false } : {}), - ...requestOptions, - ...tooling, - }; - }, - mapError: spec.mapError, - }; -} - -export function createNativeExecutionDriverSpec( - input: ProviderDriverSpec, - runtimeBase: NativeChatDriverBase -): ProviderExecutionDrivers { - const driverBase = { - createBackendConfig: input.createBackendConfig, - mapError: input.mapError, - }; - const nativeDrivers = createNativeProviderDriverFactory(driverBase); - - return { - ...(input.chat !== false - ? { - chat: compileProviderChatDriver( - { ...driverBase, ...input.chat }, - runtimeBase - ), - } - : {}), - ...(input.structured !== false - ? { - structured: nativeDrivers.structured(input.structured ?? undefined), - } - : {}), - ...(input.embedding !== false - ? { - embedding: nativeDrivers.embedding(input.embedding ?? undefined), - } - : {}), - ...(input.rerank !== false - ? { rerank: nativeDrivers.rerank(input.rerank ?? undefined) } - : {}), - ...(input.image !== false - ? { image: nativeDrivers.image(input.image ?? undefined) } - : {}), - }; -} diff --git a/packages/backend/server/src/plugins/copilot/providers/provider-tokens.ts b/packages/backend/server/src/plugins/copilot/providers/provider-tokens.ts deleted file mode 100644 index 098cbee5c7..0000000000 --- a/packages/backend/server/src/plugins/copilot/providers/provider-tokens.ts +++ /dev/null @@ -1,18 +0,0 @@ -import { - AnthropicOfficialProvider, - AnthropicVertexProvider, -} from './anthropic'; -import { CloudflareWorkersAIProvider } from './cloudflare'; -import { FalProvider } from './fal'; -import { GeminiGenerativeProvider, GeminiVertexProvider } from './gemini'; -import { OpenAIProvider } from './openai'; - -export const CopilotProviders = [ - OpenAIProvider, - CloudflareWorkersAIProvider, - FalProvider, - GeminiGenerativeProvider, - GeminiVertexProvider, - AnthropicOfficialProvider, - AnthropicVertexProvider, -]; diff --git a/packages/backend/server/src/plugins/copilot/providers/provider.ts b/packages/backend/server/src/plugins/copilot/providers/provider.ts deleted file mode 100644 index 5c46fc212a..0000000000 --- a/packages/backend/server/src/plugins/copilot/providers/provider.ts +++ /dev/null @@ -1,249 +0,0 @@ -import { Inject, Injectable, Logger } from '@nestjs/common'; - -import { Config } from '../../../base'; -import type { NodeTextMiddleware, ProviderMiddlewareConfig } from '../config'; -import { ToolExecutorHost } from '../runtime/hosts/tool-executor-host'; -import { mapNativeSemanticError } from '../runtime/native-errors'; -import type { ToolLoopBackend } from '../runtime/tool/bridge'; -import type { CopilotTool, CopilotToolSet } from '../tools'; -import { resolveProviderMiddleware } from './provider-middleware'; -import { - checkProviderParams, - getAttachCapability as getAttachCapabilityHelper, - matchProviderModel as matchProviderModelHelper, - type ProviderModelRuntimeContext, - requireProviderModelSelection, - resolveProviderModel, -} from './provider-model-runtime'; -import { - type CopilotProviderExecution, - createNativeExecutionDriverSpec, - type ProviderDriverSpec, - type ProviderExecutionDrivers, - type ProviderRuntimeHostSeed, -} from './provider-runtime-contract'; -import { - type CopilotChatOptions, - CopilotChatTools, - type CopilotImageOptions, - type CopilotModelBackendKind, - CopilotProviderModel, - CopilotProviderType, - type CopilotStructuredOptions, - type ModelAttachmentCapability, - ModelFullConditions, - ModelOutputType, - type PromptMessage, -} from './types'; -export type { - CopilotProviderExecution, - ProviderDriverSpec, - ProviderExecutionDrivers, - ProviderRuntimeHostSeed, -} from './provider-runtime-contract'; - -@Injectable() -export abstract class CopilotProvider { - protected readonly logger = new Logger(this.constructor.name); - protected readonly MAX_STEPS = 20; - - abstract readonly type: CopilotProviderType; - protected abstract resolveModelBackendKind( - execution?: CopilotProviderExecution - ): CopilotModelBackendKind; - abstract configured(execution?: CopilotProviderExecution): boolean; - - @Inject() protected readonly AFFiNEConfig!: Config; - @Inject() protected readonly toolExecutorHost!: ToolExecutorHost; - - get maxSteps() { - return this.MAX_STEPS; - } - - protected resolveModelRuntimeContext( - execution?: CopilotProviderExecution - ): ProviderModelRuntimeContext { - return { - type: this.type, - backendKind: this.resolveModelBackendKind(execution), - }; - } - - protected get modelRuntimeContext(): ProviderModelRuntimeContext { - return this.resolveModelRuntimeContext(); - } - - getDriverSpec(): ProviderDriverSpec | undefined { - return undefined; - } - - getExecutionDrivers(): ProviderExecutionDrivers | undefined { - const spec = this.getDriverSpec(); - return spec ? this.createDriverSpec(spec) : undefined; - } - - protected createDriverSpec( - spec: ProviderDriverSpec - ): ProviderExecutionDrivers { - return createNativeExecutionDriverSpec(spec, { - createBackendConfig: spec.createBackendConfig, - mapError: error => { - const mapped = mapNativeSemanticError(error); - return mapped === error ? spec.mapError(error) : mapped; - }, - checkParams: input => - checkProviderParams( - this.resolveModelRuntimeContext(input.execution), - input - ), - selectModel: (cond, execution) => - requireProviderModelSelection( - this.resolveModelRuntimeContext(execution), - cond - ), - getTools: this.getTools.bind(this), - getActiveProviderMiddleware: this.getActiveProviderMiddleware.bind(this), - }); - } - - selectModel( - cond: ModelFullConditions, - execution?: CopilotProviderExecution - ): CopilotProviderModel { - return requireProviderModelSelection( - this.resolveModelRuntimeContext(execution), - cond - ); - } - - checkParams(input: { - cond: ModelFullConditions; - messages?: PromptMessage[]; - embeddings?: string[]; - options?: - | CopilotChatOptions - | CopilotStructuredOptions - | CopilotImageOptions; - withAttachment?: boolean; - execution?: CopilotProviderExecution; - }) { - return checkProviderParams( - this.resolveModelRuntimeContext(input.execution), - input - ); - } - - getRuntimeHostSeed(): ProviderRuntimeHostSeed { - return { - model: this.resolveModelRuntimeContext(), - resolveExecutionDrivers: () => this.getExecutionDrivers(), - selectModel: this.selectModel.bind(this), - checkParams: this.checkParams.bind(this), - getAttachCapability: this.getAttachCapability.bind(this), - getActiveProviderMiddleware: this.getActiveProviderMiddleware.bind(this), - getTools: this.getTools.bind(this), - metricLabels: this.metricLabels.bind(this), - }; - } - - protected getExecutionProfile(execution?: CopilotProviderExecution) { - return execution?.profile?.type === this.type - ? execution.profile - : undefined; - } - - getActiveProviderMiddleware( - execution?: CopilotProviderExecution - ): ProviderMiddlewareConfig { - return ( - this.getExecutionProfile(execution)?.middleware ?? - resolveProviderMiddleware(this.type) - ); - } - - metricLabels( - model: string, - labels: Record = {}, - execution?: CopilotProviderExecution - ) { - return { - model, - providerId: execution?.providerId ?? `${this.type}-default`, - ...labels, - }; - } - - protected get config(): C { - return this.AFFiNEConfig.copilot.providers[this.type] as C; - } - - protected getConfig(execution?: CopilotProviderExecution): C { - const profile = this.getExecutionProfile(execution); - if (profile) { - return profile.config as C; - } - return this.config; - } - getAttachCapability( - model: CopilotProviderModel, - outputType: ModelOutputType - ): ModelAttachmentCapability | undefined { - return getAttachCapabilityHelper(model, outputType); - } - - // make it async to allow dynamic check available models in some providers - async match( - cond: ModelFullConditions = {}, - execution?: CopilotProviderExecution - ): Promise { - return ( - this.configured(execution) && - matchProviderModelHelper(this.resolveModelRuntimeContext(execution), cond) - ); - } - - resolveModel( - modelId: string, - execution?: CopilotProviderExecution - ): CopilotProviderModel | undefined { - return resolveProviderModel( - this.resolveModelRuntimeContext(execution), - modelId - ); - } - - protected getProviderSpecificTools( - _toolName: CopilotChatTools, - _model: string - ): [string, CopilotTool?] | undefined { - return; - } - - // use for tool use, shared between providers - async getTools( - options: CopilotChatOptions, - model: string - ): Promise { - this.logger.debug(`getTools: ${JSON.stringify(options?.tools ?? [])}`); - return await this.toolExecutorHost.getTools( - options, - model, - this.getProviderSpecificTools.bind(this) - ); - } - - createNativeAdapter( - backend: ToolLoopBackend, - tools: CopilotToolSet, - nodeTextMiddleware?: NodeTextMiddleware[], - options: { - maxSteps?: number; - nodeTextMiddleware?: NodeTextMiddleware[]; - } = {} - ) { - return this.toolExecutorHost.createNativeAdapter(backend, tools, { - ...options, - nodeTextMiddleware: nodeTextMiddleware ?? options.nodeTextMiddleware, - }); - } -} diff --git a/packages/backend/server/src/plugins/copilot/providers/registry-service.ts b/packages/backend/server/src/plugins/copilot/providers/registry-service.ts deleted file mode 100644 index da537cab0b..0000000000 --- a/packages/backend/server/src/plugins/copilot/providers/registry-service.ts +++ /dev/null @@ -1,28 +0,0 @@ -import { Injectable } from '@nestjs/common'; - -import { Config } from '../../../base'; -import { - buildProviderRegistry, - type CopilotProviderRegistry, - type CopilotProvidersConfigInput, -} from './provider-registry'; - -@Injectable() -export class CopilotProviderRegistryService { - private lastConfig?: CopilotProvidersConfigInput; - private lastRegistry?: CopilotProviderRegistry; - - constructor(private readonly config: Config) {} - - getRegistry(): CopilotProviderRegistry { - const providerConfig = this.config.copilot.providers; - if (this.lastConfig === providerConfig && this.lastRegistry) { - return this.lastRegistry; - } - - const registry = buildProviderRegistry(providerConfig); - this.lastConfig = providerConfig; - this.lastRegistry = registry; - return registry; - } -} diff --git a/packages/backend/server/src/plugins/copilot/providers/types.ts b/packages/backend/server/src/plugins/copilot/providers/types.ts index bee81d208c..55ccb4dd35 100644 --- a/packages/backend/server/src/plugins/copilot/providers/types.ts +++ b/packages/backend/server/src/plugins/copilot/providers/types.ts @@ -1,4 +1,4 @@ -import { AiPromptRole } from '@prisma/client'; +import { AiSessionMessageRole } from '@prisma/client'; import { z } from 'zod'; import { JSONSchema } from '../../../base'; @@ -7,7 +7,6 @@ import type { CapabilityModelCapability, ModelConditionsContract, } from '../../../native'; -import type { CopilotModelBackendKind } from '../runtime/contracts'; import { type StreamObject, StreamObjectSchema, @@ -97,7 +96,6 @@ export const PromptToolsSchema = z export const PromptConfigStrictSchema = z.object({ tools: PromptToolsSchema.nullable().optional(), - proModels: z.array(z.string()).nullable().optional(), // params requirements requireContent: z.boolean().nullable().optional(), requireAttachment: z.boolean().nullable().optional(), @@ -108,7 +106,7 @@ export const PromptConfigStrictSchema = z.object({ presencePenalty: z.number().nullable().optional(), temperature: z.number().nullable().optional(), topP: z.number().nullable().optional(), - maxTokens: z.number().nullable().optional(), + maxOutputTokens: z.number().nullable().optional(), // fal modelName: z.string().nullable().optional(), loras: z @@ -132,7 +130,7 @@ export type PromptTools = z.infer; export const EmbeddingMessage = z.array(z.string().trim().min(1)).min(1); -export const ChatMessageRole = Object.values(AiPromptRole) as [ +export const ChatMessageRole = Object.values(AiSessionMessageRole) as [ 'system', 'assistant', 'user', @@ -268,6 +266,8 @@ const CopilotProviderOptionsSchema = z.object({ billingUnitId: z.string().optional(), taskId: z.string().optional(), actionId: z.string().optional(), + builtInRouteId: z.string().optional(), + managedTargetId: z.string().optional(), quotaBackedRoutesAllowed: z.boolean().optional(), featureKind: z .enum([ @@ -380,8 +380,8 @@ export interface CopilotProviderModel { capabilities: ModelCapability[]; } -export type { CopilotModelBackendKind }; - -export type ModelConditions = Omit; +export type ModelConditions = Omit & { + profileId?: string; +}; export type ModelFullConditions = ModelConditionsContract; diff --git a/packages/backend/server/src/plugins/copilot/resolver.ts b/packages/backend/server/src/plugins/copilot/resolver.ts index 3d882c85c2..d5912ed57e 100644 --- a/packages/backend/server/src/plugins/copilot/resolver.ts +++ b/packages/backend/server/src/plugins/copilot/resolver.ts @@ -31,11 +31,12 @@ import { CurrentUser } from '../../core/auth'; import { DocAction, PermissionAccess } from '../../core/permission'; import { UserType } from '../../core/user'; import type { ListSessionOptions, UpdateChatSession } from '../../models'; +import { llmGetBuiltInRouteOptions } from '../../native'; +import { ByokEntitlementPolicy } from './byok'; import { CompatHistoryProjector } from './compat/history-projector'; import { ConversationInboxService } from './conversation/inbox'; -import { PromptService } from './prompt/service'; -import { CopilotProviderFactory } from './providers/factory'; -import { ModelOutputType, type StreamObject } from './providers/types'; +import { CopilotEnabled } from './feature'; +import type { StreamObject } from './providers/types'; import { ChatSessionService } from './session'; import { type ChatHistory, type ChatMessage, SubmittedMessage } from './types'; @@ -256,12 +257,6 @@ class CopilotHistoriesType implements Omit { @Field(() => String) promptName!: string; - @Field(() => String) - model!: string; - - @Field(() => [String]) - optionalModels!: string[]; - @Field(() => String, { description: 'An mark identifying which view to use to display the session', nullable: true, @@ -274,11 +269,6 @@ class CopilotHistoriesType implements Omit { @Field(() => String, { nullable: true }) title!: string | null; - @Field(() => Number, { - description: 'The number of tokens used in the session', - }) - tokens!: number; - @Field(() => [ChatMessageType]) messages!: ChatMessageType[]; @@ -303,27 +293,6 @@ class CopilotQuotaType { used!: number; } -@ObjectType() -class CopilotModelType { - @Field(() => String) - id!: string; - - @Field(() => String) - name!: string; -} - -@ObjectType() -export class CopilotModelsType { - @Field(() => String) - defaultModel!: string; - - @Field(() => [CopilotModelType]) - optionalModels!: CopilotModelType[]; - - @Field(() => [CopilotModelType]) - proModels!: CopilotModelType[]; -} - @ObjectType() export class CopilotSessionType { @Field(() => ID) @@ -343,12 +312,33 @@ export class CopilotSessionType { @Field(() => String) promptName!: string; +} + +@ObjectType('CopilotRouteTarget') +class CopilotRouteTargetType { + @Field(() => String) + id!: string; @Field(() => String) - model!: string; + displayName!: string; - @Field(() => [String]) - optionalModels!: string[]; + @Field(() => String) + minimumTier!: string; + + @Field(() => Boolean) + available!: boolean; +} + +@ObjectType('CopilotRouteOptions') +class CopilotRouteOptionsType { + @Field(() => String) + routeId!: string; + + @Field(() => String, { nullable: true }) + defaultTargetId!: string | null; + + @Field(() => [CopilotRouteTargetType]) + choices!: CopilotRouteTargetType[]; } // ================== Resolver ================== @@ -360,20 +350,47 @@ export class CopilotType { } @Throttle() +@CopilotEnabled() @Resolver(() => CopilotType) export class CopilotResolver { - private readonly modelNames = new Map(); - constructor( private readonly ac: PermissionAccess, private readonly mutex: RequestMutex, - private readonly prompt: PromptService, private readonly chatSession: ChatSessionService, private readonly historyProjector: CompatHistoryProjector, private readonly inbox: ConversationInboxService, - private readonly providerFactory: CopilotProviderFactory + private readonly entitlement: ByokEntitlementPolicy ) {} + @ResolveField(() => CopilotRouteOptionsType, { + nullable: true, + description: 'List native built-in route choices for a prompt', + complexity: 2, + }) + async routeOptions( + @CurrentUser() user: CurrentUser, + @Args('promptName') promptName: string + ): Promise { + const options = llmGetBuiltInRouteOptions(promptName); + if (!options) return null; + if (env.selfhosted) { + return { routeId: options.routeId, defaultTargetId: null, choices: [] }; + } + const premium = await this.entitlement.hasAiPlan(user.id); + return { + routeId: options.routeId, + defaultTargetId: premium + ? (options.premiumDefaultTargetId ?? null) + : (options.standardDefaultTargetId ?? null), + choices: options.choices.map(choice => ({ + id: choice.id, + displayName: choice.displayName, + minimumTier: choice.minimumTier, + available: premium || choice.minimumTier === 'Standard', + })), + }; + } + @ResolveField(() => CopilotQuotaType, { name: 'quota', description: 'Get the quota of the user in the workspace', @@ -408,51 +425,6 @@ export class CopilotResolver { return { userId: user.id, workspaceId, docId: docId || undefined }; } - @ResolveField(() => CopilotModelsType, { - description: - 'List available models for a prompt, with human-readable names', - complexity: 2, - }) - async models( - @Args('promptName') promptName: string - ): Promise { - const prompt = await this.prompt.get(promptName); - if (!prompt) { - throw new NotFoundException('Prompt not found'); - } - const convertModels = async (ids: string[]) => { - const models = await Promise.all( - ids.map(async id => { - const cachedName = this.modelNames.get(id); - if (cachedName) return { id, name: cachedName }; - - const resolved = await this.providerFactory.resolveProvider({ - modelId: id, - outputType: ModelOutputType.Text, - }); - const name = resolved?.provider.resolveModel( - resolved.modelId ?? id, - resolved.execution - )?.name; - if (name) { - this.modelNames.set(id, name); - return { id, name }; - } - return null; - }) - ); - - return models.filter(model => !!model) as CopilotModelType[]; - }; - const proModels = prompt.config?.proModels || []; - - return { - defaultModel: prompt.model, - optionalModels: await convertModels(prompt.optionalModels), - proModels: await convertModels(proModels), - }; - } - @ResolveField(() => CopilotSessionType, { description: 'Get the session by id', complexity: 2, @@ -813,6 +785,7 @@ export class CopilotResolver { } @Throttle() +@CopilotEnabled() @Resolver(() => UserType) export class UserCopilotResolver { constructor(private readonly ac: PermissionAccess) {} diff --git a/packages/backend/server/src/plugins/copilot/runtime/action-output-projector.ts b/packages/backend/server/src/plugins/copilot/runtime/action-output-projector.ts index 859382c602..7dd9e9db17 100644 --- a/packages/backend/server/src/plugins/copilot/runtime/action-output-projector.ts +++ b/packages/backend/server/src/plugins/copilot/runtime/action-output-projector.ts @@ -110,7 +110,7 @@ function isImageAction(actionId: string) { } function resolveProjector(actionId: string): ActionResultProjector | null { - if (actionId.startsWith('transcript.audio.')) { + if (actionId === 'transcript.audio') { return null; } if (isImageAction(actionId)) { diff --git a/packages/backend/server/src/plugins/copilot/runtime/action-runtime-bridge.ts b/packages/backend/server/src/plugins/copilot/runtime/action-runtime-bridge.ts index 62fc425d0f..2a29758d82 100644 --- a/packages/backend/server/src/plugins/copilot/runtime/action-runtime-bridge.ts +++ b/packages/backend/server/src/plugins/copilot/runtime/action-runtime-bridge.ts @@ -1,15 +1,10 @@ -import { Injectable, Optional } from '@nestjs/common'; +import { Injectable } from '@nestjs/common'; import { Models } from '../../../models'; import type { AiActionRunStatus } from '../../../models/copilot-action-run'; -import { - type NativeActionEvent, - type NativeActionRuntimeInput, - runNativeActionRecipePreparedStream, -} from '../../../native'; +import { type NativeActionEvent } from '../../../native'; import type { CopilotImageOptions, - CopilotProviderType, CopilotStructuredOptions, PromptMessage, } from '../providers/types'; @@ -18,18 +13,10 @@ import { projectActionResultToAssistantTurn, summarizeActionResult, } from './action-output-projector'; -import { - buildStructuredResponseFromSchemaJson, - type RequiredStructuredOutputContract, -} from './contracts'; -import { ExecutionPlanBuilder } from './execution-plan'; +import { CapabilityRuntime } from './capability-runtime'; +import { type RequiredStructuredOutputContract } from './contracts'; import { TurnPersistence } from './hosts/turn-persistence'; -type ActionRuntimeBridgeNativeInput = Omit< - NativeActionRuntimeInput, - 'recipeId' | 'recipeVersion' ->; - export type ActionRuntimeBridgeInput = { userId: string; workspaceId: string; @@ -42,26 +29,18 @@ export type ActionRuntimeBridgeInput = { attempt?: number; retryOf?: string | null; inputSnapshot?: unknown; - nativeInput?: ActionRuntimeBridgeNativeInput; onRunCreated?: ( context: ActionRuntimeBridgeRunContext ) => Promise | void; - prepareStructuredRoutes?: { - stepId?: string; + step: { + slot: string; + builtInRouteId: string; + profileId?: string; modelId?: string; messages: PromptMessage[]; - options?: CopilotStructuredOptions; - prefer?: CopilotProviderType; - responseSchemaJson?: Record; + options?: CopilotStructuredOptions | CopilotImageOptions; responseContract?: RequiredStructuredOutputContract; }; - prepareImageRoutes?: { - stepId?: string; - modelId?: string; - messages: PromptMessage[]; - options?: CopilotImageOptions; - prefer?: CopilotProviderType; - }; persistAttachment?: (attachment: unknown) => Promise | unknown; signal?: AbortSignal; }; @@ -107,94 +86,41 @@ export class ActionRuntimeBridge { constructor( private readonly models: Models, private readonly turnPersistence: TurnPersistence, - @Optional() private readonly plans?: ExecutionPlanBuilder + private readonly runtime: CapabilityRuntime ) {} - protected runNativeStream( - input: NativeActionRuntimeInput, - signal?: AbortSignal - ) { - return runNativeActionRecipePreparedStream(input, signal); - } - - private async prepareNativeInput( - input: ActionRuntimeBridgeInput - ): Promise { - const nativeInput = { - ...input.nativeInput, - input: input.nativeInput?.input ?? {}, - }; - const structured = input.prepareStructuredRoutes; - const image = input.prepareImageRoutes; - if (!structured && !image) { - return nativeInput; - } - if (!this.plans) { - throw new Error('Action route preparation is not available'); - } - const state = - nativeInput.input && typeof nativeInput.input === 'object' - ? { ...(nativeInput.input as Record) } - : {}; - - if (structured) { - const responseContract = - structured.responseContract ?? - (buildStructuredResponseFromSchemaJson( - structured.responseSchemaJson ?? { type: 'object' } - ) as RequiredStructuredOutputContract); - const plan = await this.plans.buildStructuredPlan( - { modelId: structured.modelId }, - structured.messages, - structured.options, - structured.prefer ? { prefer: structured.prefer } : undefined, - responseContract + private async execute(input: ActionRuntimeBridgeInput) { + const step = input.step; + if (step.responseContract) { + const output = await this.runtime.generateStructuredValue( + { profileId: step.profileId, modelId: step.modelId }, + step.messages, + { + ...(step.options as CopilotStructuredOptions | undefined), + builtInRouteId: step.builtInRouteId, + }, + step.responseContract, + undefined, + step.slot ); - const preparedRoutes = plan.nativeDispatch?.structured?.routes; - if (!preparedRoutes?.length) { - throw new Error('No native structured provider route prepared'); - } - - const existingPreparedRoutes = - state.preparedRoutes && - typeof state.preparedRoutes === 'object' && - !Array.isArray(state.preparedRoutes) - ? (state.preparedRoutes as Record) - : {}; - state.preparedRoutes = { - ...existingPreparedRoutes, - [structured.stepId ?? 'generate']: preparedRoutes, - }; + return { result: output.value, attachments: [] }; } - - if (image) { - const plan = await this.plans.buildImagePlan( - { modelId: image.modelId }, - image.messages, - image.options, - image.prefer ? { prefer: image.prefer } : undefined - ); - const preparedRoutes = plan.nativeDispatch?.image?.routes; - if (!preparedRoutes?.length) { - throw new Error('No native image provider route prepared'); - } - - const existingPreparedRoutes = - state.preparedRoutes && - typeof state.preparedRoutes === 'object' && - !Array.isArray(state.preparedRoutes) - ? (state.preparedRoutes as Record) - : {}; - state.preparedRoutes = { - ...existingPreparedRoutes, - [image.stepId ?? 'generate-image']: preparedRoutes, - }; + const images = []; + for await (const image of this.runtime.streamImageArtifacts( + { profileId: step.profileId, modelId: step.modelId }, + step.messages, + { + ...(step.options as CopilotImageOptions | undefined), + builtInRouteId: step.builtInRouteId, + }, + undefined, + step.slot + )) { + images.push(image); } - - return { - ...nativeInput, - input: state, - }; + const result = images[0]; + if (!result) throw new Error('Action image generation produced no image'); + return { result, attachments: [result] }; } private async projectAssistantResult( @@ -273,28 +199,36 @@ export class ActionRuntimeBridge { let finalEvent: NativeActionEvent | undefined; const attachments: unknown[] = []; try { - const nativeInput = await this.prepareNativeInput({ - ...inputWithBillingUnit, - }); - for await (const event of this.runNativeStream( - { - ...nativeInput, - recipeId: inputWithBillingUnit.actionId, - recipeVersion: inputWithBillingUnit.actionVersion, - }, - inputWithBillingUnit.signal - )) { - finalEvent = event; - let projectedEvent = event; - if (event.type === 'attachment') { - const attachment = input.persistAttachment - ? await input.persistAttachment(event.attachment) - : event.attachment; - attachments.push(attachment); - projectedEvent = { ...event, attachment }; - } - yield { ...projectedEvent, runId: run.id }; + const actionStart: NativeActionEvent = { + type: 'action_start', + actionId: input.actionId, + actionVersion: input.actionVersion, + status: 'running', + }; + yield { ...actionStart, runId: run.id }; + const output = await this.execute(inputWithBillingUnit); + for (const artifact of output.attachments) { + const attachment = input.persistAttachment + ? await input.persistAttachment(artifact) + : artifact; + attachments.push(attachment); + yield { + type: 'attachment', + actionId: input.actionId, + actionVersion: input.actionVersion, + status: 'running', + attachment, + runId: run.id, + }; } + finalEvent = { + type: 'action_done', + actionId: input.actionId, + actionVersion: input.actionVersion, + status: 'succeeded', + result: output.result, + }; + yield { ...finalEvent, runId: run.id }; } catch (error) { finalEvent = { type: 'error', @@ -351,33 +285,14 @@ export class ActionRuntimeBridge { ): ActionRuntimeBridgeInput { return { ...input, - prepareStructuredRoutes: input.prepareStructuredRoutes - ? { - ...input.prepareStructuredRoutes, - options: { - ...input.prepareStructuredRoutes.options, - actionId: - input.prepareStructuredRoutes.options?.actionId ?? - input.actionId, - billingUnitId: - input.prepareStructuredRoutes.options?.billingUnitId ?? - billingUnitId, - }, - } - : undefined, - prepareImageRoutes: input.prepareImageRoutes - ? { - ...input.prepareImageRoutes, - options: { - ...input.prepareImageRoutes.options, - actionId: - input.prepareImageRoutes.options?.actionId ?? input.actionId, - billingUnitId: - input.prepareImageRoutes.options?.billingUnitId ?? - billingUnitId, - }, - } - : undefined, + step: { + ...input.step, + options: { + ...input.step.options, + actionId: input.step.options?.actionId ?? input.actionId, + billingUnitId: input.step.options?.billingUnitId ?? billingUnitId, + }, + }, }; } } diff --git a/packages/backend/server/src/plugins/copilot/runtime/capability-runtime.ts b/packages/backend/server/src/plugins/copilot/runtime/capability-runtime.ts index 27e5043602..b9c2c60d77 100644 --- a/packages/backend/server/src/plugins/copilot/runtime/capability-runtime.ts +++ b/packages/backend/server/src/plugins/copilot/runtime/capability-runtime.ts @@ -1,7 +1,27 @@ +/* oxlint-disable import/no-cycle -- Tool callbacks can invoke nested Copilot prompts. */ import { Injectable } from '@nestjs/common'; -import { CopilotPromptInvalid } from '../../../base'; -import { ValidatedStructuredValueSchema } from '../core'; +import { Config } from '../../../base/config'; +import { CopilotPromptInvalid } from '../../../base/error/errors.gen'; +import { BackendRuntimeProvider } from '../../../core/backend-runtime'; +import { + buildLlmEmbeddingRequest, + buildLlmImageRequestFromMessages, + buildLlmRerankRequest, + type LlmImageResponse, + type LlmToolCallbackRequest, + type LlmToolLoopStreamEvent, + llmValidateJsonSchema, +} from '../../../native'; +import { + getByokSourceCoverage, + getCopilotFeatureAccess, +} from '../access/feature-coverage'; +import { assertCopilotEnabled } from '../availability'; +import { ByokEntitlementPolicy } from '../byok/policy'; +import type { ByokFeatureKind } from '../byok/types'; +import { ConversationPolicy } from '../conversation/policy'; +import { ValidatedStructuredValueSchema } from '../core/types'; import { type CopilotChatOptions, type CopilotEmbeddingOptions, @@ -9,112 +29,294 @@ import { type CopilotProviderType, type CopilotRerankRequest, type CopilotStructuredOptions, + type ModelAttachmentCapability, type ModelConditions, type PromptMessage, type StreamObject, } from '../providers/types'; import { + buildToolContracts, type RequiredStructuredOutputContract, requireStructuredOutputContract, } from './contracts'; import { - ExecutionPlanBuilder, - type ExecutionPlanForKind, -} from './execution-plan'; + type CopilotRuntimeEvent, + CopilotRuntimeEventConsumer, +} from './copilot-runtime-event-consumer'; +import { mapNativeSemanticError } from './native-errors'; import { - NativeExecutionEngine, - type NativeImageArtifact, -} from './native-execution-engine'; + buildCanonicalNativeRequest, + buildCanonicalNativeStructuredRequest, + preparePromptMessagesForNativeRequest, +} from './native-request-runtime'; +import { executeToolCall } from './tool/bridge'; +import { NativeProviderAdapter } from './tool/native-adapter'; +import { ToolRuntime } from './tool-runtime'; -type ProviderFilter = { - prefer?: CopilotProviderType; +type ProviderFilter = { prefer?: CopilotProviderType }; +type RuntimeOptions = NonNullable & { + dimensions?: number; + responseSchemaJson?: Record; + schemaHash?: string; + strict?: boolean; + profileId?: string; }; -const providerModelId = (modelId?: string) => modelId ?? 'auto'; +export type NativeImageArtifact = LlmImageResponse['images'][number]; + +const attachmentCapability = { + kinds: ['image', 'audio', 'file'], + sourceKinds: ['url', 'data', 'bytes', 'file_handle'], + allowRemoteUrls: true, +} satisfies ModelAttachmentCapability; @Injectable() export class CapabilityRuntime { constructor( - private readonly plans: ExecutionPlanBuilder, - private readonly engine: NativeExecutionEngine + private readonly backend: BackendRuntimeProvider, + private readonly entitlement: ByokEntitlementPolicy, + private readonly conversations: ConversationPolicy, + private readonly tools: ToolRuntime, + private readonly events: CopilotRuntimeEventConsumer, + private readonly config: Config ) {} - private async executePlan( - build: () => Promise, - execute: (plan: TPlan) => Promise - ) { - return await execute(await build()); + private async access(options: RuntimeOptions) { + assertCopilotEnabled(this.config); + const workspaceId = options.workspace; + const featureKind = (options.featureKind ?? 'chat') as ByokFeatureKind; + const coverage = getByokSourceCoverage(featureKind); + const [serverByok, localByok, premium] = workspaceId + ? await Promise.all([ + coverage.server && this.entitlement.hasServerEntitlement(workspaceId), + coverage.local && + this.entitlement.hasLocalEntitlement(workspaceId, options.user), + this.entitlement.hasAiPlan(options.user), + ]) + : [false, false, await this.entitlement.hasAiPlan(options.user)]; + const routeAllowed = + options.quotaBackedRoutesAllowed ?? + (!getCopilotFeatureAccess(featureKind).quotaMetered || + !options.user || + (await this.conversations.hasQuota(options.user))); + return { + routeAllowed, + managedTier: premium ? ('Premium' as const) : ('Standard' as const), + serverByok, + localByok, + }; } - private executeStreamPlan( - build: () => Promise, - execute: (plan: TPlan) => AsyncIterableIterator - ): AsyncIterableIterator { - return (async function* () { - yield* execute(await build()); - })(); + private eventContext(options: RuntimeOptions) { + return { + workspaceId: options.workspace, + userId: options.user, + sessionId: options.session, + taskId: options.taskId, + actionId: options.actionId, + billingUnitId: options.billingUnitId, + featureKind: (options.featureKind ?? 'chat') as ByokFeatureKind, + }; } - private hasNativeDispatch( - plan: ExecutionPlanForKind<'embedding'> | ExecutionPlanForKind<'rerank'>, - kind: 'embedding' | 'rerank' + private targetOverride(cond: ModelConditions) { + return cond.profileId && cond.modelId + ? { profileId: cond.profileId, modelId: cond.modelId } + : undefined; + } + + async assertRoute( + slot: string, + cond: ModelConditions, + options: CopilotChatOptions = {} ) { - return !!plan.nativeDispatch?.[kind]; + try { + await this.backend.assertCopilotRoute({ + slot, + builtInRouteId: options.builtInRouteId, + workspaceId: options.workspace, + userId: options.user, + localLeaseId: options.byokLeaseId, + access: await this.access(options), + managedTargetId: options.managedTargetId, + targetOverride: this.targetOverride(cond), + }); + } catch (error) { + throw mapNativeSemanticError(error); + } + } + + private async execute( + slot: string, + request: unknown, + cond: ModelConditions, + options: RuntimeOptions + ) { + try { + const output = await this.backend.executeCopilot({ + slot, + builtInRouteId: options.builtInRouteId, + workspaceId: options.workspace, + userId: options.user, + localLeaseId: options.byokLeaseId, + access: await this.access(options), + managedTargetId: options.managedTargetId, + targetOverride: this.targetOverride(cond), + request, + }); + await this.events.consume( + output.events as CopilotRuntimeEvent[], + this.eventContext(options) + ); + return output.result; + } catch (error) { + throw mapNativeSemanticError(error); + } + } + + private async prepareChat( + messages: PromptMessage[], + options: RuntimeOptions + ) { + const toolSet = await this.tools.getTools(options, ''); + const { request } = await buildCanonicalNativeRequest({ + model: 'route-selected', + messages, + options, + toolContracts: buildToolContracts(toolSet), + attachmentCapability, + include: options.reasoning ? ['reasoning'] : undefined, + reasoning: options.reasoning ? { effort: 'medium' } : undefined, + }); + return { request: { ...request, stream: true }, toolSet }; + } + + private async stream( + slot: string, + cond: ModelConditions, + messages: PromptMessage[], + options: RuntimeOptions + ) { + const { request, toolSet } = await this.prepareChat(messages, options); + const rawStream = this.backend.streamCopilot< + LlmToolLoopStreamEvent | CopilotRuntimeEvent + >( + { + slot, + builtInRouteId: options.builtInRouteId, + workspaceId: options.workspace, + userId: options.user, + localLeaseId: options.byokLeaseId, + access: await this.access(options), + managedTargetId: options.managedTargetId, + targetOverride: this.targetOverride(cond), + request, + }, + async requestJson => { + const toolRequest = JSON.parse(requestJson) as LlmToolCallbackRequest; + return JSON.stringify( + await executeToolCall(toolSet, toolRequest, { + signal: options.signal, + messages, + }) + ); + }, + { maxSteps: 20, signal: options.signal } + ); + const runtimeEvents = this.events; + const eventContext = this.eventContext(options); + async function* productEvents() { + for await (const event of rawStream) { + if ('route' in event) { + await runtimeEvents.consume([event], eventContext); + } else if (event.type === 'error') { + throw mapNativeSemanticError( + new Error( + typeof event.message === 'string' + ? event.message + : 'native runtime stream error' + ) + ); + } else { + yield event; + } + } + } + return { request, stream: productEvents() }; } async text( cond: ModelConditions, messages: PromptMessage[], - options?: CopilotChatOptions, - filter?: ProviderFilter + options: CopilotChatOptions = {}, + _filter?: ProviderFilter ) { - return await this.executePlan( - () => this.plans.buildTextPlan(cond, messages, options, filter), - plan => this.engine.execute(plan) + const prepared = await this.stream('prompt.text', cond, messages, options); + return await new NativeProviderAdapter(() => prepared.stream).text( + prepared.request, + options.signal, + messages ); } async *streamText( cond: ModelConditions, messages: PromptMessage[], - options?: CopilotChatOptions, - filter?: ProviderFilter + options: CopilotChatOptions = {}, + _filter?: ProviderFilter ): AsyncIterableIterator { - yield* this.executeStreamPlan( - () => this.plans.buildStreamTextPlan(cond, messages, options, filter), - plan => this.engine.executeStream(plan) + const prepared = await this.stream('chat.default', cond, messages, options); + yield* new NativeProviderAdapter(() => prepared.stream).streamText( + prepared.request, + options.signal, + messages ); } async *streamObject( cond: ModelConditions, messages: PromptMessage[], - options?: CopilotChatOptions, - filter?: ProviderFilter + options: CopilotChatOptions = {}, + _filter?: ProviderFilter ): AsyncIterableIterator { - yield* this.executeStreamPlan( - () => this.plans.buildStreamObjectPlan(cond, messages, options, filter), - plan => this.engine.executeStream(plan) + const prepared = await this.stream('chat.default', cond, messages, options); + yield* new NativeProviderAdapter(() => prepared.stream).streamObject( + prepared.request, + options.signal, + messages ); } async generateStructured( cond: ModelConditions, messages: PromptMessage[], - options?: CopilotStructuredOptions, - filter?: ProviderFilter, - responseContract?: RequiredStructuredOutputContract + options: CopilotStructuredOptions = {}, + _filter?: ProviderFilter, + responseContract?: RequiredStructuredOutputContract, + slot = 'prompt.structured' ) { - return await this.executePlan( - () => - this.plans.buildStructuredPlan( - cond, - messages, - options, - filter, - responseContract - ), - plan => this.engine.execute(plan) + const contract = requireStructuredOutputContract(responseContract); + if (!contract) { + throw new CopilotPromptInvalid('Structured schema contract is required'); + } + const { request } = await buildCanonicalNativeStructuredRequest({ + model: 'route-selected', + messages, + options, + responseContract: contract, + attachmentCapability, + }); + const result = (await this.execute(slot, request, cond, options)) as { + output_json?: unknown; + output_text: string; + }; + if (result.output_json === undefined) { + throw new CopilotPromptInvalid( + 'Structured response is missing output_json' + ); + } + return JSON.stringify( + llmValidateJsonSchema(request.schema, result.output_json) ); } @@ -123,87 +325,90 @@ export class CapabilityRuntime { messages: PromptMessage[], options: CopilotStructuredOptions, responseContract?: RequiredStructuredOutputContract, - filter?: ProviderFilter + filter?: ProviderFilter, + slot = 'prompt.structured' ) { - const validatedResponseContract = - requireStructuredOutputContract(responseContract); - if (!options || !validatedResponseContract) { + const contract = requireStructuredOutputContract(responseContract); + if (!contract) { throw new CopilotPromptInvalid('Structured schema contract is required'); } - - const output = await this.generateStructured( - cond, - messages, - options, - filter, - validatedResponseContract + const value = JSON.parse( + await this.generateStructured( + cond, + messages, + options, + filter, + contract, + slot + ) ); - const value = JSON.parse(output); return ValidatedStructuredValueSchema.parse({ value, - schemaHash: validatedResponseContract.schemaHash, + schemaHash: contract.schemaHash, schemaValidationVersion: 'json-schema-v1', - provider: filter?.prefer ?? 'auto', - model: providerModelId(cond.modelId), + provider: 'auto', + model: 'route-selected', }); } - async embeddingConfigured(modelId: string) { - try { - return this.hasNativeDispatch( - await this.plans.buildEmbeddingPlan(modelId, 'ping'), - 'embedding' - ); - } catch { - return false; - } + async embeddingConfigured(_modelId: string) { + return this.config.copilot.enabled; } async embed( - modelId: string, + _modelId: string, input: string | string[], - options?: CopilotEmbeddingOptions + options: CopilotEmbeddingOptions = {} ) { - return await this.executePlan( - () => this.plans.buildEmbeddingPlan(modelId, input, options), - plan => this.engine.execute(plan) - ); + const result = (await this.execute( + 'index.embedding', + buildLlmEmbeddingRequest({ + model: 'route-selected', + inputs: Array.isArray(input) ? input : [input], + dimensions: options.dimensions, + }), + {}, + options + )) as { embeddings: number[][] }; + return result.embeddings; } - async rerankConfigured(modelId: string) { - try { - return this.hasNativeDispatch( - await this.plans.buildRerankPlan(modelId, { - query: 'ping', - candidates: [{ text: 'ping' }], - }), - 'rerank' - ); - } catch { - return false; - } + async rerankConfigured(_modelId: string) { + return this.config.copilot.enabled; } async rerank( - modelId: string, + _modelId: string, request: CopilotRerankRequest, - options?: CopilotChatOptions + options: CopilotChatOptions = {} ) { - return await this.executePlan( - () => this.plans.buildRerankPlan(modelId, request, options), - plan => this.engine.execute(plan) - ); + const result = (await this.execute( + 'search.rerank', + buildLlmRerankRequest('route-selected', request), + {}, + options + )) as { scores: number[] }; + return result.scores; } async *streamImageArtifacts( cond: ModelConditions, messages: PromptMessage[], - options?: CopilotImageOptions, - filter?: ProviderFilter + options: CopilotImageOptions = {}, + _filter?: ProviderFilter, + slot = 'image.generate' ): AsyncIterableIterator { - yield* this.executeStreamPlan( - () => this.plans.buildImagePlan(cond, messages, options, filter), - plan => this.engine.executeImageArtifacts(plan) - ); + const { quality, seed } = options; + const result = (await this.execute( + slot, + buildLlmImageRequestFromMessages({ + model: 'route-selected', + messages: preparePromptMessagesForNativeRequest(messages, true), + options: { quality, seed }, + }), + cond, + options + )) as LlmImageResponse; + yield* result.images; } } diff --git a/packages/backend/server/src/plugins/copilot/runtime/contracts/execution-plan-contract.ts b/packages/backend/server/src/plugins/copilot/runtime/contracts/execution-plan-contract.ts deleted file mode 100644 index 641abee9fe..0000000000 --- a/packages/backend/server/src/plugins/copilot/runtime/contracts/execution-plan-contract.ts +++ /dev/null @@ -1,104 +0,0 @@ -import { - type LlmBackendConfig, - llmCompileExecutionPlan, - type LlmEmbeddingRequest, - type LlmImageRequest, - type LlmProtocol, - type LlmRequest, - type LlmRerankRequest, - type LlmStructuredRequest, -} from '../../../../native'; -import type { - CopilotProviderType, - ModelConditions, - PromptMessage, -} from '../../providers/types'; - -// Owner: runtime core mirror facade. -// The semantic source of truth is the native/Rust execution-plan contract -// behind llmCompileExecutionPlan(); this file only keeps the TypeScript shape -// needed by Node live-plan assembly until generated/native TS types replace it. -export type ExecutionRequestKind = - | 'text' - | 'streamText' - | 'streamObject' - | 'structured' - | 'embedding' - | 'rerank' - | 'image'; - -export type ExecutionRoute = { - providerId: string; - protocol: LlmProtocol; - model: string; - backendConfig: LlmBackendConfig; -}; - -export type ExecutionTransportContract = - | { kind: 'chat'; request: LlmRequest } - | { kind: 'structured'; request: LlmStructuredRequest } - | { kind: 'embedding'; request: LlmEmbeddingRequest } - | { kind: 'rerank'; request: LlmRerankRequest } - | { kind: 'image'; request: LlmImageRequest }; - -export type SerializableExecutionPlanRequest = - | { - kind: 'text' | 'streamText' | 'streamObject'; - cond: ModelConditions; - messages: PromptMessage[]; - options?: Record; - } - | { - kind: 'structured'; - cond: ModelConditions; - messages: PromptMessage[]; - options?: Record; - } - | { - kind: 'image'; - cond: ModelConditions; - messages: PromptMessage[]; - options?: Record; - } - | { - kind: 'embedding'; - cond: ModelConditions; - modelId: string; - input: string | string[]; - options?: Record; - } - | { - kind: 'rerank'; - cond: ModelConditions; - modelId: string; - request: { - query: string; - candidates: { id?: string; text: string }[]; - topK?: number; - }; - options?: Record; - }; - -export type SerializableExecutionPlan = { - routes: ExecutionRoute[]; - request: SerializableExecutionPlanRequest; - transport?: ExecutionTransportContract; - routePolicy: { fallbackOrder: string[] }; - runtimePolicy: { - prefer?: CopilotProviderType; - maxSteps?: number; - }; - attachmentPolicy: { - materializeRemoteAttachments: boolean; - }; - responsePostprocess: { - mode: ExecutionRequestKind; - }; - hostContext?: { - currentMessages?: PromptMessage[]; - }; -}; - -export function parseExecutionPlan(value: unknown) { - return llmCompileExecutionPlan(value); -} diff --git a/packages/backend/server/src/plugins/copilot/runtime/contracts/index.ts b/packages/backend/server/src/plugins/copilot/runtime/contracts/index.ts index 6429ae84ca..68b3e32982 100644 --- a/packages/backend/server/src/plugins/copilot/runtime/contracts/index.ts +++ b/packages/backend/server/src/plugins/copilot/runtime/contracts/index.ts @@ -1,5 +1,3 @@ -export * from './execution-plan-contract'; -export * from './native-contract'; export * from './prompt-contract'; export * from './runtime-event-contract'; export * from './shared'; diff --git a/packages/backend/server/src/plugins/copilot/runtime/contracts/native-contract.ts b/packages/backend/server/src/plugins/copilot/runtime/contracts/native-contract.ts deleted file mode 100644 index 336e07a7c4..0000000000 --- a/packages/backend/server/src/plugins/copilot/runtime/contracts/native-contract.ts +++ /dev/null @@ -1,97 +0,0 @@ -import serverNativeModule, { - type CapabilityMatchRequest, - type CapabilityMatchResponse, - type ModelRegistryMatchRequest, - type ModelRegistryMatchResponse, - type ModelRegistryResolveRequest, - type ModelRegistryResolveResponse, - type ModelRegistryVariantContract, - type ProviderDriverSpec, - type RequestedModelMatchRequest, - type RequestedModelMatchResponse, -} from '@affine/server-native'; - -// Owner: native/Rust contract facade. -// These types and validators intentionally proxy @affine/server-native and -// must not grow independent runtime semantics in Node. -export type { - CapabilityMatchRequest, - CapabilityMatchResponse, - ProviderDriverSpec, - RequestedModelMatchRequest, - RequestedModelMatchResponse, -}; - -export type CopilotModelBackendKind = ModelRegistryMatchRequest['backendKind']; -export type ModelRegistryVariant = ModelRegistryVariantContract; -export type ResolveModelRegistryVariantRequest = ModelRegistryResolveRequest; -export type ResolveModelRegistryVariantResponse = ModelRegistryResolveResponse; -export type MatchModelRegistryRequest = ModelRegistryMatchRequest; -export type MatchModelRegistryResponse = ModelRegistryMatchResponse; - -function validateNativeContract(name: string, value: unknown): T { - return serverNativeModule.llmValidateContract(name, value) as T; -} - -export function parseCapabilityMatchRequest(value: unknown) { - return validateNativeContract( - 'capabilityMatchRequest', - value - ); -} - -export function parseCapabilityMatchResponse(value: unknown) { - return validateNativeContract( - 'capabilityMatchResponse', - value - ); -} - -export function parseResolveModelRegistryVariantRequest(value: unknown) { - return validateNativeContract( - 'modelRegistryResolveRequest', - value - ); -} - -export function parseResolveModelRegistryVariantResponse(value: unknown) { - return validateNativeContract( - 'modelRegistryResolveResponse', - value - ); -} - -export function parseMatchModelRegistryRequest(value: unknown) { - return validateNativeContract( - 'modelRegistryMatchRequest', - value - ); -} - -export function parseMatchModelRegistryResponse(value: unknown) { - return validateNativeContract( - 'modelRegistryMatchResponse', - value - ); -} - -export function parseProviderDriverSpec(value: unknown) { - return validateNativeContract( - 'providerDriverSpec', - value - ); -} - -export function parseRequestedModelMatchRequest(value: unknown) { - return validateNativeContract( - 'requestedModelMatchRequest', - value - ); -} - -export function parseRequestedModelMatchResponse(value: unknown) { - return validateNativeContract( - 'requestedModelMatchResponse', - value - ); -} diff --git a/packages/backend/server/src/plugins/copilot/runtime/contracts/prompt-contract.ts b/packages/backend/server/src/plugins/copilot/runtime/contracts/prompt-contract.ts index 7f3c05b7e7..ca98f4047d 100644 --- a/packages/backend/server/src/plugins/copilot/runtime/contracts/prompt-contract.ts +++ b/packages/backend/server/src/plugins/copilot/runtime/contracts/prompt-contract.ts @@ -1,15 +1,6 @@ -import { - llmValidateContract, - type NativePromptCountTokensRequest, - type NativePromptCountTokensResponse, - type NativePromptMetadataRequest, - type NativePromptMetadataResponse, - type NativePromptRenderRequest, - type NativePromptRenderResponse, - type NativePromptSessionRenderRequest, - type NativePromptSessionRenderResponse, - type PromptMessageContract as NativePromptMessageContract, - type PromptStructuredResponseContract as NativePromptStructuredResponseContract, +import type { + PromptMessageContract as NativePromptMessageContract, + PromptStructuredResponseContract as NativePromptStructuredResponseContract, } from '../../../../native'; import { normalizePromptResponseFormat } from './structured-output-contract'; @@ -32,14 +23,6 @@ type PromptMessageInput = { params?: Record | null; responseFormat?: PromptResponseFormat | null; }; -export type PromptRenderContract = NativePromptRenderRequest; -export type PromptRenderResult = NativePromptRenderResponse; -export type PromptTokenCountContract = NativePromptCountTokensRequest; -export type PromptTokenCountResult = NativePromptCountTokensResponse; -export type PromptMetadataContract = NativePromptMetadataRequest; -export type PromptMetadataResult = NativePromptMetadataResponse; -export type PromptSessionContract = NativePromptSessionRenderRequest; -export type PromptSessionResult = NativePromptSessionRenderResponse; export type NativePromptResponseFormatProjection = { nativeResponseFormat?: PromptStructuredResponseContract; }; @@ -82,17 +65,3 @@ export function projectPromptMessageForNative( return { message: nativeMessage, nativeResponseFormat }; } - -export function parsePromptRenderContract(value: unknown) { - return llmValidateContract( - 'promptRenderContract', - value - ); -} - -export function parsePromptSessionContract(value: unknown) { - return llmValidateContract( - 'promptSessionContract', - value - ); -} diff --git a/packages/backend/server/src/plugins/copilot/runtime/copilot-runtime-event-consumer.ts b/packages/backend/server/src/plugins/copilot/runtime/copilot-runtime-event-consumer.ts new file mode 100644 index 0000000000..1999b108ae --- /dev/null +++ b/packages/backend/server/src/plugins/copilot/runtime/copilot-runtime-event-consumer.ts @@ -0,0 +1,129 @@ +import { Injectable, Logger } from '@nestjs/common'; + +import { metrics } from '../../../base'; +import { Models } from '../../../models'; +import { type ByokFeatureKind, ByokProviderSource } from '../byok/types'; + +export type CopilotRuntimeRouteIdentity = { + profileId: string; + source: 'server' | 'local' | 'affine_cloud'; + provider: string; + model: string; +}; + +export type CopilotRuntimeEvent = + | { type: 'route_selected'; route: CopilotRuntimeRouteIdentity } + | { + type: 'route_failed'; + route: CopilotRuntimeRouteIdentity; + errorKind: string; + } + | { + type: 'usage'; + route: CopilotRuntimeRouteIdentity; + usage: { + prompt_tokens?: number; + completion_tokens?: number; + total_tokens?: number; + cached_tokens?: number; + input_tokens?: number; + output_tokens?: number; + }; + }; + +export type CopilotRuntimeEventContext = { + workspaceId?: string; + userId?: string; + sessionId?: string; + taskId?: string; + actionId?: string; + billingUnitId?: string; + featureKind: ByokFeatureKind; +}; + +@Injectable() +export class CopilotRuntimeEventConsumer { + private readonly logger = new Logger(CopilotRuntimeEventConsumer.name); + + constructor(private readonly models: Models) {} + + async consume( + events: CopilotRuntimeEvent[], + context: CopilotRuntimeEventContext + ) { + for (const event of events) { + try { + if (event.type === 'usage') { + await this.recordUsage(event, context); + } else if (event.type === 'route_failed') { + await this.recordFailure(event, context); + } + } catch (error) { + this.logger.warn( + `Failed to consume copilot runtime event: ${ + error instanceof Error ? error.message : String(error) + }` + ); + } + } + } + + private async recordUsage( + event: Extract, + context: CopilotRuntimeEventContext + ) { + if (!context.workspaceId || event.route.source === 'affine_cloud') { + return; + } + const usage = event.usage; + metrics.ai.counter('byok_usage').add(1, { + provider: event.route.provider, + source: event.route.source, + feature: context.featureKind, + }); + await this.models.copilotUsage.create({ + workspaceId: context.workspaceId, + userId: context.userId, + provider: event.route.provider, + providerSource: + event.route.source === 'server' + ? ByokProviderSource.Server + : ByokProviderSource.Local, + featureKind: context.featureKind, + model: event.route.model, + sessionId: context.sessionId, + taskId: context.taskId, + actionId: context.actionId, + billingUnitId: context.billingUnitId, + promptTokens: usage.prompt_tokens ?? usage.input_tokens ?? 0, + completionTokens: usage.completion_tokens ?? usage.output_tokens ?? 0, + totalTokens: usage.total_tokens ?? 0, + cachedTokens: usage.cached_tokens ?? 0, + }); + if (event.route.source === 'server') { + await this.models.copilotWorkspaceByokConfig.touchUsed( + context.workspaceId, + event.route.profileId + ); + } + } + + private async recordFailure( + event: Extract, + context: CopilotRuntimeEventContext + ) { + metrics.ai.counter('byok_route_failure').add(1, { + provider: event.route.provider, + source: event.route.source, + feature: context.featureKind, + reason: event.errorKind, + }); + if (context.workspaceId && event.route.source === 'server') { + await this.models.copilotWorkspaceByokConfig.markFailure( + context.workspaceId, + event.route.profileId, + event.errorKind + ); + } + } +} diff --git a/packages/backend/server/src/plugins/copilot/runtime/execution-metrics.ts b/packages/backend/server/src/plugins/copilot/runtime/execution-metrics.ts deleted file mode 100644 index df5595ee25..0000000000 --- a/packages/backend/server/src/plugins/copilot/runtime/execution-metrics.ts +++ /dev/null @@ -1,65 +0,0 @@ -import { Injectable } from '@nestjs/common'; - -import { metrics } from '../../../base'; -import type { ResolvedCopilotProvider } from '../providers/factory'; -import type { CopilotProviderType } from '../providers/types'; -import type { ExecutionRequestKind } from './execution-plan'; - -type ExecutionDispatchPath = 'prepared_routes'; - -export function summarizePreparedRoutes( - routes: Array> -) { - const preparedCount = routes.filter(route => !!route.prepared).length; - return { - routeCount: routes.length, - preparedCount, - preparedMode: - preparedCount === 0 - ? 'none' - : preparedCount === routes.length - ? 'all' - : 'partial', - } as const; -} - -function planAttrs( - kind: ExecutionRequestKind, - prefer?: CopilotProviderType, - routes?: ResolvedCopilotProvider[] -) { - const summary = summarizePreparedRoutes(routes ?? []); - return { - kind, - prefer: prefer ?? 'auto', - prepared: summary.preparedMode, - route_count: summary.routeCount, - }; -} - -@Injectable() -export class CopilotExecutionMetrics { - recordPlan( - kind: ExecutionRequestKind, - routes: ResolvedCopilotProvider[], - prefer?: CopilotProviderType - ) { - const attrs = planAttrs(kind, prefer, routes); - metrics.ai.counter('execution_plan_total').add(1, attrs); - metrics.ai.histogram('execution_plan_routes').record(attrs.route_count, { - kind: attrs.kind, - prefer: attrs.prefer, - prepared: attrs.prepared, - }); - } - - recordDispatch( - kind: ExecutionRequestKind, - path: ExecutionDispatchPath, - routeCount: number - ) { - const attrs = { kind, path }; - metrics.ai.counter('execution_dispatch_total').add(1, attrs); - metrics.ai.histogram('execution_dispatch_routes').record(routeCount, attrs); - } -} diff --git a/packages/backend/server/src/plugins/copilot/runtime/execution-plan.ts b/packages/backend/server/src/plugins/copilot/runtime/execution-plan.ts deleted file mode 100644 index cf97306808..0000000000 --- a/packages/backend/server/src/plugins/copilot/runtime/execution-plan.ts +++ /dev/null @@ -1,827 +0,0 @@ -import { Injectable } from '@nestjs/common'; - -import type { - LlmPreparedDispatchRoute, - LlmPreparedEmbeddingDispatchRoute, - LlmPreparedImageDispatchRoute, - LlmPreparedRerankDispatchRoute, - LlmPreparedStructuredDispatchRoute, -} from '../../../native'; -import { llmNormalizePreparedRoutes } from '../../../native'; -import { - CopilotProviderFactory, - type ResolvedCopilotProvider, -} from '../providers/factory'; -import type { - PreparedNativeEmbeddingExecution, - PreparedNativeExecution, - PreparedNativeImageExecution, - PreparedNativeRerankExecution, - PreparedNativeStructuredExecution, -} from '../providers/provider-runtime-contract'; -import type { - CopilotChatOptions, - CopilotEmbeddingOptions, - CopilotImageOptions, - CopilotProviderType, - CopilotRerankRequest, - CopilotStructuredOptions, - ModelConditions, - PromptMessage, -} from '../providers/types'; -import { ModelOutputType } from '../providers/types'; -import type { RequiredStructuredOutputContract } from './contracts'; -import { - type ExecutionRequestKind, - type ExecutionRoute, - type ExecutionTransportContract, - parseExecutionPlan, - type SerializableExecutionPlan, - type SerializableExecutionPlanRequest, -} from './contracts/execution-plan-contract'; -import { CopilotExecutionMetrics } from './execution-metrics'; - -export type { ExecutionRequestKind }; - -type ProviderFilter = { - prefer?: CopilotProviderType; -}; - -type BaseExecutionRequest = { - kind: TKind; - cond: ModelConditions; -}; - -type TextExecutionRequest = BaseExecutionRequest<'text'> & { - messages: PromptMessage[]; - options?: CopilotChatOptions; -}; - -type StreamTextExecutionRequest = BaseExecutionRequest<'streamText'> & { - messages: PromptMessage[]; - options?: CopilotChatOptions; -}; - -type StreamObjectExecutionRequest = BaseExecutionRequest<'streamObject'> & { - messages: PromptMessage[]; - options?: CopilotChatOptions; -}; - -type StructuredExecutionRequest = BaseExecutionRequest<'structured'> & { - messages: PromptMessage[]; - options?: CopilotStructuredOptions; -}; - -type ImageExecutionRequest = BaseExecutionRequest<'image'> & { - messages: PromptMessage[]; - options?: CopilotImageOptions; -}; - -type EmbeddingExecutionRequest = BaseExecutionRequest<'embedding'> & { - modelId: string; - input: string | string[]; - options?: CopilotEmbeddingOptions; -}; - -type RerankExecutionRequest = BaseExecutionRequest<'rerank'> & { - modelId: string; - request: CopilotRerankRequest; - options?: CopilotChatOptions; -}; - -export type ExecutionPlanRequest = - | TextExecutionRequest - | StreamTextExecutionRequest - | StreamObjectExecutionRequest - | StructuredExecutionRequest - | ImageExecutionRequest - | EmbeddingExecutionRequest - | RerankExecutionRequest; - -export type ExecutionPlanForKind = - ExecutionPlan & { - request: Extract; - }; - -type NativePreparedDispatchPlan = { - routes: TRoute[]; - prepared: TPrepared; -}; - -export type NativeChatDispatchPlan = NativePreparedDispatchPlan< - LlmPreparedDispatchRoute, - PreparedNativeExecution -> & { - hasTools: boolean; -}; - -export type NativeStructuredDispatchPlan = NativePreparedDispatchPlan< - LlmPreparedStructuredDispatchRoute, - PreparedNativeStructuredExecution ->; - -export type NativeEmbeddingDispatchPlan = NativePreparedDispatchPlan< - LlmPreparedEmbeddingDispatchRoute, - PreparedNativeEmbeddingExecution ->; - -export type NativeRerankDispatchPlan = NativePreparedDispatchPlan< - LlmPreparedRerankDispatchRoute, - PreparedNativeRerankExecution ->; - -export type NativeImageDispatchPlan = NativePreparedDispatchPlan< - LlmPreparedImageDispatchRoute, - PreparedNativeImageExecution ->; - -export type ExecutionPlan = { - nativeDispatch?: { - chat?: NativeChatDispatchPlan; - structured?: NativeStructuredDispatchPlan; - embedding?: NativeEmbeddingDispatchPlan; - rerank?: NativeRerankDispatchPlan; - image?: NativeImageDispatchPlan; - }; - serializable?: SerializableExecutionPlan; - transport?: ExecutionTransportContract; - request: ExecutionPlanRequest; - routePolicy: { fallbackOrder: string[] }; - runtimePolicy: { - prefer?: CopilotProviderType; - }; - attachmentPolicy: { - materializeRemoteAttachments: boolean; - }; - responsePostprocess: { mode: ExecutionRequestKind }; - hostPersistence: { - persistAssistantTurn: boolean; - outputKind: ExecutionRequestKind; - }; - hostContext: { - signal?: AbortSignal; - currentMessages?: PromptMessage[]; - }; -}; - -type PreparedRouteLike = { - route: { - providerId: string; - protocol: PreparedNativeExecution['route']['protocol']; - model: string; - backendConfig: PreparedNativeExecution['route']['backendConfig']; - }; - request: TRequest; -}; - -function buildPreparedTransport< - TKind extends ExecutionTransportContract['kind'], - TPrepared extends PreparedRouteLike, ->( - kind: TKind, - routes: ResolvedCopilotProvider[], - getPrepared: (route: ResolvedCopilotProvider) => TPrepared | undefined -): ExecutionTransportContract | undefined { - const prepared = - routes.length === 1 ? routes[0] && getPrepared(routes[0]) : undefined; - if (!prepared) { - return; - } - - return { - kind, - request: prepared.request, - } as ExecutionTransportContract; -} - -function collectPreparedRoutes( - routes: ResolvedCopilotProvider[], - getPrepared: (route: ResolvedCopilotProvider) => TPrepared | undefined, - mapPreparedRoute: (prepared: TPrepared) => TRoute -): TRoute[] | undefined { - if (!routes.length) { - return; - } - - const preparedRoutes: TRoute[] = []; - for (const route of routes) { - const prepared = getPrepared(route); - if (!prepared) { - return; - } - preparedRoutes.push(mapPreparedRoute(prepared)); - } - - return preparedRoutes; -} - -function buildPreparedDispatchPlan< - TPrepared extends PreparedRouteLike, - TRoute, - TDispatch extends NativePreparedDispatchPlan, ->( - routes: ResolvedCopilotProvider[], - getPrepared: (route: ResolvedCopilotProvider) => TPrepared | undefined, - mapPreparedRoute: (prepared: TPrepared) => TRoute, - buildPreparedDispatchResult?: ( - preparedRoutes: TRoute[], - prepared: TPrepared - ) => TDispatch -): TDispatch | undefined { - const preparedRoutes = collectPreparedRoutes( - routes, - getPrepared, - mapPreparedRoute - ); - const prepared = routes[0] && getPrepared(routes[0]); - if (!preparedRoutes || !prepared) { - return; - } - - const normalizedRoutes = llmNormalizePreparedRoutes(preparedRoutes); - - return buildPreparedDispatchResult - ? buildPreparedDispatchResult(normalizedRoutes, prepared) - : ({ routes: normalizedRoutes, prepared } as TDispatch); -} - -type DispatchPreparedRoute = { - provider_id: string; - protocol: PreparedNativeExecution['route']['protocol']; - model: string; - config: PreparedNativeExecution['route']['backendConfig']; - request: TRequest; -}; - -function mapPreparedDispatchRoute( - prepared: PreparedRouteLike -): DispatchPreparedRoute { - return { - provider_id: prepared.route.providerId, - protocol: prepared.route.protocol, - model: prepared.route.model, - config: prepared.route.backendConfig, - request: prepared.request, - }; -} - -type PreparedExecutionArtifactSpec< - TKind extends ExecutionTransportContract['kind'], - TPrepared extends PreparedRouteLike, - TRoute, - TDispatch extends NativePreparedDispatchPlan, -> = { - transportKind: TKind; - getPrepared: (route: ResolvedCopilotProvider) => TPrepared | undefined; - mapPreparedRoute: (prepared: TPrepared) => TRoute; - buildPreparedDispatch?: ( - preparedRoutes: TRoute[], - prepared: TPrepared - ) => TDispatch; -}; - -type PreparedExecutionArtifacts = { - dispatch?: TDispatch; - transport?: ExecutionTransportContract; -}; - -function buildPreparedExecutionArtifacts< - TKind extends ExecutionTransportContract['kind'], - TPrepared extends PreparedRouteLike, - TRoute, - TDispatch extends NativePreparedDispatchPlan, ->( - routes: ResolvedCopilotProvider[], - spec: PreparedExecutionArtifactSpec -): PreparedExecutionArtifacts { - return { - dispatch: buildPreparedDispatchPlan( - routes, - spec.getPrepared, - spec.mapPreparedRoute, - spec.buildPreparedDispatch - ), - transport: buildPreparedTransport( - spec.transportKind, - routes, - spec.getPrepared - ), - }; -} - -const chatArtifactSpec: PreparedExecutionArtifactSpec< - 'chat', - PreparedNativeExecution, - LlmPreparedDispatchRoute, - NativeChatDispatchPlan -> = { - transportKind: 'chat', - getPrepared: route => route.prepared, - mapPreparedRoute: mapPreparedDispatchRoute, - buildPreparedDispatch: (preparedRoutes, prepared) => ({ - routes: preparedRoutes, - prepared, - hasTools: Object.keys(prepared.tools).length > 0, - }), -}; - -const structuredArtifactSpec: PreparedExecutionArtifactSpec< - 'structured', - PreparedNativeStructuredExecution, - LlmPreparedStructuredDispatchRoute, - NativeStructuredDispatchPlan -> = { - transportKind: 'structured', - getPrepared: route => route.preparedStructured, - mapPreparedRoute: mapPreparedDispatchRoute, -}; - -const embeddingArtifactSpec: PreparedExecutionArtifactSpec< - 'embedding', - PreparedNativeEmbeddingExecution, - LlmPreparedEmbeddingDispatchRoute, - NativeEmbeddingDispatchPlan -> = { - transportKind: 'embedding', - getPrepared: route => route.preparedEmbedding, - mapPreparedRoute: mapPreparedDispatchRoute, -}; - -const rerankArtifactSpec: PreparedExecutionArtifactSpec< - 'rerank', - PreparedNativeRerankExecution, - LlmPreparedRerankDispatchRoute, - NativeRerankDispatchPlan -> = { - transportKind: 'rerank', - getPrepared: route => route.preparedRerank, - mapPreparedRoute: mapPreparedDispatchRoute, -}; - -const imageArtifactSpec: PreparedExecutionArtifactSpec< - 'image', - PreparedNativeImageExecution, - LlmPreparedImageDispatchRoute, - NativeImageDispatchPlan -> = { - transportKind: 'image', - getPrepared: route => route.preparedImage, - mapPreparedRoute: mapPreparedDispatchRoute, -}; - -function buildFallbackOrder(routes: ResolvedCopilotProvider[]) { - return routes.map(route => route.providerId); -} - -function mapExecutionRoute(route: ResolvedCopilotProvider): ExecutionRoute { - const preparedRoute = - route.prepared?.route ?? - route.preparedStructured?.route ?? - route.preparedEmbedding?.route ?? - route.preparedRerank?.route ?? - route.preparedImage?.route; - - if (preparedRoute) { - return { - providerId: preparedRoute.providerId, - protocol: preparedRoute.protocol, - model: preparedRoute.model, - backendConfig: preparedRoute.backendConfig, - }; - } - - const rawRoute = route as unknown as ExecutionRoute; - return { - providerId: rawRoute.providerId, - protocol: rawRoute.protocol, - model: rawRoute.model, - backendConfig: rawRoute.backendConfig, - }; -} - -function stripHostOnlyOptions( - options: TOptions -): Record | undefined { - if (!options) { - return; - } - - const { - signal: _signal, - user: _user, - session: _session, - workspace: _workspace, - quotaBackedRoutesAllowed: _quotaBackedRoutesAllowed, - ...serializable - } = options as Record; - - return Object.keys(serializable).length ? serializable : undefined; -} - -function buildSerializableRequest( - request: ExecutionPlanRequest -): SerializableExecutionPlanRequest { - switch (request.kind) { - case 'text': - case 'streamText': - case 'streamObject': - case 'structured': - case 'image': - return { - ...request, - options: stripHostOnlyOptions(request.options), - } as SerializableExecutionPlanRequest; - case 'embedding': - case 'rerank': - return { - ...request, - options: stripHostOnlyOptions(request.options), - }; - } -} - -function buildSerializableExecutionPlan( - routes: ResolvedCopilotProvider[], - input: Omit< - ExecutionPlan, - 'nativeDispatch' | 'serializable' | 'hostContext' - > & - Pick -): SerializableExecutionPlan { - return parseExecutionPlan({ - routes: routes.map(mapExecutionRoute), - request: buildSerializableRequest(input.request), - transport: input.transport, - routePolicy: input.routePolicy, - runtimePolicy: input.runtimePolicy, - attachmentPolicy: input.attachmentPolicy, - responsePostprocess: input.responsePostprocess, - hostContext: input.hostContext.currentMessages - ? { currentMessages: input.hostContext.currentMessages } - : undefined, - }); -} - -type MessagePlanArtifacts = Pick; - -function buildMessagePlanArtifacts( - kind: Extract< - ExecutionRequestKind, - 'text' | 'streamText' | 'streamObject' | 'structured' | 'image' - >, - routes: ResolvedCopilotProvider[] -): MessagePlanArtifacts { - const chatArtifacts = - kind === 'text' || kind === 'streamText' || kind === 'streamObject' - ? buildPreparedExecutionArtifacts(routes, chatArtifactSpec) - : undefined; - const structuredArtifacts = - kind === 'structured' - ? buildPreparedExecutionArtifacts(routes, structuredArtifactSpec) - : undefined; - const imageArtifacts = - kind === 'image' - ? buildPreparedExecutionArtifacts(routes, imageArtifactSpec) - : undefined; - const nativeDispatch = { - chat: - kind === 'text' || kind === 'streamText' || kind === 'streamObject' - ? chatArtifacts?.dispatch - : undefined, - structured: - kind === 'structured' ? structuredArtifacts?.dispatch : undefined, - image: kind === 'image' ? imageArtifacts?.dispatch : undefined, - }; - - return { - nativeDispatch, - transport: - kind === 'text' || kind === 'streamText' || kind === 'streamObject' - ? chatArtifacts?.transport - : kind === 'structured' - ? structuredArtifacts?.transport - : kind === 'image' - ? imageArtifacts?.transport - : undefined, - }; -} - -function buildEmbeddingPlanArtifacts( - routes: ResolvedCopilotProvider[] -): Pick { - const embeddingArtifacts = buildPreparedExecutionArtifacts( - routes, - embeddingArtifactSpec - ); - return { - nativeDispatch: { - embedding: embeddingArtifacts.dispatch, - }, - transport: embeddingArtifacts.transport, - }; -} - -function buildRerankPlanArtifacts( - routes: ResolvedCopilotProvider[] -): Pick { - const rerankArtifacts = buildPreparedExecutionArtifacts( - routes, - rerankArtifactSpec - ); - return { - nativeDispatch: { - rerank: rerankArtifacts.dispatch, - }, - transport: rerankArtifacts.transport, - }; -} - -@Injectable() -export class ExecutionPlanBuilder { - constructor( - private readonly providers: CopilotProviderFactory, - private readonly executionMetrics: CopilotExecutionMetrics - ) {} - - private async buildMessagePlan< - TKind extends Extract< - ExecutionRequestKind, - 'text' | 'streamText' | 'streamObject' | 'structured' | 'image' - >, - >( - kind: TKind, - cond: ModelConditions, - messages: PromptMessage[], - options?: - | CopilotChatOptions - | CopilotStructuredOptions - | CopilotImageOptions, - filter: ProviderFilter = {} - ): Promise> { - const outputType = - kind === 'image' - ? ModelOutputType.Image - : kind === 'streamObject' - ? ModelOutputType.Object - : kind === 'structured' - ? ModelOutputType.Structured - : ModelOutputType.Text; - - const routes = - kind === 'text' || kind === 'streamText' || kind === 'streamObject' - ? await this.providers.prepareRoutes( - kind, - { ...cond, outputType }, - messages, - (options as CopilotChatOptions | undefined) ?? {}, - filter - ) - : kind === 'structured' - ? await this.providers.prepareStructuredRoutes( - { ...cond, outputType }, - messages, - (options as CopilotStructuredOptions | undefined) ?? {}, - filter - ) - : await this.providers.prepareImageRoutes( - { ...cond, outputType }, - messages, - (options as CopilotImageOptions | undefined) ?? {}, - filter - ); - this.executionMetrics.recordPlan(kind, routes, filter.prefer); - const { nativeDispatch, transport } = buildMessagePlanArtifacts( - kind, - routes - ); - const plan = { - transport, - request: { - kind, - cond: { ...cond, modelId: cond.modelId }, - messages, - options, - } as Extract, - routePolicy: { - fallbackOrder: buildFallbackOrder(routes), - }, - runtimePolicy: { prefer: filter.prefer }, - attachmentPolicy: { materializeRemoteAttachments: true }, - responsePostprocess: { mode: kind }, - hostPersistence: { - persistAssistantTurn: true, - outputKind: kind, - }, - hostContext: { - signal: options?.signal, - currentMessages: messages, - }, - } as Omit, 'nativeDispatch' | 'serializable'>; - - return { - nativeDispatch, - serializable: buildSerializableExecutionPlan(routes, plan), - ...plan, - }; - } - - async buildTextPlan( - cond: ModelConditions, - messages: PromptMessage[], - options?: CopilotChatOptions, - filter?: ProviderFilter - ): Promise> { - return await this.buildMessagePlan('text', cond, messages, options, filter); - } - - async buildStreamTextPlan( - cond: ModelConditions, - messages: PromptMessage[], - options?: CopilotChatOptions, - filter?: ProviderFilter - ): Promise> { - return await this.buildMessagePlan( - 'streamText', - cond, - messages, - options, - filter - ); - } - - async buildStreamObjectPlan( - cond: ModelConditions, - messages: PromptMessage[], - options?: CopilotChatOptions, - filter?: ProviderFilter - ): Promise> { - return await this.buildMessagePlan( - 'streamObject', - cond, - messages, - options, - filter - ); - } - - async buildStructuredPlan( - cond: ModelConditions, - messages: PromptMessage[], - options?: CopilotStructuredOptions, - filter?: ProviderFilter, - responseContract?: RequiredStructuredOutputContract - ): Promise> { - const outputType = ModelOutputType.Structured; - const routes = await this.providers.prepareStructuredRoutes( - { ...cond, outputType }, - messages, - options ?? {}, - filter ?? {}, - responseContract - ); - this.executionMetrics.recordPlan('structured', routes, filter?.prefer); - const { nativeDispatch, transport } = buildMessagePlanArtifacts( - 'structured', - routes - ); - const plan = { - transport, - request: { - kind: 'structured', - cond: { ...cond, modelId: cond.modelId }, - messages, - options, - }, - routePolicy: { - fallbackOrder: buildFallbackOrder(routes), - }, - runtimePolicy: { prefer: filter?.prefer }, - attachmentPolicy: { materializeRemoteAttachments: true }, - responsePostprocess: { mode: 'structured' }, - hostPersistence: { - persistAssistantTurn: true, - outputKind: 'structured', - }, - hostContext: { - signal: options?.signal, - currentMessages: messages, - }, - } as Omit< - ExecutionPlanForKind<'structured'>, - 'nativeDispatch' | 'serializable' - >; - - return { - nativeDispatch, - serializable: buildSerializableExecutionPlan(routes, plan), - ...plan, - }; - } - - async buildImagePlan( - cond: ModelConditions, - messages: PromptMessage[], - options?: CopilotImageOptions, - filter?: ProviderFilter - ): Promise> { - return await this.buildMessagePlan( - 'image', - cond, - messages, - options, - filter - ); - } - - async buildEmbeddingPlan( - modelId: string, - input: string | string[], - options?: CopilotEmbeddingOptions - ): Promise> { - const routes = await this.providers.prepareEmbeddingRoutes( - modelId, - input, - options - ); - this.executionMetrics.recordPlan('embedding', routes); - const { nativeDispatch, transport } = buildEmbeddingPlanArtifacts(routes); - const plan = { - transport, - request: { - kind: 'embedding', - cond: { modelId }, - modelId, - input, - options, - }, - routePolicy: { - fallbackOrder: buildFallbackOrder(routes), - }, - runtimePolicy: {}, - attachmentPolicy: { materializeRemoteAttachments: false }, - responsePostprocess: { mode: 'embedding' }, - hostPersistence: { - persistAssistantTurn: false, - outputKind: 'embedding', - }, - hostContext: { - signal: options?.signal, - }, - } as Omit< - ExecutionPlanForKind<'embedding'>, - 'nativeDispatch' | 'serializable' - >; - - return { - nativeDispatch, - serializable: buildSerializableExecutionPlan(routes, plan), - ...plan, - }; - } - - async buildRerankPlan( - modelId: string, - request: CopilotRerankRequest, - options?: CopilotChatOptions - ): Promise> { - const routes = await this.providers.prepareRerankRoutes( - modelId, - request, - options - ); - this.executionMetrics.recordPlan('rerank', routes); - const { nativeDispatch, transport } = buildRerankPlanArtifacts(routes); - const plan = { - transport, - request: { - kind: 'rerank', - cond: { modelId }, - modelId, - request, - options, - }, - routePolicy: { - fallbackOrder: buildFallbackOrder(routes), - }, - runtimePolicy: {}, - attachmentPolicy: { materializeRemoteAttachments: false }, - responsePostprocess: { mode: 'rerank' }, - hostPersistence: { - persistAssistantTurn: false, - outputKind: 'rerank', - }, - hostContext: { - signal: options?.signal, - }, - } as Omit< - ExecutionPlanForKind<'rerank'>, - 'nativeDispatch' | 'serializable' - >; - - return { - nativeDispatch, - serializable: buildSerializableExecutionPlan(routes, plan), - ...plan, - }; - } -} diff --git a/packages/backend/server/src/plugins/copilot/runtime/hosts/action-stream-host.ts b/packages/backend/server/src/plugins/copilot/runtime/hosts/action-stream-host.ts index 523e9894ad..6b04b53e68 100644 --- a/packages/backend/server/src/plugins/copilot/runtime/hosts/action-stream-host.ts +++ b/packages/backend/server/src/plugins/copilot/runtime/hosts/action-stream-host.ts @@ -1,6 +1,9 @@ import { Injectable } from '@nestjs/common'; -import type { LlmImageResponse } from '../../../../native'; +import { + getCopilotActionRecipe, + type LlmImageResponse, +} from '../../../../native'; import { PromptService } from '../../prompt'; import type { PromptMessage } from '../../providers/types'; import type { ChatSession } from '../../session'; @@ -8,6 +11,10 @@ import { ChatQuerySchema } from '../../types'; import { projectActionEventToChatEvent } from '../action-output-projector'; import type { ActionRuntimeBridgeEvent } from '../action-runtime-bridge'; import { ActionRuntimeBridge } from '../action-runtime-bridge'; +import { + buildStructuredResponseFromSchemaJson, + requireStructuredOutputContract, +} from '../contracts'; import { ConversationHost } from './conversation-host'; import { ImageResultHost } from './image-result-host'; @@ -17,32 +24,6 @@ function firstQueryValue(value: string | string[] | undefined) { return Array.isArray(value) ? value[0] : value; } -const ACTION_PROMPTS: Record = { - 'mindmap.generate': 'mindmap.generate', - 'slides.outline': 'slides.outline', -}; - -type ImageActionRoutePreparation = { - modelId?: string; - messages: PromptMessage[]; - options: Record; -}; - -function isImageAction(id: string) { - return id.startsWith('image.filter.'); -} - -function actionTextResultSchema() { - return { - type: 'object', - properties: { - result: { type: 'string' }, - }, - required: ['result'], - additionalProperties: false, - }; -} - @Injectable() export class ActionStreamHost { constructor( @@ -73,6 +54,7 @@ export class ActionStreamHost { firstQueryValue(query.actionId) ?? prepared.session.config.promptName; const actionId = requestedActionId; const actionVersion = firstQueryValue(query.actionVersion) ?? 'v1'; + const recipe = getCopilotActionRecipe(actionId, actionVersion); const retryOf = parsedQuery.retry ? firstQueryValue(query.runId) : undefined; @@ -81,19 +63,16 @@ export class ActionStreamHost { ...this.conversations.buildLatestTurnPromptParams(prepared.latestTurn), }; const finalMessage = await this.preparePromptMessages( - actionId, + recipe.promptRef, prepared.session, params ); - const imageRoutes = await this.prepareImageRoutes( - actionId, - prepared.session, - params, - userId, - parsedQuery.byokLeaseId, - prepared.quotaBackedRoutesAllowed, - signal - ); + const responseContract = recipe.responseContract + ? requireStructuredOutputContract( + buildStructuredResponseFromSchemaJson(recipe.responseContract.schema) + ) + : undefined; + const producesImage = recipe.outputProjection === 'first_image'; const runStream = this.bridge.runStream({ userId, workspaceId: prepared.session.config.workspaceId, @@ -108,7 +87,7 @@ export class ActionStreamHost { params, messageId: prepared.messageId, }, - persistAttachment: isImageAction(actionId) + persistAttachment: producesImage ? attachment => this.persistImageAttachment( userId, @@ -116,35 +95,25 @@ export class ActionStreamHost { attachment ) : undefined, - prepareStructuredRoutes: isImageAction(actionId) - ? undefined - : { - stepId: 'generate', - modelId: - typeof query.modelId === 'string' && query.modelId - ? query.modelId - : undefined, - messages: finalMessage, - responseSchemaJson: actionTextResultSchema(), - options: { - ...prepared.session.config.promptConfig, - signal, - user: userId, - workspace: prepared.session.config.workspaceId, - session: sessionId, - byokLeaseId: parsedQuery.byokLeaseId, - quotaBackedRoutesAllowed: prepared.quotaBackedRoutesAllowed, - featureKind: 'action', - }, - }, - prepareImageRoutes: imageRoutes - ? { - stepId: 'generate-image', - modelId: imageRoutes.modelId, - messages: imageRoutes.messages, - options: imageRoutes.options, - } - : undefined, + step: { + slot: recipe.slot, + builtInRouteId: recipe.promptRef, + profileId: parsedQuery.profileId, + modelId: parsedQuery.modelId, + messages: finalMessage, + responseContract, + options: { + ...prepared.session.config.promptConfig, + signal, + user: userId, + workspace: prepared.session.config.workspaceId, + session: sessionId, + byokLeaseId: parsedQuery.byokLeaseId, + managedTargetId: parsedQuery.routeTargetId, + quotaBackedRoutesAllowed: prepared.quotaBackedRoutesAllowed, + featureKind: producesImage ? 'image' : 'action', + }, + }, signal, }); @@ -157,18 +126,13 @@ export class ActionStreamHost { } private async preparePromptMessages( - actionId: string, + promptRef: string, session: ChatSession, params: Record ): Promise { - const promptName = ACTION_PROMPTS[actionId]; - if (!promptName) { - return session.finish(params); - } - - const prompt = await this.prompts.get(promptName); + const prompt = await this.prompts.get(promptRef); if (!prompt) { - throw new Error(`Prompt ${promptName} not found`); + throw new Error(`Prompt ${promptRef} not found`); } return this.prompts.finish( prompt, @@ -177,44 +141,6 @@ export class ActionStreamHost { ); } - private async prepareImageRoutes( - actionId: string, - session: ChatSession, - params: Record, - userId: string, - byokLeaseId?: string, - quotaBackedRoutesAllowed?: boolean, - signal?: AbortSignal - ): Promise { - if (!isImageAction(actionId)) { - return undefined; - } - - const prompt = await this.prompts.get(actionId); - if (!prompt) { - throw new Error(`Prompt ${actionId} not found`); - } - const finalMessage = this.prompts.finish( - prompt, - params as Record, - session.config.sessionId - ); - return { - modelId: prompt.model, - messages: finalMessage, - options: { - ...prompt.config, - signal, - user: userId, - workspace: session.config.workspaceId, - session: session.config.sessionId, - byokLeaseId, - quotaBackedRoutesAllowed, - featureKind: 'image', - }, - }; - } - private async persistImageAttachment( userId: string, workspaceId: string, diff --git a/packages/backend/server/src/plugins/copilot/runtime/hosts/attachment-materialization-planner.ts b/packages/backend/server/src/plugins/copilot/runtime/hosts/attachment-materialization-planner.ts deleted file mode 100644 index f1eb1ceaae..0000000000 --- a/packages/backend/server/src/plugins/copilot/runtime/hosts/attachment-materialization-planner.ts +++ /dev/null @@ -1,133 +0,0 @@ -import type { LlmBackendConfig, LlmProtocol } from '../../../../native'; -import { llmPlanAttachmentReference } from '../../../../native'; -import type { PromptAttachment } from '../../providers/types'; -import { - type AdmittedAttachmentSource, - admittedAttachmentToPromptAttachment, -} from './attachment-admission'; - -export type AdmittedAttachmentMaterializationPlan = { - mode: 'inline'; - reason: 'admitted_bytes'; - attachment: PromptAttachment; -}; - -export type HostAttachmentMaterializationRequest = { - attachmentId: string; - target: 'bytes' | 'data'; - providerConstraint?: string; - maxSize: number; - timeoutMs: number; - redirectPolicy: 'follow-safe'; - expectedMime?: string; - url: string; -}; - -type RemoteReferenceReason = - | 'generic_remote_reference' - | 'gemini_api_file_uri' - | 'gemini_api_youtube_url'; - -type MaterializationRequestReason = - | 'generic_remote_reference' - | 'gemini_api_inline_http_url' - | 'unsupported_scheme' - | 'non_url_source'; - -function assertRemoteReferenceReason( - reason: string -): asserts reason is RemoteReferenceReason { - if ( - reason !== 'generic_remote_reference' && - reason !== 'gemini_api_file_uri' && - reason !== 'gemini_api_youtube_url' - ) { - throw new Error(`Unexpected remote attachment reference reason: ${reason}`); - } -} - -function assertMaterializationRequestReason( - reason: string -): asserts reason is MaterializationRequestReason { - if ( - reason !== 'gemini_api_inline_http_url' && - reason !== 'generic_remote_reference' && - reason !== 'unsupported_scheme' && - reason !== 'non_url_source' - ) { - throw new Error(`Unexpected attachment materialization reason: ${reason}`); - } -} - -export async function planHostUrlAttachmentMaterialization( - protocol: LlmProtocol, - backendConfig: LlmBackendConfig, - input: { - attachmentId: string; - url: string; - expectedMime?: string; - maxSize: number; - timeoutMs?: number; - } -): Promise< - | { - mode: 'remote_reference'; - reason: - | 'generic_remote_reference' - | 'gemini_api_file_uri' - | 'gemini_api_youtube_url'; - url: string; - } - | { - mode: 'materialization_request'; - reason: - | 'generic_remote_reference' - | 'gemini_api_inline_http_url' - | 'unsupported_scheme' - | 'non_url_source'; - request: HostAttachmentMaterializationRequest; - } -> { - const plan = await llmPlanAttachmentReference(protocol, backendConfig, { - url: input.url, - }); - const forceHostMaterialization = - protocol === 'gemini' && - backendConfig.request_layer === 'gemini_vertex' && - plan.reason === 'generic_remote_reference'; - - if (plan.mode === 'remote' && !forceHostMaterialization) { - assertRemoteReferenceReason(plan.reason); - return { - mode: 'remote_reference', - reason: plan.reason, - url: input.url, - }; - } - - assertMaterializationRequestReason(plan.reason); - return { - mode: 'materialization_request', - reason: plan.reason, - request: { - attachmentId: input.attachmentId, - target: 'bytes', - providerConstraint: protocol, - maxSize: input.maxSize, - timeoutMs: input.timeoutMs ?? 15_000, - redirectPolicy: 'follow-safe', - expectedMime: input.expectedMime, - url: input.url, - }, - }; -} - -export function planAdmittedAttachmentMaterialization( - source: AdmittedAttachmentSource -): AdmittedAttachmentMaterializationPlan { - return { - mode: 'inline', - reason: 'admitted_bytes', - attachment: admittedAttachmentToPromptAttachment(source), - }; -} diff --git a/packages/backend/server/src/plugins/copilot/runtime/hosts/capability-policy-host.ts b/packages/backend/server/src/plugins/copilot/runtime/hosts/capability-policy-host.ts deleted file mode 100644 index 86f0f5e28c..0000000000 --- a/packages/backend/server/src/plugins/copilot/runtime/hosts/capability-policy-host.ts +++ /dev/null @@ -1,124 +0,0 @@ -import { Injectable } from '@nestjs/common'; -import { ModuleRef } from '@nestjs/core'; - -import { ServerFeature, ServerService } from '../../../../core'; -import { QuotaStateService } from '../../../../core/quota/state'; -import type { ChatSession } from '../../session'; -import { type ToolsConfig } from '../../types'; -import { getTools } from '../../utils'; -import { - ModelSelectionPolicy, - type ResolveModelInput, -} from '../model-selection-policy'; - -export type ChatSelectionOptions = { - responseMode: 'text' | 'object' | 'image'; - modelId?: string; - reasoning?: boolean; - webSearch?: boolean; - toolsConfig?: ToolsConfig; - byokLeaseId?: string; - billingUnitId?: string; - featureKind?: 'chat' | 'action' | 'image'; - quotaBackedRoutesAllowed?: boolean; -}; - -type ResolvePolicyModelInput = ResolveModelInput & { - proModels?: string[] | null; - userId?: string; - paymentEnabled?: boolean; -}; - -@Injectable() -export class CapabilityPolicyHost { - constructor( - private readonly server: ServerService, - private readonly moduleRef: ModuleRef, - private readonly modelSelection: ModelSelectionPolicy - ) {} - - private async hasAiProAccess( - userId: string | undefined, - paymentEnabled: boolean | undefined - ) { - if (!paymentEnabled || !userId) { - return false; - } - - try { - const state = await this.moduleRef - .get(QuotaStateService, { strict: false }) - .reconcileUserQuotaState(userId); - const flags = state.flags as { unlimitedCopilot?: boolean }; - return ( - !!flags.unlimitedCopilot || - ['pro', 'lifetime_pro', 'ai'].includes(state.plan) - ); - } catch { - return false; - } - } - - private async resolveModel(input: ResolvePolicyModelInput) { - const resolved = this.modelSelection.resolveRequestedModel(input); - if (!resolved.matchedOptionalModel) { - return resolved.selectedModel; - } - - if ( - input.paymentEnabled && - this.modelSelection.matchesModelList( - input.proModels ?? [], - input.requestedModelId - ) && - !(await this.hasAiProAccess(input.userId, input.paymentEnabled)) - ) { - return input.defaultModel; - } - - return resolved.selectedModel; - } - - async selectChat(session: ChatSession, options: ChatSelectionOptions) { - const model = await this.resolveChatModel({ - userId: session.config.userId, - defaultModel: session.model, - optionalModels: session.optionalModels, - proModels: session.config.promptConfig?.proModels, - requestedModelId: options.modelId, - paymentEnabled: this.server.features.includes(ServerFeature.Payment), - }); - const tools = getTools( - session.config.promptConfig?.tools, - options.toolsConfig - ); - return { - model, - providerOptions: { - ...session.config.promptConfig, - user: session.config.userId, - session: session.config.sessionId, - workspace: session.config.workspaceId, - byokLeaseId: options.byokLeaseId, - billingUnitId: options.billingUnitId, - featureKind: options.featureKind ?? 'chat', - quotaBackedRoutesAllowed: options.quotaBackedRoutesAllowed, - reasoning: options.reasoning, - webSearch: options.webSearch, - tools, - }, - }; - } - - async resolveChatModel(input: ResolvePolicyModelInput) { - return await this.resolveModel(input); - } - - async resolvePromptModel(input: ResolveModelInput) { - return this.modelSelection.resolveRequestedModel(input).selectedModel; - } - - async resolveFixedTaskModel(input: ResolveModelInput) { - return this.modelSelection.resolveRequestedModel(input).selectedModel; - } -} diff --git a/packages/backend/server/src/plugins/copilot/runtime/hosts/conversation-host.ts b/packages/backend/server/src/plugins/copilot/runtime/hosts/conversation-host.ts index 3b94cea13c..21ac416850 100644 --- a/packages/backend/server/src/plugins/copilot/runtime/hosts/conversation-host.ts +++ b/packages/backend/server/src/plugins/copilot/runtime/hosts/conversation-host.ts @@ -5,8 +5,8 @@ import { CopilotSessionNotFound, Mutex, } from '../../../../base'; -import { CopilotAccessPolicy } from '../../access'; import { CompatSubmissionStore } from '../../compat/submission-store'; +import { ConversationPolicy } from '../../conversation/policy'; import { canonicalizeTurnTrace, type Turn, @@ -35,7 +35,7 @@ export class ConversationHost { private readonly sessions: ChatSessionService, private readonly submissions: CompatSubmissionStore, private readonly mutex: Mutex, - private readonly access: CopilotAccessPolicy + private readonly policy: ConversationPolicy ) {} private async loadAcceptedTurn( @@ -109,31 +109,22 @@ export class ConversationHost { session: ChatSession, sessionId: string, messageId?: string, - retry = false, - byokLeaseId?: string + retry = false ): Promise { - const resolveChatRouteAccess = () => - this.access.resolveTurnRouteAccess({ - userId, - workspaceId: session.config.workspaceId, - byokLeaseId, - featureKind: 'chat', - }); + const quotaBackedRoutesAllowed = () => this.policy.hasQuota(userId); if (!messageId) { await this.sessions.revertLatestMessage(sessionId, false); session.revertLatestMessage(false); if (!session.latestUserTurn) { - const routeAccess = await resolveChatRouteAccess(); return { turn: session.latestUserTurn, - quotaBackedRoutesAllowed: routeAccess.quotaBackedRoutesAllowed, + quotaBackedRoutesAllowed: await quotaBackedRoutesAllowed(), }; } - const routeAccess = await resolveChatRouteAccess(); return { turn: session.latestUserTurn, - quotaBackedRoutesAllowed: routeAccess.quotaBackedRoutesAllowed, + quotaBackedRoutesAllowed: await quotaBackedRoutesAllowed(), }; } @@ -177,7 +168,7 @@ export class ConversationHost { }; } - const routeAccess = await resolveChatRouteAccess(); + const quotaAllowed = await quotaBackedRoutesAllowed(); const submission = await this.submissions.get(messageId); if (!submission || submission.sessionId !== sessionId) { @@ -192,7 +183,6 @@ export class ConversationHost { const turn = await this.sessions.appendTurn({ sessionId, userId: session.config.userId, - prompt: { model: session.model }, compatSubmissionId: messageId, turn: { conversationId: sessionId, @@ -213,7 +203,7 @@ export class ConversationHost { session.pushPersistedTurn(turn); return { turn, - quotaBackedRoutesAllowed: routeAccess.quotaBackedRoutesAllowed, + quotaBackedRoutesAllowed: quotaAllowed, }; } @@ -222,8 +212,7 @@ export class ConversationHost { sessionId: string, query: Record ): Promise { - const { messageId, retry, params, byokLeaseId } = - ChatQuerySchema.parse(query); + const { messageId, retry, params } = ChatQuerySchema.parse(query); const session = await this.sessions.get(sessionId); if (!session || session.config.userId !== userId) { throw new CopilotSessionNotFound(); @@ -233,8 +222,7 @@ export class ConversationHost { session, sessionId, messageId, - retry, - byokLeaseId + retry ); const currentUserMessage = session.stashTurns.findLast(turn => turn.role === 'user') ?? @@ -280,7 +268,6 @@ export class ConversationHost { const persisted = await this.sessions.appendTurn({ sessionId: session.config.sessionId, userId: session.config.userId, - prompt: { model: session.model }, turn: assistantTurn, }); session.pushPersistedTurn(persisted); diff --git a/packages/backend/server/src/plugins/copilot/runtime/hosts/tool-executor-host.ts b/packages/backend/server/src/plugins/copilot/runtime/hosts/tool-executor-host.ts deleted file mode 100644 index c64e1d2337..0000000000 --- a/packages/backend/server/src/plugins/copilot/runtime/hosts/tool-executor-host.ts +++ /dev/null @@ -1,43 +0,0 @@ -import { Injectable } from '@nestjs/common'; - -import type { NodeTextMiddleware } from '../../config'; -import type { - CopilotChatOptions, - CopilotChatTools, -} from '../../providers/types'; -import type { CopilotTool, CopilotToolSet } from '../../tools'; -import type { ToolLoopBackend } from '../tool/bridge'; -import { ToolRuntime } from '../tool-runtime'; - -export type ProviderSpecificToolResolver = ( - toolName: CopilotChatTools, - model: string -) => [string, CopilotTool?] | undefined; - -@Injectable() -export class ToolExecutorHost { - constructor(private readonly runtime: ToolRuntime) {} - - async getTools( - options: CopilotChatOptions, - model: string, - resolveProviderSpecificTool?: ProviderSpecificToolResolver - ): Promise { - return await this.runtime.getTools( - options, - model, - resolveProviderSpecificTool - ); - } - - createNativeAdapter( - backend: ToolLoopBackend, - tools: CopilotToolSet, - options: { - maxSteps?: number; - nodeTextMiddleware?: NodeTextMiddleware[]; - } = {} - ) { - return this.runtime.createNativeAdapter(backend, tools, options); - } -} diff --git a/packages/backend/server/src/plugins/copilot/runtime/model-selection-policy.ts b/packages/backend/server/src/plugins/copilot/runtime/model-selection-policy.ts deleted file mode 100644 index 10a9f173f0..0000000000 --- a/packages/backend/server/src/plugins/copilot/runtime/model-selection-policy.ts +++ /dev/null @@ -1,55 +0,0 @@ -import { Injectable } from '@nestjs/common'; - -import { CopilotSessionInvalidInput } from '../../../base'; -import { llmResolveRequestedModelMatch } from '../../../native'; -import { CopilotProviderRegistryService } from '../providers/registry-service'; - -export type ResolveModelInput = { - defaultModel: string; - optionalModels?: string[] | null; - requestedModelId?: string; -}; - -@Injectable() -export class ModelSelectionPolicy { - constructor(private readonly registries: CopilotProviderRegistryService) {} - - private getRegistry() { - return this.registries.getRegistry(); - } - - private matchRequestedModel( - optionalModels: string[], - requestedModelId?: string, - defaultModel?: string - ) { - return llmResolveRequestedModelMatch({ - providerIds: [...this.getRegistry().profiles.keys()], - optionalModels, - requestedModelId, - defaultModel, - }); - } - - resolveRequestedModel(input: ResolveModelInput): { - selectedModel: string; - matchedOptionalModel: boolean; - } { - if (!input.defaultModel) { - throw new CopilotSessionInvalidInput('Model is required'); - } - const matched = this.matchRequestedModel( - input.optionalModels ?? [], - input.requestedModelId, - input.defaultModel - ); - return { - selectedModel: matched.selectedModel ?? input.defaultModel, - matchedOptionalModel: matched.matchedOptionalModel, - }; - } - - matchesModelList(models: string[], modelId?: string) { - return this.matchRequestedModel(models, modelId).matchedOptionalModel; - } -} diff --git a/packages/backend/server/src/plugins/copilot/runtime/native-errors.ts b/packages/backend/server/src/plugins/copilot/runtime/native-errors.ts index df67855615..c122a3e098 100644 --- a/packages/backend/server/src/plugins/copilot/runtime/native-errors.ts +++ b/packages/backend/server/src/plugins/copilot/runtime/native-errors.ts @@ -1,4 +1,4 @@ -import { NetworkError } from '../../../base'; +import { CopilotQuotaExceeded, NetworkError } from '../../../base'; const LLM_TIMEOUT_ERROR_PREFIX = 'llm_timeout:'; @@ -18,6 +18,9 @@ function nativeErrorMessage(error: unknown) { export function mapNativeSemanticError(error: unknown): unknown { const message = nativeErrorMessage(error); + if (message === 'access_unavailable') { + return new CopilotQuotaExceeded(); + } if (message?.startsWith(LLM_TIMEOUT_ERROR_PREFIX)) { return new NetworkError( message.slice(LLM_TIMEOUT_ERROR_PREFIX.length).trim() || diff --git a/packages/backend/server/src/plugins/copilot/runtime/native-execution-engine.ts b/packages/backend/server/src/plugins/copilot/runtime/native-execution-engine.ts deleted file mode 100644 index f67dbd6ab7..0000000000 --- a/packages/backend/server/src/plugins/copilot/runtime/native-execution-engine.ts +++ /dev/null @@ -1,536 +0,0 @@ -import { Injectable, Logger } from '@nestjs/common'; - -import { NoCopilotProviderAvailable } from '../../../base'; -import { - llmDispatchPlan, - llmDispatchPlanStream, - type LlmDispatchResponse, - llmEmbeddingDispatchPlan, - llmImageDispatchPlan, - type LlmImageResponse, - llmRerankDispatchPlan, - llmStructuredDispatchPlan, - llmValidateJsonSchema, - parseNativeStructuredOutput, -} from '../../../native'; -import { type ByokFeatureKind, ByokService } from '../byok'; -import { type StreamObject } from '../providers/types'; -import { CopilotExecutionMetrics } from './execution-metrics'; -import { - type ExecutionPlan, - type ExecutionPlanForKind, - type NativeChatDispatchPlan, - type NativeImageDispatchPlan, -} from './execution-plan'; -import { mapNativeSemanticError } from './native-errors'; -import { - createNativeToolLoopAdapter, - NativeProviderAdapter, - type NativeProviderAdapterOptions, -} from './tool/native-adapter'; - -const logger = new Logger('NativeExecutionEngine'); - -function modelIdForError(modelId?: string) { - return modelId ?? 'auto'; -} - -type ExecutionPlanKind = ExecutionPlan['request']['kind']; -type ValueExecutionKind = Exclude< - ExecutionPlanKind, - 'streamText' | 'streamObject' | 'image' ->; -type StreamExecutionKind = Extract< - ExecutionPlanKind, - 'streamText' | 'streamObject' ->; -export type NativeImageArtifact = LlmImageResponse['images'][number]; - -function resolveAbortSignal( - signalOrOptions?: AbortSignal | { signal?: AbortSignal } -) { - return signalOrOptions && - typeof signalOrOptions === 'object' && - 'aborted' in signalOrOptions - ? signalOrOptions - : signalOrOptions?.signal; -} - -function extractTextResponse(response: LlmDispatchResponse) { - return response.message.content - .filter(part => part.type === 'text' || part.type === 'reasoning') - .map(part => part.text) - .join('') - .trim(); -} - -function getUsageContext(plan: ExecutionPlan) { - const options = 'options' in plan.request ? plan.request.options : undefined; - const requestFeatureKind = - plan.request.kind === 'text' || - plan.request.kind === 'streamText' || - plan.request.kind === 'streamObject' - ? 'chat' - : plan.request.kind; - return { - workspaceId: options?.workspace, - userId: options?.user, - sessionId: options?.session, - taskId: options?.taskId, - actionId: options?.actionId, - billingUnitId: options?.billingUnitId, - featureKind: options?.featureKind ?? requestFeatureKind, - }; -} - -async function recordByokUsage( - byok: ByokService, - plan: ExecutionPlan, - input: { - providerId?: string; - model?: string | null; - usage?: LlmDispatchResponse['usage']; - } -) { - const context = getUsageContext(plan); - try { - await byok.recordUsage({ - workspaceId: context.workspaceId, - userId: context.userId, - sessionId: context.sessionId, - taskId: context.taskId, - actionId: context.actionId, - billingUnitId: context.billingUnitId, - featureKind: context.featureKind as ByokFeatureKind, - providerId: input.providerId, - model: input.model, - usage: input.usage, - }); - } catch (error) { - logger.warn( - `Failed to record BYOK usage: ${ - error instanceof Error ? error.message : String(error) - }` - ); - } -} - -async function recordSingleByokRouteFailure( - byok: ByokService, - plan: ExecutionPlan, - error: unknown -) { - const [providerId] = plan.routePolicy.fallbackOrder; - if (plan.routePolicy.fallbackOrder.length !== 1 || !providerId) { - return; - } - const context = getUsageContext(plan); - try { - await byok.recordProviderFailure({ - workspaceId: context.workspaceId, - providerId, - featureKind: context.featureKind as ByokFeatureKind, - error, - }); - } catch (recordError) { - logger.warn( - `Failed to record BYOK provider failure: ${ - recordError instanceof Error ? recordError.message : String(recordError) - }` - ); - } -} - -function recordPreparedDispatch( - executionMetrics: CopilotExecutionMetrics | undefined, - plan: ExecutionPlan, - routeCount: number -) { - executionMetrics?.recordDispatch( - plan.request.kind, - 'prepared_routes', - routeCount - ); -} - -function createNativeChatAdapter( - dispatch: NativeChatDispatchPlan, - options?: { - onUsage?: NativeProviderAdapterOptions['onUsage']; - } -) { - if (dispatch.hasTools) { - return createNativeToolLoopAdapter( - { preparedRoutes: dispatch.routes }, - dispatch.prepared.tools, - { - maxSteps: dispatch.prepared.maxSteps, - nodeTextMiddleware: dispatch.prepared.postprocess?.nodeTextMiddleware, - onUsage: options?.onUsage, - } - ); - } - - const nativeDispatch = ( - _nativeRequest: typeof dispatch.prepared.request, - signalOrOptions?: AbortSignal | { signal?: AbortSignal } - ) => - llmDispatchPlanStream({ - preparedRoutes: dispatch.routes, - signal: resolveAbortSignal(signalOrOptions), - }); - - return new NativeProviderAdapter(nativeDispatch, { - nodeTextMiddleware: dispatch.prepared.postprocess?.nodeTextMiddleware, - onUsage: options?.onUsage, - }); -} - -async function runPreparedValuePlan( - plan: ExecutionPlan, - routeCount: number, - executionMetrics: CopilotExecutionMetrics | undefined, - run: () => Promise, - byok: ByokService -) { - recordPreparedDispatch(executionMetrics, plan, routeCount); - try { - return await run(); - } catch (error) { - const mapped = mapNativeSemanticError(error); - await recordSingleByokRouteFailure(byok, plan, mapped); - throw mapped; - } -} - -async function* mapPreparedStreamErrors( - source: AsyncIterable, - plan: ExecutionPlan, - byok: ByokService -): AsyncIterableIterator { - try { - yield* source; - } catch (error) { - const mapped = mapNativeSemanticError(error); - await recordSingleByokRouteFailure(byok, plan, mapped); - throw mapped; - } -} - -async function runChatValuePlan( - plan: ExecutionPlan, - dispatch: NativeChatDispatchPlan, - executionMetrics: CopilotExecutionMetrics | undefined, - byok: ByokService -) { - const adapter = createNativeChatAdapter(dispatch); - return await runPreparedValuePlan( - plan, - dispatch.routes.length, - executionMetrics, - async () => { - if ( - !dispatch.hasTools && - !dispatch.prepared.postprocess?.nodeTextMiddleware?.length - ) { - const result = await llmDispatchPlan({ - preparedRoutes: dispatch.routes, - }); - await recordByokUsage(byok, plan, { - providerId: result.provider_id, - model: result.response.model, - usage: result.response.usage, - }); - return extractTextResponse(result.response); - } - - if (plan.request.kind !== 'text') { - throw new Error('chat value dispatch requires text plan'); - } - - return await adapter.text( - dispatch.prepared.request, - plan.hostContext.signal, - plan.request.messages - ); - }, - byok - ); -} - -async function* runChatStreamPlan( - plan: ExecutionPlan, - dispatch: NativeChatDispatchPlan, - executionMetrics: CopilotExecutionMetrics | undefined, - byok: ByokService -): AsyncIterableIterator { - const adapter = createNativeChatAdapter(dispatch, { - onUsage: async usage => { - await recordByokUsage(byok, plan, { - providerId: usage.providerId, - model: usage.model, - usage: usage.usage, - }); - }, - }); - recordPreparedDispatch(executionMetrics, plan, dispatch.routes.length); - - if (plan.request.kind === 'streamText') { - yield* mapPreparedStreamErrors( - adapter.streamText( - dispatch.prepared.request, - plan.hostContext.signal, - plan.request.messages - ), - plan, - byok - ); - return; - } - - if (plan.request.kind === 'streamObject') { - yield* mapPreparedStreamErrors( - adapter.streamObject( - dispatch.prepared.request, - plan.hostContext.signal, - plan.request.messages - ), - plan, - byok - ); - return; - } - - throw new Error('chat stream dispatch requires streamText/streamObject plan'); -} - -async function* runPreparedImageArtifactPlan( - dispatch: NativeImageDispatchPlan, - plan: ExecutionPlan, - executionMetrics: CopilotExecutionMetrics | undefined, - byok: ByokService -): AsyncIterableIterator { - if (plan.request.kind !== 'image') { - throw new Error('image dispatch requires image plan'); - } - - recordPreparedDispatch(executionMetrics, plan, dispatch.routes.length); - let result; - try { - result = await llmImageDispatchPlan({ - preparedRoutes: dispatch.routes, - }); - await recordByokUsage(byok, plan, { - providerId: result.provider_id, - model: dispatch.prepared.route.model, - usage: result.response.usage - ? { - prompt_tokens: result.response.usage.input_tokens ?? 0, - completion_tokens: result.response.usage.output_tokens ?? 0, - total_tokens: result.response.usage.total_tokens ?? 0, - } - : undefined, - }); - } catch (error) { - const mapped = mapNativeSemanticError(error); - await recordSingleByokRouteFailure(byok, plan, mapped); - throw mapped; - } - for (const artifact of result.response.images) { - yield artifact; - } -} - -async function executePreparedPlan( - plan: ExecutionPlan, - executionMetrics: CopilotExecutionMetrics | undefined, - byok: ByokService -): Promise { - switch (plan.request.kind) { - case 'text': { - const dispatch = plan.nativeDispatch?.chat; - return dispatch - ? await runChatValuePlan(plan, dispatch, executionMetrics, byok) - : null; - } - case 'structured': { - const dispatch = plan.nativeDispatch?.structured; - if (!dispatch) { - return null; - } - return await runPreparedValuePlan( - plan, - dispatch.routes.length, - executionMetrics, - async () => { - const result = await llmStructuredDispatchPlan({ - preparedRoutes: dispatch.routes, - }); - await recordByokUsage(byok, plan, { - providerId: result.provider_id, - model: result.response.model, - usage: result.response.usage, - }); - const parsed = parseNativeStructuredOutput(result.response); - const validated = llmValidateJsonSchema( - dispatch.prepared.request.schema, - parsed - ); - return JSON.stringify(validated); - }, - byok - ); - } - case 'embedding': { - const dispatch = plan.nativeDispatch?.embedding; - if (!dispatch) { - return null; - } - return await runPreparedValuePlan( - plan, - dispatch.routes.length, - executionMetrics, - async () => { - const result = await llmEmbeddingDispatchPlan({ - preparedRoutes: dispatch.routes, - }); - await recordByokUsage(byok, plan, { - providerId: result.provider_id, - model: result.response.model, - usage: result.response.usage - ? { - prompt_tokens: result.response.usage.prompt_tokens, - completion_tokens: 0, - total_tokens: result.response.usage.total_tokens, - } - : undefined, - }); - return result.response.embeddings; - }, - byok - ); - } - case 'rerank': { - const dispatch = plan.nativeDispatch?.rerank; - if (!dispatch) { - return null; - } - return await runPreparedValuePlan( - plan, - dispatch.routes.length, - executionMetrics, - async () => { - const result = await llmRerankDispatchPlan({ - preparedRoutes: dispatch.routes, - }); - await recordByokUsage(byok, plan, { - providerId: result.provider_id, - model: result.response.model, - }); - return result.response.scores; - }, - byok - ); - } - default: - return null; - } -} - -function executePreparedStreamPlan( - plan: ExecutionPlan, - executionMetrics: CopilotExecutionMetrics | undefined, - byok: ByokService -): AsyncIterableIterator | null { - switch (plan.request.kind) { - case 'streamText': - case 'streamObject': { - const dispatch = plan.nativeDispatch?.chat; - return dispatch - ? runChatStreamPlan(plan, dispatch, executionMetrics, byok) - : null; - } - default: - return null; - } -} - -function noRouteStream(plan: ExecutionPlan) { - return (async function* (): AsyncIterableIterator { - yield* [] as T[]; - throw new NoCopilotProviderAvailable({ - modelId: modelIdForError(plan.request.cond.modelId), - }); - })(); -} - -@Injectable() -export class NativeExecutionEngine { - constructor( - private readonly byok: ByokService, - private readonly executionMetrics?: CopilotExecutionMetrics - ) {} - - private noRoute(plan: ExecutionPlan): never { - throw new NoCopilotProviderAvailable({ - modelId: modelIdForError(plan.request.cond.modelId), - }); - } - - async execute( - plan: ExecutionPlanForKind<'text' | 'structured'> - ): Promise; - async execute(plan: ExecutionPlanForKind<'embedding'>): Promise; - async execute(plan: ExecutionPlanForKind<'rerank'>): Promise; - async execute( - plan: ExecutionPlanForKind - ): Promise { - const result = await executePreparedPlan( - plan, - this.executionMetrics, - this.byok - ); - if (result === null) { - return this.noRoute(plan); - } - - return result; - } - - executeStream( - plan: ExecutionPlanForKind<'streamText'> - ): AsyncIterableIterator; - executeStream( - plan: ExecutionPlanForKind<'streamObject'> - ): AsyncIterableIterator; - executeStream( - plan: ExecutionPlanForKind - ): AsyncIterableIterator { - const result = executePreparedStreamPlan( - plan, - this.executionMetrics, - this.byok - ); - if (result) { - return result; - } - - return noRouteStream(plan); - } - - executeImageArtifacts( - plan: ExecutionPlanForKind<'image'> - ): AsyncIterableIterator { - const dispatch = plan.nativeDispatch?.image; - if (dispatch) { - return runPreparedImageArtifactPlan( - dispatch, - plan, - this.executionMetrics, - this.byok - ); - } - - return noRouteStream(plan); - } -} diff --git a/packages/backend/server/src/plugins/copilot/runtime/native-request-runtime.ts b/packages/backend/server/src/plugins/copilot/runtime/native-request-runtime.ts index 3d1bf636b9..5cf6666c77 100644 --- a/packages/backend/server/src/plugins/copilot/runtime/native-request-runtime.ts +++ b/packages/backend/server/src/plugins/copilot/runtime/native-request-runtime.ts @@ -114,7 +114,7 @@ export async function buildCanonicalNativeRequest({ request = llmBuildCanonicalRequest({ model, messages: normalizedMessages, - maxTokens: options.maxTokens ?? undefined, + maxTokens: options.maxOutputTokens ?? undefined, temperature: options.temperature ?? undefined, tools: toolContracts, include, @@ -171,7 +171,7 @@ export async function buildCanonicalNativeStructuredRequest({ model, messages: normalizedMessages, schema: explicitResponseContract?.responseSchemaJson, - maxTokens: options.maxTokens ?? undefined, + maxTokens: options.maxOutputTokens ?? undefined, temperature: options.temperature ?? undefined, reasoning, strict: options.strict, diff --git a/packages/backend/server/src/plugins/copilot/runtime/prompt-runtime.ts b/packages/backend/server/src/plugins/copilot/runtime/prompt-runtime.ts index 418e5e8e52..65945062e9 100644 --- a/packages/backend/server/src/plugins/copilot/runtime/prompt-runtime.ts +++ b/packages/backend/server/src/plugins/copilot/runtime/prompt-runtime.ts @@ -1,4 +1,5 @@ -import { Injectable } from '@nestjs/common'; +/* oxlint-disable import/no-cycle -- Prompt execution delegates to the capability runtime. */ +import { forwardRef, Inject, Injectable } from '@nestjs/common'; import { CopilotPromptNotFound } from '../../../base'; import { PromptService } from '../prompt/service'; @@ -11,7 +12,6 @@ import { } from '../providers/types'; import { CapabilityRuntime } from './capability-runtime'; import type { RequiredStructuredOutputContract } from './contracts'; -import { CapabilityPolicyHost } from './hosts/capability-policy-host'; type PromptRuntimeStructuredContract = RequiredStructuredOutputContract; @@ -24,8 +24,11 @@ type PromptRuntimeStructuredProviderOptions = Omit< export class PromptRuntime { constructor( private readonly prompts: PromptService, - private readonly capabilityPolicy: CapabilityPolicyHost, - private readonly runtime: CapabilityRuntime + @Inject(forwardRef(() => CapabilityRuntime)) + private readonly runtime: Pick< + CapabilityRuntime, + 'text' | 'generateStructuredValue' + > ) {} private async preparePrompt( @@ -44,11 +47,8 @@ export class PromptRuntime { return { prompt, - modelId: await this.capabilityPolicy.resolvePromptModel({ - defaultModel: prompt.model, - optionalModels: prompt.optionalModels, - requestedModelId: options.modelId, - }), + builtInRouteId: prompt.name, + modelId: 'route-selected', finalMessages: [ ...this.prompts.finish(prompt, params), ...(options.appendMessages ?? []), @@ -75,6 +75,7 @@ export class PromptRuntime { { ...prepared.prompt.config, ...options.providerOptions, + builtInRouteId: prepared.builtInRouteId, }, { prefer: prepared.prefer } ); @@ -100,6 +101,7 @@ export class PromptRuntime { { ...prepared.prompt.config, ...options.providerOptions, + builtInRouteId: prepared.builtInRouteId, responseSchemaJson: options.responseContract.responseSchemaJson, schemaHash: options.responseContract.schemaHash, strict: options.strict, diff --git a/packages/backend/server/src/plugins/copilot/runtime/provider-chat-runtime.ts b/packages/backend/server/src/plugins/copilot/runtime/provider-chat-runtime.ts deleted file mode 100644 index 6786a5e9eb..0000000000 --- a/packages/backend/server/src/plugins/copilot/runtime/provider-chat-runtime.ts +++ /dev/null @@ -1,234 +0,0 @@ -import type { - CopilotProviderExecution, - PreparedNativeExecution, - PreparedNativeRequestOptions, - ProviderChatDriver, - ProviderChatDriverPrepareInput, -} from '../providers/provider-runtime-contract'; -import type { - CopilotChatOptions, - CopilotProviderModel, - CopilotProviderType, - ModelConditions, - ModelFullConditions, - PromptMessage, - StreamObject, -} from '../providers/types'; -import { ModelOutputType } from '../providers/types'; -import { - resolveDriverOrThrow, - resolvePreparedModelId, - runPreparedExecution, -} from './provider-driver-runtime'; -import type { NativeProviderAdapter } from './tool/native-adapter'; - -type MetricLabels = Record; - -export type ChatRuntimeContext = { - type: CopilotProviderType; - resolveChatDriver: () => ProviderChatDriver | undefined; - selectModel: (cond: ModelFullConditions) => CopilotProviderModel; - metricLabels: ( - model: string, - labels?: MetricLabels, - execution?: CopilotProviderExecution - ) => MetricLabels; - createPreparedExecutionAdapter: ( - prepared: PreparedNativeExecution - ) => NativeProviderAdapter; -}; - -type ChatExecutionMode = { - kind: ProviderChatDriverPrepareInput['kind']; - outputType: ModelOutputType; - unsupportedKind: 'text' | 'object'; - callMetric: string; - errorMetric: string; -}; - -export async function prepareNativeChatExecution( - resolveChatDriver: () => ProviderChatDriver | undefined, - buildPreparedNativeExecution: ( - options: PreparedNativeRequestOptions - ) => Promise, - input: ProviderChatDriverPrepareInput -): Promise { - const driver = resolveChatDriver(); - if (!driver) { - return null; - } - - const prepared = await driver.prepare(input); - if (!prepared) { - return null; - } - - return await buildPreparedNativeExecution({ - ...prepared, - execution: input.execution, - options: input.options, - }); -} - -async function runNativeChat( - context: ChatRuntimeContext, - prepareNativeExecution: ( - kind: ProviderChatDriverPrepareInput['kind'], - cond: ModelConditions, - messages: PromptMessage[], - options?: CopilotChatOptions, - execution?: CopilotProviderExecution - ) => Promise, - mode: ChatExecutionMode, - model: ModelConditions, - messages: PromptMessage[], - options: CopilotChatOptions | undefined, - execution: CopilotProviderExecution | undefined, - run: ( - adapter: NativeProviderAdapter, - prepared: PreparedNativeExecution, - signal: AbortSignal | undefined, - promptMessages: PromptMessage[] - ) => - | Promise - | AsyncIterableIterator - | AsyncIterableIterator -) { - const driver = resolveDriverOrThrow( - context.type, - mode.unsupportedKind, - context.resolveChatDriver - ); - const chatOptions = options ?? {}; - const prepared = await prepareNativeExecution( - mode.kind, - model, - messages, - chatOptions, - execution - ); - const modelId = resolvePreparedModelId( - context, - model, - mode.outputType, - prepared - ); - - return await runPreparedExecution({ - driver, - prepared, - modelId, - execution, - metricContext: context, - metricsName: { - call: mode.callMetric, - error: mode.errorMetric, - }, - execute: async preparedExecution => - await run( - context.createPreparedExecutionAdapter(preparedExecution), - preparedExecution, - chatOptions.signal, - messages - ), - }); -} - -export async function runNativeText( - context: ChatRuntimeContext, - prepareNativeExecution: ( - kind: ProviderChatDriverPrepareInput['kind'], - cond: ModelConditions, - messages: PromptMessage[], - options?: CopilotChatOptions, - execution?: CopilotProviderExecution - ) => Promise, - model: ModelConditions, - messages: PromptMessage[], - options?: CopilotChatOptions, - execution?: CopilotProviderExecution -) { - return (await runNativeChat( - context, - prepareNativeExecution, - { - kind: 'text', - outputType: ModelOutputType.Text, - unsupportedKind: 'text', - callMetric: 'chat_text_calls', - errorMetric: 'chat_text_errors', - }, - model, - messages, - options, - execution, - (adapter, prepared, signal, promptMessages) => - adapter.text(prepared.request, signal, promptMessages) - )) as string; -} - -export async function* runNativeStreamText( - context: ChatRuntimeContext, - prepareNativeExecution: ( - kind: ProviderChatDriverPrepareInput['kind'], - cond: ModelConditions, - messages: PromptMessage[], - options?: CopilotChatOptions, - execution?: CopilotProviderExecution - ) => Promise, - model: ModelConditions, - messages: PromptMessage[], - options?: CopilotChatOptions, - execution?: CopilotProviderExecution -): AsyncIterableIterator { - yield* (await runNativeChat( - context, - prepareNativeExecution, - { - kind: 'streamText', - outputType: ModelOutputType.Text, - unsupportedKind: 'text', - callMetric: 'chat_text_stream_calls', - errorMetric: 'chat_text_stream_errors', - }, - model, - messages, - options, - execution, - (adapter, prepared, signal, promptMessages) => - adapter.streamText(prepared.request, signal, promptMessages) - )) as AsyncIterableIterator; -} - -export async function* runNativeStreamObject( - context: ChatRuntimeContext, - prepareNativeExecution: ( - kind: ProviderChatDriverPrepareInput['kind'], - cond: ModelConditions, - messages: PromptMessage[], - options?: CopilotChatOptions, - execution?: CopilotProviderExecution - ) => Promise, - model: ModelConditions, - messages: PromptMessage[], - options?: CopilotChatOptions, - execution?: CopilotProviderExecution -): AsyncIterableIterator { - yield* (await runNativeChat( - context, - prepareNativeExecution, - { - kind: 'streamObject', - outputType: ModelOutputType.Object, - unsupportedKind: 'object', - callMetric: 'chat_object_stream_calls', - errorMetric: 'chat_object_stream_errors', - }, - model, - messages, - options, - execution, - (adapter, prepared, signal, promptMessages) => - adapter.streamObject(prepared.request, signal, promptMessages) - )) as AsyncIterableIterator; -} diff --git a/packages/backend/server/src/plugins/copilot/runtime/provider-driver-runtime.ts b/packages/backend/server/src/plugins/copilot/runtime/provider-driver-runtime.ts deleted file mode 100644 index fcd3f9eb63..0000000000 --- a/packages/backend/server/src/plugins/copilot/runtime/provider-driver-runtime.ts +++ /dev/null @@ -1,682 +0,0 @@ -import { - CopilotPromptInvalid, - CopilotProviderNotSupported, - metrics, -} from '../../../base'; -import { - buildLlmEmbeddingRequest, - buildLlmRerankRequest, - type LlmBackendConfig, - type LlmEmbeddingRequest, - type LlmProtocol, - type LlmRerankRequest, - type LlmStructuredRequest, - type LlmStructuredResponse, - llmValidateJsonSchema, - parseNativeStructuredOutput, -} from '../../../native'; -import type { ProviderMiddlewareConfig } from '../config'; -import { resolveProviderModelRoute } from '../providers/provider-model-runtime'; -import type { - CopilotProviderExecution, - EmbeddingProviderDriver, - ImageProviderDriver, - PreparedNativeEmbeddingExecution, - PreparedNativeImageExecution, - PreparedNativeRerankExecution, - PreparedNativeStructuredExecution, - RerankProviderDriver, - StructuredProviderDriver, -} from '../providers/provider-runtime-contract'; -import type { - CopilotChatOptions, - CopilotEmbeddingOptions, - CopilotImageOptions, - CopilotProviderModel, - CopilotProviderType, - CopilotRerankRequest, - CopilotStructuredOptions, - ModelAttachmentCapability, - ModelConditions, - ModelFullConditions, - PromptMessage, -} from '../providers/types'; -import { ModelOutputType } from '../providers/types'; -import { type RequiredStructuredOutputContract } from './contracts'; -import { buildNativeStructuredRequest } from './native-request-runtime'; - -const DEFAULT_EMBEDDING_TASK_TYPE = 'RETRIEVAL_DOCUMENT'; - -type MetricLabels = Record; -type DriverMetricNames = { - call: string; - error: string; -}; - -export type StructuredRuntimeContext = { - type: CopilotProviderType; - resolveStructuredDriver: () => StructuredProviderDriver | undefined; - checkParams: (input: { - cond: ModelFullConditions; - messages?: PromptMessage[]; - embeddings?: string[]; - options?: - | CopilotChatOptions - | CopilotStructuredOptions - | CopilotImageOptions; - withAttachment?: boolean; - execution?: CopilotProviderExecution; - }) => Promise; - selectModel: ( - cond: ModelFullConditions, - execution?: CopilotProviderExecution - ) => CopilotProviderModel; - getAttachCapability: ( - model: CopilotProviderModel, - outputType: ModelOutputType - ) => ModelAttachmentCapability | undefined; - getActiveProviderMiddleware: ( - execution?: CopilotProviderExecution - ) => ProviderMiddlewareConfig; - buildPreparedNativeStructuredExecution: ( - protocol: LlmProtocol, - backendConfig: LlmBackendConfig, - model: string, - request: LlmStructuredRequest, - execution?: CopilotProviderExecution - ) => PreparedNativeStructuredExecution; - createNativeStructuredDispatch: ( - backendConfig: LlmBackendConfig, - protocol: LlmProtocol, - execution?: CopilotProviderExecution - ) => (request: LlmStructuredRequest) => Promise; - metricLabels: ( - model: string, - labels?: MetricLabels, - execution?: CopilotProviderExecution - ) => MetricLabels; -}; - -export type EmbeddingRuntimeContext = { - type: CopilotProviderType; - resolveEmbeddingDriver: () => EmbeddingProviderDriver | undefined; - checkParams: (input: { - cond: ModelFullConditions; - messages?: PromptMessage[]; - embeddings?: string[]; - options?: - | CopilotChatOptions - | CopilotStructuredOptions - | CopilotImageOptions; - withAttachment?: boolean; - execution?: CopilotProviderExecution; - }) => Promise; - selectModel: ( - cond: ModelFullConditions, - execution?: CopilotProviderExecution - ) => CopilotProviderModel; - buildPreparedNativeEmbeddingExecution: ( - protocol: LlmProtocol, - backendConfig: LlmBackendConfig, - model: string, - request: LlmEmbeddingRequest, - execution?: CopilotProviderExecution - ) => PreparedNativeEmbeddingExecution; - createNativeEmbeddingDispatch: ( - backendConfig: LlmBackendConfig, - protocol: LlmProtocol, - execution?: CopilotProviderExecution - ) => (request: LlmEmbeddingRequest) => Promise<{ embeddings: number[][] }>; - metricLabels: ( - model: string, - labels?: MetricLabels, - execution?: CopilotProviderExecution - ) => MetricLabels; -}; - -export type RerankRuntimeContext = { - type: CopilotProviderType; - resolveRerankDriver: () => RerankProviderDriver | undefined; - checkParams: (input: { - cond: ModelFullConditions; - messages?: PromptMessage[]; - embeddings?: string[]; - options?: - | CopilotChatOptions - | CopilotStructuredOptions - | CopilotImageOptions; - withAttachment?: boolean; - execution?: CopilotProviderExecution; - }) => Promise; - selectModel: ( - cond: ModelFullConditions, - execution?: CopilotProviderExecution - ) => CopilotProviderModel; - buildPreparedNativeRerankExecution: ( - protocol: LlmProtocol, - backendConfig: LlmBackendConfig, - model: string, - request: LlmRerankRequest, - execution?: CopilotProviderExecution - ) => PreparedNativeRerankExecution; - createNativeRerankDispatch: ( - backendConfig: LlmBackendConfig, - protocol: LlmProtocol, - execution?: CopilotProviderExecution - ) => (request: LlmRerankRequest) => Promise<{ scores: number[] }>; -}; - -export type ImageRuntimeContext = { - type: CopilotProviderType; - resolveImageDriver: () => ImageProviderDriver | undefined; - checkParams: (input: { - cond: ModelFullConditions; - messages?: PromptMessage[]; - options?: CopilotImageOptions; - withAttachment?: boolean; - execution?: CopilotProviderExecution; - }) => Promise; - selectModel: ( - cond: ModelFullConditions, - execution?: CopilotProviderExecution - ) => CopilotProviderModel; - buildPreparedNativeImageExecution: ( - protocol: LlmProtocol, - backendConfig: LlmBackendConfig, - model: string, - messages: PromptMessage[], - options?: CopilotImageOptions, - execution?: CopilotProviderExecution - ) => PreparedNativeImageExecution; -}; - -type NativeExecutionDriverBase = { - createBackendConfig: ( - execution?: CopilotProviderExecution - ) => Promise | LlmBackendConfig; - mapError: (error: unknown) => unknown; -}; - -type ModelSelectionContext = { - selectModel: ( - cond: ModelFullConditions, - execution?: CopilotProviderExecution - ) => CopilotProviderModel; -}; - -type MetricContext = { - metricLabels: ( - model: string, - labels?: MetricLabels, - execution?: CopilotProviderExecution - ) => MetricLabels; -}; - -type RoutedPreparedExecution = { - route: { - model: string; - backendConfig: LlmBackendConfig; - protocol: LlmProtocol; - }; -}; - -export function resolveDriverOrThrow( - type: CopilotProviderType, - kind: string, - resolveDriver: () => TDriver | undefined -) { - const driver = resolveDriver(); - if (!driver) { - throw new CopilotProviderNotSupported({ - provider: type, - kind, - }); - } - return driver; -} - -export function resolvePreparedModelId( - context: ModelSelectionContext, - cond: ModelConditions, - outputType: ModelOutputType, - prepared?: RoutedPreparedExecution | null -) { - return ( - prepared?.route.model ?? - context.selectModel({ - ...cond, - outputType, - }).id - ); -} - -async function prepareNativeExecutionBase< - TDriver extends NativeExecutionDriverBase, - TPrepared, ->({ - resolveDriver, - cond, - outputType, - checkParams, - selectModel, - execution, - checkInput, - buildPrepared, -}: { - resolveDriver: () => TDriver | undefined; - cond: ModelConditions; - outputType: ModelOutputType; - checkParams: (input: { - cond: ModelFullConditions; - messages?: PromptMessage[]; - embeddings?: string[]; - options?: - | CopilotChatOptions - | CopilotStructuredOptions - | CopilotImageOptions; - withAttachment?: boolean; - execution?: CopilotProviderExecution; - }) => Promise; - selectModel: ( - cond: ModelFullConditions, - execution?: CopilotProviderExecution - ) => CopilotProviderModel; - execution?: CopilotProviderExecution; - checkInput: { - messages?: PromptMessage[]; - embeddings?: string[]; - options?: - | CopilotChatOptions - | CopilotStructuredOptions - | CopilotImageOptions; - withAttachment?: boolean; - }; - buildPrepared: (args: { - driver: TDriver; - model: CopilotProviderModel; - backendConfig: LlmBackendConfig; - protocol: LlmProtocol; - }) => Promise | TPrepared; -}): Promise { - const driver = resolveDriver(); - if (!driver) { - return null; - } - - const normalizedCond = await checkParams({ - ...checkInput, - cond: { ...cond, outputType }, - execution, - }); - const model = selectModel(normalizedCond, execution); - const backendConfig = await driver.createBackendConfig(execution); - const route = resolveProviderModelRoute(model, outputType); - if (!route.protocol) { - throw new Error(`Missing native protocol for model ${model.id}`); - } - - return await buildPrepared({ - driver, - model, - backendConfig: - route.requestLayer === backendConfig.request_layer - ? backendConfig - : { ...backendConfig, request_layer: route.requestLayer }, - protocol: route.protocol, - }); -} - -export async function runPreparedExecution< - TPrepared extends RoutedPreparedExecution, - TResult, ->({ - driver, - prepared, - modelId, - execution, - metricContext, - metricsName, - execute, -}: { - driver: Pick; - prepared: TPrepared | null; - modelId: string; - execution?: CopilotProviderExecution; - metricContext?: MetricContext; - metricsName?: DriverMetricNames; - execute: (prepared: TPrepared) => Promise; -}): Promise { - try { - if (metricsName && metricContext) { - metrics.ai - .counter(metricsName.call) - .add(1, metricContext.metricLabels(modelId, {}, execution)); - } - if (!prepared) { - throw new Error('native route is not available'); - } - return await execute(prepared); - } catch (error) { - if (metricsName && metricContext) { - metrics.ai - .counter(metricsName.error) - .add(1, metricContext.metricLabels(modelId, {}, execution)); - } - throw driver.mapError(error); - } -} - -export async function prepareNativeStructuredExecution( - context: StructuredRuntimeContext, - cond: ModelConditions, - messages: PromptMessage[], - options: CopilotStructuredOptions = {}, - responseContract?: RequiredStructuredOutputContract, - execution?: CopilotProviderExecution -): Promise { - const driver = context.resolveStructuredDriver(); - if (!driver) { - return null; - } - - const structuredOptions = options ?? {}; - const normalizedCond = await context.checkParams({ - messages, - cond: { ...cond, outputType: ModelOutputType.Structured }, - options: structuredOptions, - execution, - }); - const model = context.selectModel(normalizedCond, execution); - const backendConfig = await driver.createBackendConfig(execution); - const route = resolveProviderModelRoute(model, ModelOutputType.Structured); - if (!route.protocol) { - throw new Error(`Missing native protocol for model ${model.id}`); - } - const preparedMessages = driver.prepareMessages - ? await driver.prepareMessages(messages, backendConfig, structuredOptions) - : messages; - if (!responseContract) { - throw new CopilotPromptInvalid('Schema is required'); - } - const { request } = await buildNativeStructuredRequest({ - model: model.id, - messages: preparedMessages, - options: structuredOptions, - responseContract, - attachmentCapability: context.getAttachCapability( - model, - ModelOutputType.Structured - ), - middleware: context.getActiveProviderMiddleware(execution), - }); - - return context.buildPreparedNativeStructuredExecution( - route.protocol, - route.requestLayer === backendConfig.request_layer - ? backendConfig - : { ...backendConfig, request_layer: route.requestLayer }, - model.id, - request, - execution - ); -} - -export async function runNativeStructured( - context: StructuredRuntimeContext, - cond: ModelConditions, - messages: PromptMessage[], - options: CopilotStructuredOptions = {}, - responseContract?: RequiredStructuredOutputContract, - execution?: CopilotProviderExecution -) { - const driver = resolveDriverOrThrow( - context.type, - 'structure', - context.resolveStructuredDriver - ); - const structuredOptions = options ?? {}; - const prepared = await prepareNativeStructuredExecution( - context, - cond, - messages, - structuredOptions, - responseContract, - execution - ); - const modelId = resolvePreparedModelId( - context, - cond, - ModelOutputType.Structured, - prepared - ); - - return await runPreparedExecution({ - driver, - prepared, - modelId, - execution, - metricContext: context, - metricsName: { - call: 'chat_text_calls', - error: 'chat_text_errors', - }, - execute: async preparedExecution => { - const dispatch = context.createNativeStructuredDispatch( - preparedExecution.route.backendConfig, - preparedExecution.route.protocol, - execution - ); - - for (let attempt = 0; ; attempt++) { - try { - const response = await dispatch(preparedExecution.request); - const parsed = parseNativeStructuredOutput(response); - const validated = llmValidateJsonSchema( - preparedExecution.request.schema, - parsed - ); - return JSON.stringify(validated); - } catch (error) { - if ( - !(await driver.shouldRetry?.({ - error, - attempt, - options: structuredOptions, - })) - ) { - throw error; - } - } - } - }, - }); -} - -export async function prepareNativeEmbeddingExecution( - context: EmbeddingRuntimeContext, - cond: ModelConditions, - input: string | string[], - options: CopilotEmbeddingOptions = {}, - execution?: CopilotProviderExecution -): Promise { - const values = Array.isArray(input) ? input : [input]; - return await prepareNativeExecutionBase({ - resolveDriver: context.resolveEmbeddingDriver, - cond, - outputType: ModelOutputType.Embedding, - checkParams: context.checkParams, - selectModel: context.selectModel, - execution, - checkInput: { - embeddings: values, - options, - }, - buildPrepared: ({ driver, model, backendConfig, protocol }) => - context.buildPreparedNativeEmbeddingExecution( - protocol, - backendConfig, - model.id, - buildLlmEmbeddingRequest({ - model: model.id, - inputs: values, - dimensions: options?.dimensions ?? driver.defaultDimensions, - taskType: driver.taskType ?? DEFAULT_EMBEDDING_TASK_TYPE, - }), - execution - ), - }); -} - -export async function runNativeEmbedding( - context: EmbeddingRuntimeContext, - cond: ModelConditions, - input: string | string[], - options?: CopilotEmbeddingOptions, - execution?: CopilotProviderExecution -) { - const driver = resolveDriverOrThrow( - context.type, - ModelOutputType.Embedding, - context.resolveEmbeddingDriver - ); - const prepared = await prepareNativeEmbeddingExecution( - context, - cond, - input, - options, - execution - ); - const modelId = resolvePreparedModelId( - context, - cond, - ModelOutputType.Embedding, - prepared - ); - - return await runPreparedExecution({ - driver, - prepared, - modelId, - execution, - metricContext: context, - metricsName: { - call: 'generate_embedding_calls', - error: 'generate_embedding_errors', - }, - execute: async preparedExecution => { - const response = await context.createNativeEmbeddingDispatch( - preparedExecution.route.backendConfig, - preparedExecution.route.protocol, - execution - )(preparedExecution.request); - return response.embeddings; - }, - }); -} - -export async function prepareNativeRerankExecution( - context: RerankRuntimeContext, - cond: ModelConditions, - request: CopilotRerankRequest, - options: CopilotChatOptions = {}, - execution?: CopilotProviderExecution -): Promise { - return await prepareNativeExecutionBase({ - resolveDriver: context.resolveRerankDriver, - cond, - outputType: ModelOutputType.Rerank, - checkParams: context.checkParams, - selectModel: context.selectModel, - execution, - checkInput: { - messages: [], - options, - }, - buildPrepared: ({ model, backendConfig, protocol }) => - context.buildPreparedNativeRerankExecution( - protocol, - backendConfig, - model.id, - buildLlmRerankRequest(model.id, request), - execution - ), - }); -} - -export async function prepareNativeImageExecution( - context: ImageRuntimeContext, - cond: ModelConditions, - messages: PromptMessage[], - options: CopilotImageOptions = {}, - execution?: CopilotProviderExecution -): Promise { - return await prepareNativeExecutionBase({ - resolveDriver: context.resolveImageDriver, - cond, - outputType: ModelOutputType.Image, - checkParams: context.checkParams, - selectModel: context.selectModel, - execution, - checkInput: { - messages, - options, - }, - buildPrepared: async ({ driver, model, backendConfig, protocol }) => { - const preparedMessages = driver.prepareMessages - ? await driver.prepareMessages(messages, backendConfig, options) - : messages; - - return context.buildPreparedNativeImageExecution( - protocol, - backendConfig, - model.id, - preparedMessages, - options, - execution - ); - }, - }); -} - -export async function runNativeRerank( - context: RerankRuntimeContext, - cond: ModelConditions, - request: CopilotRerankRequest, - options: CopilotChatOptions = {}, - execution?: CopilotProviderExecution -) { - const driver = resolveDriverOrThrow( - context.type, - ModelOutputType.Rerank, - context.resolveRerankDriver - ); - const prepared = await prepareNativeRerankExecution( - context, - cond, - request, - options, - execution - ); - - const modelId = resolvePreparedModelId( - context, - cond, - ModelOutputType.Rerank, - prepared - ); - - return await runPreparedExecution({ - driver, - prepared, - modelId, - execution, - execute: async preparedExecution => { - const response = await context.createNativeRerankDispatch( - preparedExecution.route.backendConfig, - preparedExecution.route.protocol, - execution - )(preparedExecution.request); - return response.scores; - }, - }); -} diff --git a/packages/backend/server/src/plugins/copilot/runtime/provider-runtime-context.ts b/packages/backend/server/src/plugins/copilot/runtime/provider-runtime-context.ts deleted file mode 100644 index 83531994fa..0000000000 --- a/packages/backend/server/src/plugins/copilot/runtime/provider-runtime-context.ts +++ /dev/null @@ -1,490 +0,0 @@ -import type { - LlmBackendConfig, - LlmEmbeddingRequest, - LlmProtocol, - LlmRerankRequest, - LlmStructuredRequest, - LlmStructuredResponse, -} from '../../../native'; -import type { ProviderMiddlewareConfig } from '../config'; -import type { CopilotProvider } from '../providers/provider'; -import type { ProviderModelRuntimeContext } from '../providers/provider-model-runtime'; -import { - createNativeEmbeddingDispatch as inputCreateNativeEmbeddingDispatch, - createNativeRerankDispatch as inputCreateNativeRerankDispatch, - createNativeStructuredDispatch as inputCreateNativeStructuredDispatch, - createPreparedExecutionRuntime, - type CreatePreparedExecutionRuntimeInput, - type PreparedExecutionRuntime, -} from '../providers/provider-native-runtime'; -import type { - CopilotProviderExecution, - EmbeddingProviderDriver, - ImageProviderDriver, - PreparedNativeEmbeddingExecution, - PreparedNativeExecution, - PreparedNativeImageExecution, - PreparedNativeRequestOptions, - PreparedNativeRerankExecution, - PreparedNativeStructuredExecution, - ProviderExecutionDrivers, - ProviderMetricLabels, - ProviderRuntimeHostSeed, - RerankProviderDriver, - StructuredProviderDriver, -} from '../providers/provider-runtime-contract'; -import type { - CopilotChatOptions, - CopilotEmbeddingOptions, - CopilotImageOptions, - CopilotProviderModel, - CopilotRerankRequest, - CopilotStructuredOptions, - ModelAttachmentCapability, - ModelConditions, - ModelFullConditions, - ModelOutputType, - PromptMessage, -} from '../providers/types'; -import type { CopilotToolSet } from '../tools'; -import type { RequiredStructuredOutputContract } from './contracts'; -import type { ChatRuntimeContext } from './provider-chat-runtime'; -import { - prepareNativeChatExecution, - runNativeStreamObject, - runNativeStreamText, - runNativeText, -} from './provider-chat-runtime'; -import type { - EmbeddingRuntimeContext, - ImageRuntimeContext, - RerankRuntimeContext, - StructuredRuntimeContext, -} from './provider-driver-runtime'; -import { - prepareNativeEmbeddingExecution, - prepareNativeImageExecution, - prepareNativeRerankExecution, - prepareNativeStructuredExecution, - runNativeEmbedding, - runNativeRerank, - runNativeStructured, -} from './provider-driver-runtime'; -import type { NativeProviderAdapter } from './tool/native-adapter'; - -type ProviderRuntimeContextInput = { - model: ProviderModelRuntimeContext; - resolveExecutionDrivers: () => ProviderExecutionDrivers | undefined; - selectModel: ( - cond: ModelFullConditions, - execution?: CopilotProviderExecution - ) => CopilotProviderModel; - metricLabels: ( - model: string, - labels?: ProviderMetricLabels, - execution?: CopilotProviderExecution - ) => ProviderMetricLabels; - checkParams: (input: { - cond: ModelFullConditions; - messages?: PromptMessage[]; - embeddings?: string[]; - options?: - | CopilotChatOptions - | CopilotStructuredOptions - | CopilotImageOptions; - withAttachment?: boolean; - execution?: CopilotProviderExecution; - }) => Promise; - getAttachCapability: ( - model: CopilotProviderModel, - outputType: ModelOutputType - ) => ModelAttachmentCapability | undefined; - getActiveProviderMiddleware: ( - execution?: CopilotProviderExecution - ) => ProviderMiddlewareConfig; - getTools: ( - options: CopilotChatOptions, - model: string - ) => Promise; - buildPreparedNativeExecution: ( - options: PreparedNativeRequestOptions - ) => Promise; - createPreparedExecutionAdapter: ( - prepared: PreparedNativeExecution - ) => NativeProviderAdapter; - buildPreparedNativeStructuredExecution: ( - protocol: LlmProtocol, - backendConfig: LlmBackendConfig, - model: string, - request: LlmStructuredRequest, - execution?: CopilotProviderExecution - ) => PreparedNativeStructuredExecution; - createNativeStructuredDispatch: ( - backendConfig: LlmBackendConfig, - protocol: LlmProtocol, - execution?: CopilotProviderExecution - ) => (request: LlmStructuredRequest) => Promise; - buildPreparedNativeEmbeddingExecution: ( - protocol: LlmProtocol, - backendConfig: LlmBackendConfig, - model: string, - request: LlmEmbeddingRequest, - execution?: CopilotProviderExecution - ) => PreparedNativeEmbeddingExecution; - createNativeEmbeddingDispatch: ( - backendConfig: LlmBackendConfig, - protocol: LlmProtocol, - execution?: CopilotProviderExecution - ) => (request: LlmEmbeddingRequest) => Promise<{ embeddings: number[][] }>; - buildPreparedNativeRerankExecution: ( - protocol: LlmProtocol, - backendConfig: LlmBackendConfig, - model: string, - request: LlmRerankRequest, - execution?: CopilotProviderExecution - ) => PreparedNativeRerankExecution; - createNativeRerankDispatch: ( - backendConfig: LlmBackendConfig, - protocol: LlmProtocol, - execution?: CopilotProviderExecution - ) => (request: LlmRerankRequest) => Promise<{ scores: number[] }>; - buildPreparedNativeImageExecution: ( - protocol: LlmProtocol, - backendConfig: LlmBackendConfig, - model: string, - messages: PromptMessage[], - options?: CopilotImageOptions, - execution?: CopilotProviderExecution - ) => PreparedNativeImageExecution; -}; - -export type ProviderRuntimeHostInput = Omit< - ProviderRuntimeContextInput, - | keyof ProviderRuntimeHostSeed - | 'buildPreparedNativeExecution' - | 'createPreparedExecutionAdapter' - | 'buildPreparedNativeStructuredExecution' - | 'buildPreparedNativeEmbeddingExecution' - | 'buildPreparedNativeRerankExecution' - | 'buildPreparedNativeImageExecution' -> & - ProviderRuntimeHostSeed & { - preparedExecutionRuntimeInput: CreatePreparedExecutionRuntimeInput; - createNativeStructuredDispatch: ProviderRuntimeContextInput['createNativeStructuredDispatch']; - createNativeEmbeddingDispatch: ProviderRuntimeContextInput['createNativeEmbeddingDispatch']; - createNativeRerankDispatch: ProviderRuntimeContextInput['createNativeRerankDispatch']; - }; - -type ProviderRuntimeHostOverride = { - overrideRuntimeHost?: ( - runtimeHost: ProviderRuntimeContexts - ) => ProviderRuntimeContexts; -}; - -const runtimeHosts = new WeakMap(); - -export type ProviderRuntimeContexts = { - model: ProviderModelRuntimeContext; - chat: ChatRuntimeContext; - structured: StructuredRuntimeContext; - embedding: EmbeddingRuntimeContext; - rerank: RerankRuntimeContext; - image: ImageRuntimeContext; - prepare: { - chat: ( - kind: 'text' | 'streamText' | 'streamObject', - cond: ModelConditions, - messages: PromptMessage[], - options?: CopilotChatOptions, - execution?: CopilotProviderExecution - ) => ReturnType; - structured: ( - cond: ModelConditions, - messages: PromptMessage[], - options?: CopilotStructuredOptions, - responseContract?: RequiredStructuredOutputContract, - execution?: CopilotProviderExecution - ) => ReturnType; - embedding: ( - cond: ModelConditions, - input: string | string[], - options?: CopilotEmbeddingOptions, - execution?: CopilotProviderExecution - ) => ReturnType; - rerank: ( - cond: ModelConditions, - request: CopilotRerankRequest, - options?: CopilotChatOptions, - execution?: CopilotProviderExecution - ) => ReturnType; - image: ( - cond: ModelConditions, - messages: PromptMessage[], - options?: CopilotImageOptions, - execution?: CopilotProviderExecution - ) => ReturnType; - }; - run: { - text: ( - cond: ModelConditions, - messages: PromptMessage[], - options?: CopilotChatOptions, - execution?: CopilotProviderExecution - ) => ReturnType; - streamText: ( - cond: ModelConditions, - messages: PromptMessage[], - options?: CopilotChatOptions, - execution?: CopilotProviderExecution - ) => ReturnType; - streamObject: ( - cond: ModelConditions, - messages: PromptMessage[], - options?: CopilotChatOptions, - execution?: CopilotProviderExecution - ) => ReturnType; - structured: ( - cond: ModelConditions, - messages: PromptMessage[], - options?: CopilotStructuredOptions, - responseContract?: RequiredStructuredOutputContract, - execution?: CopilotProviderExecution - ) => ReturnType; - embedding: ( - cond: ModelConditions, - input: string | string[], - options?: CopilotEmbeddingOptions, - execution?: CopilotProviderExecution - ) => ReturnType; - rerank: ( - cond: ModelConditions, - request: CopilotRerankRequest, - options?: CopilotChatOptions, - execution?: CopilotProviderExecution - ) => ReturnType; - }; -}; - -function createProviderRuntimeContexts( - input: ProviderRuntimeContextInput -): ProviderRuntimeContexts { - const resolveDriver = ( - kind: K - ): ProviderExecutionDrivers[K] | undefined => - input.resolveExecutionDrivers()?.[kind]; - const chatDriver = resolveDriver('chat'); - - const chatContext: ChatRuntimeContext = { - type: input.model.type, - resolveChatDriver: () => chatDriver, - selectModel: input.selectModel, - metricLabels: input.metricLabels, - createPreparedExecutionAdapter: input.createPreparedExecutionAdapter, - }; - - const structuredContext: StructuredRuntimeContext = { - type: input.model.type, - resolveStructuredDriver: () => - resolveDriver('structured') as StructuredProviderDriver | undefined, - checkParams: input.checkParams, - selectModel: input.selectModel, - getAttachCapability: input.getAttachCapability, - getActiveProviderMiddleware: input.getActiveProviderMiddleware, - buildPreparedNativeStructuredExecution: - input.buildPreparedNativeStructuredExecution, - createNativeStructuredDispatch: input.createNativeStructuredDispatch, - metricLabels: input.metricLabels, - }; - - const embeddingContext: EmbeddingRuntimeContext = { - type: input.model.type, - resolveEmbeddingDriver: () => - resolveDriver('embedding') as EmbeddingProviderDriver | undefined, - checkParams: input.checkParams, - selectModel: input.selectModel, - buildPreparedNativeEmbeddingExecution: - input.buildPreparedNativeEmbeddingExecution, - createNativeEmbeddingDispatch: input.createNativeEmbeddingDispatch, - metricLabels: input.metricLabels, - }; - - const rerankContext: RerankRuntimeContext = { - type: input.model.type, - resolveRerankDriver: () => - resolveDriver('rerank') as RerankProviderDriver | undefined, - checkParams: input.checkParams, - selectModel: input.selectModel, - buildPreparedNativeRerankExecution: - input.buildPreparedNativeRerankExecution, - createNativeRerankDispatch: input.createNativeRerankDispatch, - }; - - const imageContext: ImageRuntimeContext = { - type: input.model.type, - resolveImageDriver: () => - resolveDriver('image') as ImageProviderDriver | undefined, - checkParams: input.checkParams, - selectModel: input.selectModel, - buildPreparedNativeImageExecution: input.buildPreparedNativeImageExecution, - }; - - const prepare: ProviderRuntimeContexts['prepare'] = { - chat: (kind, cond, messages, options = {}, execution) => - prepareNativeChatExecution( - chatContext.resolveChatDriver, - input.buildPreparedNativeExecution, - { - kind, - cond, - messages, - options, - execution, - } - ), - structured: (cond, messages, options = {}, responseContract, execution) => - prepareNativeStructuredExecution( - structuredContext, - cond, - messages, - options, - responseContract, - execution - ), - embedding: (cond, values, options = {}, execution) => - prepareNativeEmbeddingExecution( - embeddingContext, - cond, - values, - options, - execution - ), - rerank: (cond, request, options = {}, execution) => - prepareNativeRerankExecution( - rerankContext, - cond, - request, - options, - execution - ), - image: (cond, messages, options = {}, execution) => - prepareNativeImageExecution( - imageContext, - cond, - messages, - options, - execution - ), - }; - - return { - model: input.model, - chat: chatContext, - structured: structuredContext, - embedding: embeddingContext, - rerank: rerankContext, - image: imageContext, - prepare, - run: { - text: (cond, messages, options, execution) => - runNativeText( - chatContext, - prepare.chat, - cond, - messages, - options, - execution - ), - streamText: (cond, messages, options, execution) => - runNativeStreamText( - chatContext, - prepare.chat, - cond, - messages, - options, - execution - ), - streamObject: (cond, messages, options, execution) => - runNativeStreamObject( - chatContext, - prepare.chat, - cond, - messages, - options, - execution - ), - structured: (cond, messages, options, responseContract, execution) => - runNativeStructured( - structuredContext, - cond, - messages, - options, - responseContract, - execution - ), - embedding: (cond, values, options, execution) => - runNativeEmbedding(embeddingContext, cond, values, options, execution), - rerank: (cond, request, options, execution) => - runNativeRerank(rerankContext, cond, request, options, execution), - }, - }; -} - -export function createProviderRuntimeHost( - input: ProviderRuntimeHostInput -): ProviderRuntimeContexts { - const preparedExecutionRuntime: PreparedExecutionRuntime = - createPreparedExecutionRuntime(input.preparedExecutionRuntimeInput); - - return createProviderRuntimeContexts({ - ...input, - buildPreparedNativeExecution: - preparedExecutionRuntime.buildPreparedNativeExecution, - createPreparedExecutionAdapter: - preparedExecutionRuntime.createPreparedExecutionAdapter, - buildPreparedNativeStructuredExecution: - preparedExecutionRuntime.buildPreparedNativeStructuredExecution, - buildPreparedNativeEmbeddingExecution: - preparedExecutionRuntime.buildPreparedNativeEmbeddingExecution, - buildPreparedNativeRerankExecution: - preparedExecutionRuntime.buildPreparedNativeRerankExecution, - buildPreparedNativeImageExecution: - preparedExecutionRuntime.buildPreparedNativeImageExecution, - createNativeStructuredDispatch: input.createNativeStructuredDispatch, - createNativeEmbeddingDispatch: input.createNativeEmbeddingDispatch, - createNativeRerankDispatch: input.createNativeRerankDispatch, - }); -} - -export function getProviderRuntimeHost( - provider: CopilotProvider -): ProviderRuntimeContexts { - const existingRuntimeHost = runtimeHosts.get(provider); - if (existingRuntimeHost) { - return existingRuntimeHost; - } - const runtimeHostSeed = provider.getRuntimeHostSeed(); - const runtimeHost = createProviderRuntimeHost({ - ...runtimeHostSeed, - preparedExecutionRuntimeInput: { - resolveProviderId: execution => - execution?.providerId ?? `${provider.type}-default`, - getTools: runtimeHostSeed.getTools, - getActiveProviderMiddleware: runtimeHostSeed.getActiveProviderMiddleware, - createNativeAdapter: provider.createNativeAdapter.bind(provider), - maxSteps: provider.maxSteps, - }, - createNativeStructuredDispatch: (backendConfig, protocol, _execution) => - inputCreateNativeStructuredDispatch(backendConfig, protocol), - createNativeEmbeddingDispatch: (backendConfig, protocol, _execution) => - inputCreateNativeEmbeddingDispatch(backendConfig, protocol), - createNativeRerankDispatch: (backendConfig, protocol, _execution) => - inputCreateNativeRerankDispatch(backendConfig, protocol), - }); - const resolvedRuntimeHost = - (provider as ProviderRuntimeHostOverride).overrideRuntimeHost?.( - runtimeHost - ) ?? runtimeHost; - - runtimeHosts.set(provider, resolvedRuntimeHost); - return resolvedRuntimeHost; -} diff --git a/packages/backend/server/src/plugins/copilot/runtime/task-policy.ts b/packages/backend/server/src/plugins/copilot/runtime/task-policy.ts deleted file mode 100644 index 1948e826e8..0000000000 --- a/packages/backend/server/src/plugins/copilot/runtime/task-policy.ts +++ /dev/null @@ -1,35 +0,0 @@ -import { Injectable } from '@nestjs/common'; - -import { QuotaStateService } from '../../../core/quota/state'; -import { PromptService } from '../prompt/service'; - -export const DEFAULT_EMBEDDING_MODEL = 'gemini-embedding-001'; -export const DEFAULT_RERANK_MODEL = 'gpt-4o-mini'; - -@Injectable() -export class TaskPolicy { - constructor( - private readonly quotaState: QuotaStateService, - private readonly prompts: PromptService - ) {} - - resolveEmbeddingModelId() { - return DEFAULT_EMBEDDING_MODEL; - } - - resolveRerankModelId() { - return DEFAULT_RERANK_MODEL; - } - - async resolveTranscriptionModel(userId: string) { - const prompt = await this.prompts.get('Transcript audio'); - if (!prompt) return; - - const state = await this.quotaState.reconcileUserQuotaState(userId); - const flags = state.flags as { unlimitedCopilot?: boolean }; - const hasAccess = - !!flags.unlimitedCopilot || - ['pro', 'lifetime_pro', 'ai'].includes(state.plan); - return prompt.optionalModels[hasAccess ? 1 : 0] ?? prompt.model; - } -} diff --git a/packages/backend/server/src/plugins/copilot/runtime/tool-runtime.ts b/packages/backend/server/src/plugins/copilot/runtime/tool-runtime.ts index 15ebf27808..e4315fe2c9 100644 --- a/packages/backend/server/src/plugins/copilot/runtime/tool-runtime.ts +++ b/packages/backend/server/src/plugins/copilot/runtime/tool-runtime.ts @@ -1,3 +1,4 @@ +/* oxlint-disable import/no-cycle -- Tools can invoke nested prompts and semantic search. */ import { Injectable } from '@nestjs/common'; import { Config } from '../../../base'; @@ -5,7 +6,6 @@ import { DocReader, DocWriter } from '../../../core/doc'; import { PermissionAccess } from '../../../core/permission'; import { Models } from '../../../models'; import { IndexerService } from '../../indexer'; -import type { NodeTextMiddleware } from '../config'; import { CopilotContextService } from '../context/service'; import { type CopilotChatOptions, @@ -36,8 +36,6 @@ import { createSectionEditTool, } from '../tools'; import { PromptRuntime } from './prompt-runtime'; -import type { ToolLoopBackend } from './tool/bridge'; -import { createNativeToolLoopAdapter } from './tool/native-adapter'; export type ProviderSpecificToolResolver = ( toolName: CopilotChatTools, @@ -192,15 +190,4 @@ export class ToolRuntime { return tools; } - - createNativeAdapter( - backend: ToolLoopBackend, - tools: CopilotToolSet, - options: { - maxSteps?: number; - nodeTextMiddleware?: NodeTextMiddleware[]; - } = {} - ) { - return createNativeToolLoopAdapter(backend, tools, options); - } } diff --git a/packages/backend/server/src/plugins/copilot/runtime/tool/bridge.ts b/packages/backend/server/src/plugins/copilot/runtime/tool/bridge.ts index 4c7c65c519..357a12cc35 100644 --- a/packages/backend/server/src/plugins/copilot/runtime/tool/bridge.ts +++ b/packages/backend/server/src/plugins/copilot/runtime/tool/bridge.ts @@ -1,17 +1,8 @@ import { z } from 'zod'; import { - type LlmBackendConfig, - llmDispatchToolLoopStream, - llmDispatchToolLoopStreamPrepared, - llmDispatchToolLoopStreamRouted, - type LlmPreparedDispatchRoute, - type LlmProtocol, - type LlmRequest, - type LlmRoutedBackend, type LlmToolCallbackRequest, type LlmToolCallbackResponse, - type LlmToolLoopStreamEvent, } from '../../../../native'; import type { CopilotTool, @@ -19,71 +10,11 @@ import type { CopilotToolSet, } from '../../tools'; -export type ToolLoopDispatch = ( - request: LlmRequest, - signalOrOptions?: AbortSignal | CopilotToolExecuteOptions, - maybeMessages?: CopilotToolExecuteOptions['messages'] -) => AsyncIterableIterator; - -export type ToolLoopBackend = - | { protocol: LlmProtocol; backendConfig: LlmBackendConfig } - | { routes: LlmRoutedBackend[] } - | { preparedRoutes: LlmPreparedDispatchRoute[] }; - -function normalizeToolExecuteOptions( - signalOrOptions?: AbortSignal | CopilotToolExecuteOptions, - maybeMessages?: CopilotToolExecuteOptions['messages'] -): CopilotToolExecuteOptions { - if ( - signalOrOptions && - typeof signalOrOptions === 'object' && - 'aborted' in signalOrOptions - ) { - return { - signal: signalOrOptions, - messages: maybeMessages, - }; - } - - if (!signalOrOptions) { - return maybeMessages ? { messages: maybeMessages } : {}; - } - - return { - ...signalOrOptions, - signal: signalOrOptions.signal, - messages: signalOrOptions.messages ?? maybeMessages, - }; -} - -export function createToolExecutionCallback( - tools: CopilotToolSet, - options: CopilotToolExecuteOptions = {} -) { - return async (request: LlmToolCallbackRequest) => { - return await executeToolCall(tools, request, options); - }; -} - export async function executeToolCall( tools: CopilotToolSet, request: LlmToolCallbackRequest, options: CopilotToolExecuteOptions ): Promise { - const tool = tools[request.name] as CopilotTool | undefined; - - if (!tool?.execute) { - return { - callId: request.callId, - name: request.name, - args: request.args, - rawArgumentsText: request.rawArgumentsText, - argumentParseError: request.argumentParseError, - isError: true, - output: { message: `Tool not found: ${request.name}` }, - }; - } - if (request.argumentParseError) { return { callId: request.callId, @@ -104,6 +35,19 @@ export async function executeToolCall( }; } + const tool = tools[request.name] as CopilotTool | undefined; + if (!tool?.execute) { + return { + callId: request.callId, + name: request.name, + args: request.args, + rawArgumentsText: request.rawArgumentsText, + argumentParseError: request.argumentParseError, + isError: true, + output: { message: `Tool not found: ${request.name}` }, + }; + } + try { const args = tool.inputSchema instanceof z.ZodType @@ -133,53 +77,5 @@ export async function executeToolCall( } } -export function createToolLoopBridge( - backend: ToolLoopBackend, - tools: CopilotToolSet, - maxSteps = 20 -): ToolLoopDispatch { - return ( - request: LlmRequest, - signalOrOptions?: AbortSignal | CopilotToolExecuteOptions, - maybeMessages?: CopilotToolExecuteOptions['messages'] - ) => { - const toolExecuteOptions = normalizeToolExecuteOptions( - signalOrOptions, - maybeMessages - ); - const execute = createToolExecutionCallback(tools, toolExecuteOptions); - const toolLoopRequest = { ...request, stream: true }; - - if ('routes' in backend) { - return llmDispatchToolLoopStreamRouted( - backend.routes, - toolLoopRequest, - execute, - maxSteps, - toolExecuteOptions.signal - ); - } - - if ('preparedRoutes' in backend) { - return llmDispatchToolLoopStreamPrepared( - backend.preparedRoutes, - execute, - maxSteps, - toolExecuteOptions.signal - ); - } - - return llmDispatchToolLoopStream( - backend.protocol, - backend.backendConfig, - toolLoopRequest, - execute, - maxSteps, - toolExecuteOptions.signal - ); - }; -} - -// re-export for test consumers export type { LlmToolCallbackRequest } from '../../../../native'; export type { CopilotToolExecuteOptions, CopilotToolSet } from '../../tools'; diff --git a/packages/backend/server/src/plugins/copilot/runtime/tool/native-adapter.ts b/packages/backend/server/src/plugins/copilot/runtime/tool/native-adapter.ts index c196a1836b..4190e750e1 100644 --- a/packages/backend/server/src/plugins/copilot/runtime/tool/native-adapter.ts +++ b/packages/backend/server/src/plugins/copilot/runtime/tool/native-adapter.ts @@ -7,9 +7,7 @@ import { CitationFootnoteFormatter, TextStreamParser, } from '../../providers/utils'; -import type { CopilotToolSet } from '../../tools'; import { projectRuntimeEventToStreamObject } from '../contracts/runtime-event-contract'; -import { createToolLoopBridge, type ToolLoopBackend } from './bridge'; import { type EnrichedToolCallEvent, type EnrichedToolResultEvent, @@ -425,14 +423,3 @@ export class NativeProviderAdapter { } } } - -export function createNativeToolLoopAdapter( - backend: ToolLoopBackend, - tools: CopilotToolSet, - options: NativeProviderAdapterOptions = {} -) { - return new NativeProviderAdapter( - createToolLoopBridge(backend, tools, options.maxSteps), - options - ); -} diff --git a/packages/backend/server/src/plugins/copilot/runtime/turn-orchestrator.ts b/packages/backend/server/src/plugins/copilot/runtime/turn-orchestrator.ts index 1f759af1e9..715452a9ff 100644 --- a/packages/backend/server/src/plugins/copilot/runtime/turn-orchestrator.ts +++ b/packages/backend/server/src/plugins/copilot/runtime/turn-orchestrator.ts @@ -3,14 +3,15 @@ import { Injectable } from '@nestjs/common'; import { CopilotContextService } from '../context/service'; import { type Turn } from '../core'; import { + type ModelConditions, ModelInputType, type PromptParams, type StreamObject, } from '../providers/types'; import { ChatSession } from '../session'; import { ChatQuerySchema } from '../types'; +import { getTools } from '../utils'; import { CapabilityRuntime } from './capability-runtime'; -import { CapabilityPolicyHost } from './hosts/capability-policy-host'; import { ConversationHost } from './hosts/conversation-host'; import { ImageResultHost } from './hosts/image-result-host'; import { TurnPersistence } from './hosts/turn-persistence'; @@ -20,7 +21,6 @@ export class TurnOrchestrator { constructor( private readonly conversations: ConversationHost, private readonly context: CopilotContextService, - private readonly capabilityPolicy: CapabilityPolicyHost, private readonly runtime: CapabilityRuntime, private readonly imageResults: ImageResultHost, private readonly turnPersistence: TurnPersistence @@ -62,8 +62,15 @@ export class TurnOrchestrator { sessionId, query ); - const { modelId, reasoning, webSearch, toolsConfig, byokLeaseId } = - ChatQuerySchema.parse(query); + const { + profileId, + modelId, + routeTargetId, + reasoning, + webSearch, + toolsConfig, + byokLeaseId, + } = ChatQuerySchema.parse(query); const promptParams = await this.buildPromptParams(sessionId, { latestTurn: prepared.latestTurn, includeContextFiles: selection.includeContextFiles, @@ -76,22 +83,34 @@ export class TurnOrchestrator { return { prepared, finalMessage, - selection: await this.capabilityPolicy.selectChat(prepared.session, { - responseMode: selection.responseMode, - modelId, - reasoning, - webSearch, - toolsConfig, - byokLeaseId, - billingUnitId: prepared.latestTurn?.id, - quotaBackedRoutesAllowed: prepared.quotaBackedRoutesAllowed, - featureKind: - selection.responseMode === 'image' - ? 'image' - : selection.responseMode === 'object' - ? 'action' - : 'chat', - }), + selection: { + model: profileId && modelId ? modelId : 'route-selected', + conditions: { profileId, modelId }, + providerOptions: { + ...prepared.session.config.promptConfig, + user: prepared.session.config.userId, + session: prepared.session.config.sessionId, + workspace: prepared.session.config.workspaceId, + profileId, + byokLeaseId, + billingUnitId: prepared.latestTurn?.id, + builtInRouteId: prepared.session.config.promptName, + managedTargetId: routeTargetId, + quotaBackedRoutesAllowed: prepared.quotaBackedRoutesAllowed, + featureKind: + selection.responseMode === 'image' + ? 'image' + : selection.responseMode === 'object' + ? 'action' + : 'chat', + reasoning, + webSearch, + tools: getTools( + prepared.session.config.promptConfig?.tools, + toolsConfig + ), + }, + }, }; } @@ -110,7 +129,7 @@ export class TurnOrchestrator { const stream = this.streamTextResult( prepared.session, - selection.model, + selection.conditions, finalMessage, { ...selection.providerOptions, @@ -129,14 +148,14 @@ export class TurnOrchestrator { private async *streamTextResult( session: ChatSession, - model: string, + conditions: ModelConditions, finalMessage: ReturnType, options: Record, wasAborted: () => boolean ) { let buffer = ''; for await (const chunk of this.runtime.streamText( - { modelId: model }, + conditions, finalMessage, options )) { @@ -165,7 +184,7 @@ export class TurnOrchestrator { finalMessage, stream: this.streamObjectResult( prepared.session, - selection.model, + selection.conditions, finalMessage, { ...selection.providerOptions, @@ -178,14 +197,14 @@ export class TurnOrchestrator { private async *streamObjectResult( session: ChatSession, - model: string, + conditions: ModelConditions, finalMessage: ReturnType, options: Record, wasAborted: () => boolean ): AsyncIterableIterator { const chunks: StreamObject[] = []; for await (const chunk of this.runtime.streamObject( - { modelId: model }, + conditions, finalMessage, options )) { @@ -223,7 +242,7 @@ export class TurnOrchestrator { userId, sessionId, prepared.session, - undefined, + selection.conditions, hasAttachment, finalMessage, { @@ -244,7 +263,7 @@ export class TurnOrchestrator { userId: string, sessionId: string, session: ChatSession, - model: string | undefined, + conditions: ModelConditions, hasAttachment: boolean, finalMessage: ReturnType, options: Record, @@ -253,7 +272,7 @@ export class TurnOrchestrator { const attachments: string[] = []; for await (const artifact of this.runtime.streamImageArtifacts( { - modelId: model, + ...conditions, inputTypes: hasAttachment ? [ModelInputType.Image] : [ModelInputType.Text], diff --git a/packages/backend/server/src/plugins/copilot/session.ts b/packages/backend/server/src/plugins/copilot/session.ts index d537bb4055..862dba76dc 100644 --- a/packages/backend/server/src/plugins/copilot/session.ts +++ b/packages/backend/server/src/plugins/copilot/session.ts @@ -2,7 +2,7 @@ import { randomUUID } from 'node:crypto'; import { Injectable, Logger } from '@nestjs/common'; import { Transactional } from '@nestjs-cls/transactional'; -import { AiPromptRole } from '@prisma/client'; +import { AiSessionMessageRole } from '@prisma/client'; import { CopilotActionTaken, @@ -20,7 +20,6 @@ import { type UpdateChatSession, UpdateChatSessionOptions, } from '../../models'; -import { CopilotAccessPolicy } from './access'; import { ConversationPolicy } from './conversation/policy'; import { ConversationStore } from './conversation/store'; import { type Conversation, promptMessageFromTurn, type Turn } from './core'; @@ -50,7 +49,6 @@ export class ChatSession implements AsyncDisposable { prompt: ResolvedPrompt, turns: PromptMessage[], params: PromptParams, - maxTokenSize: number, sessionId?: string ) => PromptMessage[]; constructor( @@ -59,23 +57,13 @@ export class ChatSession implements AsyncDisposable { prompt: ResolvedPrompt, turns: PromptMessage[], params: PromptParams, - maxTokenSize: number, sessionId?: string ) => PromptMessage[], - private readonly dispose?: (state: ChatSessionState) => Promise, - private readonly maxTokenSize = state.prompt.config?.maxTokens || 128 * 1024 + private readonly dispose?: (state: ChatSessionState) => Promise ) { this.renderPromptSession = renderPromptSession; } - get model() { - return this.state.prompt.model; - } - - get optionalModels() { - return this.state.prompt.optionalModels; - } - get config() { const { sessionId, @@ -126,7 +114,7 @@ export class ChatSession implements AsyncDisposable { revertLatestMessage(removeLatestUserMessage: boolean) { const turns = this.state.turns; turns.splice( - turns.findLastIndex(({ role }) => role === AiPromptRole.user) + + turns.findLastIndex(({ role }) => role === AiSessionMessageRole.user) + (removeLatestUserMessage ? 0 : 1) ); } @@ -136,7 +124,6 @@ export class ChatSession implements AsyncDisposable { this.state.prompt, this.state.turns.map(turn => promptMessageFromTurn(turn)), params, - this.maxTokenSize, this.state.sessionId ); } @@ -158,13 +145,11 @@ export type ConversationState = { conversation: Conversation; turns: Turn[]; prompt: ResolvedPrompt; - tokenCost: number; }; export type ConversationMetaState = { conversation: Conversation; prompt: ResolvedPrompt; - tokenCost: number; }; type StoredConversation = NonNullable< @@ -183,7 +168,6 @@ export class ChatSessionService { private readonly models: Models, private readonly jobs: JobQueue, private readonly store: ConversationStore, - private readonly access: CopilotAccessPolicy, private readonly conversationPolicy: ConversationPolicy, private readonly prompts: PromptService, private readonly promptRuntime: PromptRuntime @@ -206,14 +190,13 @@ export class ChatSessionService { private async toConversationState( session: StoredConversation ): Promise { - const { conversation, prompt, tokenCost } = + const { conversation, prompt } = await this.toConversationMetaState(session); return { conversation, turns: session.turns, prompt, - tokenCost, }; } @@ -226,7 +209,6 @@ export class ChatSessionService { return { conversation: session.conversation, prompt, - tokenCost: session.tokenCost, }; } @@ -296,11 +278,11 @@ export class ChatSessionService { } async getQuota(userId: string) { - return await this.access.getQuota(userId); + return await this.conversationPolicy.getQuota(userId); } async checkQuota(userId: string) { - await this.access.checkQuota(userId); + await this.conversationPolicy.checkQuota(userId); } async create(options: ChatSessionOptions): Promise { @@ -360,7 +342,6 @@ export class ChatSessionService { ); finalData.promptName = prompt.name; finalData.promptAction = prompt.action ?? null; - finalData.promptModel = prompt.model; } finalData.pinned = options.pinned; finalData.docId = options.docId; @@ -389,7 +370,8 @@ export class ChatSessionService { if (options.latestMessageId) { const lastMessageIdx = state.turns.findLastIndex( ({ id, role }) => - role === AiPromptRole.assistant && id === options.latestMessageId + role === AiSessionMessageRole.assistant && + id === options.latestMessageId ); if (lastMessageIdx < 0) { throw new CopilotMessageNotFound({ @@ -410,7 +392,6 @@ export class ChatSessionService { prompt: { name: state.prompt.name, action: state.prompt.action, - model: state.prompt.model, }, turns, }); @@ -434,7 +415,6 @@ export class ChatSessionService { async appendTurn(input: { sessionId: string; userId: string; - prompt: { model: string }; turn: Turn; compatSubmissionId?: string; }) { @@ -485,14 +465,8 @@ export class ChatSessionService { turns: state.turns, prompt: state.prompt, }, - (prompt, turns, params, maxTokenSize, sessionId) => - this.prompts.renderSession( - prompt, - turns, - params, - maxTokenSize, - sessionId - ), + (prompt, turns, params, sessionId) => + this.prompts.renderSession(prompt, turns, params, sessionId), async state => { await this.store.appendTurns(state); if (this.conversationPolicy.shouldScheduleTitle(state.prompt)) { @@ -538,9 +512,18 @@ export class ChatSessionService { const promptContent = this.conversationPolicy.buildTitlePromptContent(turns); const generatedTitle = this.stripNullBytes( - await this.promptRuntime.runText('Summary as title', { - content: promptContent, - }) + await this.promptRuntime.runText( + 'Summary as title', + { content: promptContent }, + { + providerOptions: { + user: conversation.userId, + workspace: conversation.workspaceId, + featureKind: 'chat', + quotaBackedRoutesAllowed: true, + }, + } + ) ).trim(); if (!generatedTitle) { diff --git a/packages/backend/server/src/plugins/copilot/tools/doc-semantic-search.ts b/packages/backend/server/src/plugins/copilot/tools/doc-semantic-search.ts index 5d39e58dcd..e7e9f55c3b 100644 --- a/packages/backend/server/src/plugins/copilot/tools/doc-semantic-search.ts +++ b/packages/backend/server/src/plugins/copilot/tools/doc-semantic-search.ts @@ -1,3 +1,4 @@ +/* oxlint-disable import/no-cycle -- Semantic search uses the shared embedding runtime. */ import { omit } from 'lodash-es'; import { z } from 'zod'; diff --git a/packages/backend/server/src/plugins/copilot/tools/index.ts b/packages/backend/server/src/plugins/copilot/tools/index.ts index f4e20e3e76..8bd1774970 100644 --- a/packages/backend/server/src/plugins/copilot/tools/index.ts +++ b/packages/backend/server/src/plugins/copilot/tools/index.ts @@ -1,3 +1,4 @@ +/* oxlint-disable import/no-cycle -- Tool exports include semantic search runtime dependencies. */ export * from './blob-read'; export * from './code-artifact'; export * from './conversation-summary'; diff --git a/packages/backend/server/src/plugins/copilot/transcript/realtime.ts b/packages/backend/server/src/plugins/copilot/transcript/realtime.ts index 1b3db913e5..70222b9f49 100644 --- a/packages/backend/server/src/plugins/copilot/transcript/realtime.ts +++ b/packages/backend/server/src/plugins/copilot/transcript/realtime.ts @@ -1,13 +1,15 @@ import { Injectable, OnModuleInit } from '@nestjs/common'; import { z } from 'zod'; -import { CopilotTranscriptionJobNotFound } from '../../../base'; +import { Config } from '../../../base/config'; +import { CopilotTranscriptionJobNotFound } from '../../../base/error/errors.gen'; import { PermissionAccess } from '../../../core/permission'; import { RealtimeRegistry, realtimeTranscriptTaskRoom, registerRealtimeLiveQuery, } from '../../../core/realtime'; +import { assertCopilotEnabled } from '../availability'; import { CopilotTranscriptionReader } from './reader'; @Injectable() @@ -15,7 +17,8 @@ export class CopilotTranscriptRealtimeProvider implements OnModuleInit { constructor( private readonly ac: PermissionAccess, private readonly transcript: CopilotTranscriptionReader, - private readonly registry: RealtimeRegistry + private readonly registry: RealtimeRegistry, + private readonly config: Config ) {} onModuleInit() { @@ -68,6 +71,7 @@ export class CopilotTranscriptRealtimeProvider implements OnModuleInit { } private async assertCopilot(userId: string, workspaceId: string) { + assertCopilotEnabled(this.config); await this.ac .user(userId) .workspace(workspaceId) diff --git a/packages/backend/server/src/plugins/copilot/transcript/resolver.ts b/packages/backend/server/src/plugins/copilot/transcript/resolver.ts index 941fe5aa43..a9cb7e0eb8 100644 --- a/packages/backend/server/src/plugins/copilot/transcript/resolver.ts +++ b/packages/backend/server/src/plugins/copilot/transcript/resolver.ts @@ -22,6 +22,7 @@ import { } from '../../../base'; import { CurrentUser } from '../../../core/auth'; import { PermissionAccess } from '../../../core/permission'; +import { CopilotEnabled } from '../feature'; import { CopilotType } from '../resolver'; import type { TranscriptionJob } from './job'; import { buildLegacyProjection } from './projection'; @@ -36,7 +37,6 @@ import type { TranscriptionQuality, TranscriptionSourceAudio, TranscriptionSubmitInput, - TranscriptProviderMeta, } from './types'; registerEnumType(AiJobStatus, { @@ -166,15 +166,6 @@ class TranscriptionQualityType implements TranscriptionQuality { overflowCount!: number | null; } -@ObjectType() -class TranscriptProviderMetaType implements TranscriptProviderMeta { - @Field(() => String, { nullable: true }) - provider!: string | null; - - @Field(() => String, { nullable: true }) - model!: string | null; -} - @InputType() class AudioSliceManifestItemInput implements AudioSliceManifestItem { @Field(() => Int) @@ -233,9 +224,6 @@ class SubmitAudioTranscriptionInput implements TranscriptionSubmitInput { @Field(() => [AudioSliceManifestItemInput], { nullable: true }) sliceManifest?: AudioSliceManifestItemInput[]; - - @Field(() => String, { nullable: true }) - strategy?: string | null; } @ObjectType() @@ -273,15 +261,9 @@ class TranscriptionResultType { @Field(() => MeetingSummaryV2Type, { nullable: true }) summaryJson!: TranscriptionPayload['summaryJson'] | null; - @Field(() => TranscriptProviderMetaType, { nullable: true }) - providerMeta!: TranscriptionPayload['providerMeta'] | null; - @Field(() => String, { nullable: true }) version!: string | null; - @Field(() => String, { nullable: true }) - strategy!: string | null; - @Field(() => AiJobStatus) status!: AiJobStatus; } @@ -292,6 +274,7 @@ const FinishedStatus: Set = new Set([ ]); @Injectable() +@CopilotEnabled() @Resolver(() => CopilotType) export class CopilotTranscriptionResolver { constructor( @@ -318,9 +301,7 @@ export class CopilotTranscriptionResolver { normalizedSegments: null, normalizedTranscript: null, summaryJson: null, - providerMeta: null, version: null, - strategy: null, }; if (FinishedStatus.has(finalJob.status)) { finalJob.title = legacy?.title ?? null; @@ -333,9 +314,7 @@ export class CopilotTranscriptionResolver { finalJob.normalizedSegments = ret?.normalizedSegments ?? null; finalJob.normalizedTranscript = ret?.normalizedTranscript ?? null; finalJob.summaryJson = ret?.summaryJson ?? null; - finalJob.providerMeta = ret?.providerMeta ?? null; finalJob.version = ret?.version ?? null; - finalJob.strategy = ret?.strategy ?? null; } return finalJob; } diff --git a/packages/backend/server/src/plugins/copilot/transcript/schema.ts b/packages/backend/server/src/plugins/copilot/transcript/schema.ts index 9185001192..6bb8b21198 100644 --- a/packages/backend/server/src/plugins/copilot/transcript/schema.ts +++ b/packages/backend/server/src/plugins/copilot/transcript/schema.ts @@ -73,11 +73,6 @@ export const TranscriptionQualitySchema = z.object({ overflowCount: z.number().nullable().optional(), }); -export const TranscriptProviderMetaSchema = z.object({ - provider: z.string().nullable().optional(), - model: z.string().nullable().optional(), -}); - export const TranscriptionLegacyProjectionSchema = z.object({ title: z.string().nullable().optional(), summary: z.string().nullable().optional(), @@ -96,9 +91,7 @@ export const TranscriptionPayloadV2Schema = z.object({ .optional(), normalizedTranscript: z.string().nullable().optional(), summaryJson: MeetingSummaryV2Schema.nullable().optional(), - providerMeta: TranscriptProviderMetaSchema.nullable().optional(), version: z.string().optional(), - strategy: z.string().optional(), }); export const TranscriptionSubmitInputSchema = TranscriptionPayloadV2Schema.pick( @@ -135,9 +128,7 @@ const CanonicalTranscriptPayloadSchema = TranscriptionPayloadV2Schema.refine( payload.normalizedSegments !== undefined || payload.normalizedTranscript !== undefined || payload.summaryJson !== undefined || - payload.providerMeta !== undefined || - payload.version !== undefined || - payload.strategy !== undefined, + payload.version !== undefined, { message: 'canonical transcript payload must contain canonical transcript fields', diff --git a/packages/backend/server/src/plugins/copilot/transcript/service.ts b/packages/backend/server/src/plugins/copilot/transcript/service.ts index 10d59f9e2f..2a6557cd59 100644 --- a/packages/backend/server/src/plugins/copilot/transcript/service.ts +++ b/packages/backend/server/src/plugins/copilot/transcript/service.ts @@ -14,11 +14,9 @@ import { realtimeTranscriptTaskRoom, } from '../../../core/realtime'; import { Models } from '../../../models'; -import { CopilotAccessPolicy } from '../access'; import { PromptService } from '../prompt'; -import { CopilotProviderType } from '../providers/types'; import { ActionRuntimeBridge } from '../runtime/action-runtime-bridge'; -import { TaskPolicy } from '../runtime/task-policy'; +import { CapabilityRuntime } from '../runtime/capability-runtime'; import { CopilotStorage } from '../storage'; import { taskToJob, type TranscriptionJob } from './job'; import { @@ -32,9 +30,9 @@ import type { } from './types'; import { readStream } from './utils'; -const TRANSCRIPT_ACTION_ID = 'transcript.audio.gemini'; +const TRANSCRIPT_ACTION_ID = 'transcript.audio'; +const TRANSCRIPT_PROMPT_REF = 'Transcript audio structured'; const TRANSCRIPT_ACTION_VERSION = 'v1'; -const TRANSCRIPT_STRATEGY = 'gemini'; @Injectable() export class CopilotTranscriptionService { @@ -42,10 +40,9 @@ export class CopilotTranscriptionService { private readonly models: Models, private readonly job: JobQueue, private readonly storage: CopilotStorage, - private readonly tasks: TaskPolicy, private readonly prompts: PromptService, private readonly actionBridge: ActionRuntimeBridge, - private readonly access: CopilotAccessPolicy, + private readonly runtime: CapabilityRuntime, private readonly realtime: RealtimePublisher ) {} @@ -58,27 +55,10 @@ export class CopilotTranscriptionService { sourceAudio: payload.sourceAudio, quality: payload.quality, sliceManifest: payload.sliceManifest, - providerMeta: payload.providerMeta, version: 'transcript-result-v1', - strategy: TRANSCRIPT_STRATEGY, }; } - private async resolveTranscriptStrategy(userId: string, strategy?: string) { - if (strategy && strategy !== TRANSCRIPT_STRATEGY) { - throw new BadRequestException( - `Transcript strategy ${strategy} is not available` - ); - } - const model = await this.tasks.resolveTranscriptionModel(userId); - if (!model) { - throw new BadRequestException( - 'Transcript strategy gemini is not available' - ); - } - return { model, strategy: TRANSCRIPT_STRATEGY }; - } - private async persistUploads( userId: string, workspaceId: string, @@ -123,11 +103,8 @@ export class CopilotTranscriptionService { } satisfies TranscriptionPayloadV2; } - private async buildTranscriptActionMessages( - payload: TranscriptionPayloadV2, - modelId?: string - ) { - const prompt = await this.prompts.get('Transcript audio structured'); + private async buildTranscriptActionMessages(payload: TranscriptionPayloadV2) { + const prompt = await this.prompts.get(TRANSCRIPT_PROMPT_REF); if (!prompt) { throw new Error('Transcript action prompt not found'); } @@ -140,10 +117,6 @@ export class CopilotTranscriptionService { mimeType: info.mimeType, index: info.index ?? null, })) ?? null, - providerMeta: { - provider: CopilotProviderType.Gemini, - model: modelId ?? payload.providerMeta?.model ?? null, - }, }; const attachments = payload.infos?.map(info => ({ @@ -165,7 +138,7 @@ export class CopilotTranscriptionService { workspaceId: string, blobId: string, blobs: FileUpload[], - input?: TranscriptionSubmitInput & { strategy?: string | null } + input?: TranscriptionSubmitInput ): Promise { const existingTask = await this.models.copilotTranscriptTask.getWithUser( userId, @@ -180,15 +153,15 @@ export class CopilotTranscriptionService { throw new CopilotTranscriptionJobExists(); } - await this.access.assertQuotaOrByok({ - userId, - workspaceId, - featureKind: 'transcript', - }); - - const { model, strategy } = await this.resolveTranscriptStrategy( - userId, - input?.strategy ?? undefined + await this.runtime.assertRoute( + 'transcript.audio', + {}, + { + user: userId, + workspace: workspaceId, + featureKind: 'transcript', + builtInRouteId: TRANSCRIPT_PROMPT_REF, + } ); const infos = await this.persistUploads(userId, workspaceId, blobId, blobs); const payload = this.createCanonicalPayload(blobId, infos, input); @@ -196,7 +169,6 @@ export class CopilotTranscriptionService { userId, workspaceId, blobId, - strategy, recipeId: TRANSCRIPT_ACTION_ID, recipeVersion: TRANSCRIPT_ACTION_VERSION, inputSnapshot: payload, @@ -206,7 +178,6 @@ export class CopilotTranscriptionService { await this.job.add('copilot.transcript.task.submit', { taskId: task.id, payload, - modelId: model, }); await this.models.copilotTranscriptTask.markRunning(task.id); this.publishTaskChanged(workspaceId, task.id, AiJobStatus.running); @@ -234,21 +205,20 @@ export class CopilotTranscriptionService { ); } - await this.access.assertQuotaOrByok({ - userId, - workspaceId, - featureKind: 'transcript', - }); - const payload = this.parseTaskPayload(task.protectedResult); - const { model } = await this.resolveTranscriptStrategy( - userId, - task.strategy + await this.runtime.assertRoute( + 'transcript.audio', + {}, + { + user: userId, + workspace: workspaceId, + featureKind: 'transcript', + builtInRouteId: TRANSCRIPT_PROMPT_REF, + } ); await this.job.add('copilot.transcript.task.submit', { taskId, payload, - modelId: model, retryOf: task.actionRunId ?? undefined, }); await this.models.copilotTranscriptTask.markRunning(taskId); @@ -282,12 +252,6 @@ export class CopilotTranscriptionService { return taskToJob(task); } - await this.access.assertQuotaOrByok({ - userId, - workspaceId, - featureKind: 'transcript', - }); - const settled = await this.models.copilotTranscriptTask.settle(task.id); return taskToJob(settled); } @@ -311,7 +275,6 @@ export class CopilotTranscriptionService { async transcriptTask({ taskId, payload, - modelId, retryOf, }: Jobs['copilot.transcript.task.submit']) { const task = await this.models.copilotTranscriptTask.get(taskId); @@ -324,10 +287,7 @@ export class CopilotTranscriptionService { let bridgeFailed = false; let bridgeError = 'transcript native recipe failed'; let finalResult: unknown = null; - const messages = await this.buildTranscriptActionMessages( - payload, - modelId - ); + const messages = await this.buildTranscriptActionMessages(payload); for await (const event of this.actionBridge.runStream({ userId: task.userId, workspaceId: task.workspaceId, @@ -335,14 +295,6 @@ export class CopilotTranscriptionService { actionVersion: TRANSCRIPT_ACTION_VERSION, retryOf: retryOf ?? null, inputSnapshot: payload, - nativeInput: { - input: { - sourceAudio: payload.sourceAudio ?? null, - quality: payload.quality ?? null, - infos: payload.infos ?? null, - sliceManifest: payload.sliceManifest ?? null, - }, - }, onRunCreated: async ({ runId }) => { await this.models.copilotTranscriptTask.markRunning(taskId, runId); this.publishTaskChanged( @@ -351,9 +303,9 @@ export class CopilotTranscriptionService { AiJobStatus.running ); }, - prepareStructuredRoutes: { - stepId: 'transcribe', - modelId, + step: { + slot: 'transcript.audio', + builtInRouteId: TRANSCRIPT_PROMPT_REF, messages, options: { user: task.userId, @@ -362,7 +314,6 @@ export class CopilotTranscriptionService { billingUnitId: taskId, featureKind: 'transcript', }, - prefer: CopilotProviderType.Gemini, responseContract: TranscriptActionResultContract, }, })) { diff --git a/packages/backend/server/src/plugins/copilot/transcript/types.ts b/packages/backend/server/src/plugins/copilot/transcript/types.ts index cbdd85749b..38115dfc71 100644 --- a/packages/backend/server/src/plugins/copilot/transcript/types.ts +++ b/packages/backend/server/src/plugins/copilot/transcript/types.ts @@ -15,7 +15,6 @@ import { TranscriptionQualitySchema, TranscriptionSourceAudioSchema, TranscriptionSubmitInputSchema, - TranscriptProviderMetaSchema, } from './schema'; export type LegacyTranscriptionSegment = z.infer< @@ -36,9 +35,6 @@ export type TranscriptionSourceAudio = z.infer< typeof TranscriptionSourceAudioSchema >; export type TranscriptionQuality = z.infer; -export type TranscriptProviderMeta = z.infer< - typeof TranscriptProviderMetaSchema ->; export type TranscriptionLegacyProjection = z.infer< typeof TranscriptionLegacyProjectionSchema >; @@ -57,7 +53,6 @@ declare global { 'copilot.transcript.task.submit': { taskId: string; payload: TranscriptionPayloadV2; - modelId?: string; retryOf?: string; }; } diff --git a/packages/backend/server/src/plugins/copilot/types.ts b/packages/backend/server/src/plugins/copilot/types.ts index 352166d154..8e9d5bc48b 100644 --- a/packages/backend/server/src/plugins/copilot/types.ts +++ b/packages/backend/server/src/plugins/copilot/types.ts @@ -36,7 +36,9 @@ export type ToolsConfig = z.infer; export const ChatQuerySchema = z .object({ messageId: zMaybeString, + profileId: zMaybeString, modelId: zMaybeString, + routeTargetId: zMaybeString, byokLeaseId: zMaybeString, retry: zBool, reasoning: zBool, @@ -44,10 +46,29 @@ export const ChatQuerySchema = z toolsConfig: ToolsConfigSchema, }) .catchall(z.string()) + .superRefine((value, context) => { + if (!!value.profileId !== !!value.modelId) { + context.addIssue({ + code: 'custom', + message: 'profileId and modelId must be provided together', + }); + } + for (const field of ['requirements', 'deployment', 'profiles', 'presets']) { + if (Object.hasOwn(value, field)) { + context.addIssue({ + code: 'custom', + path: [field], + message: `${field} is owned by the native route policy`, + }); + } + } + }) .transform( ({ messageId, + profileId, modelId, + routeTargetId, byokLeaseId, retry, reasoning, @@ -56,7 +77,9 @@ export const ChatQuerySchema = z ...params }) => ({ messageId, + profileId, modelId, + routeTargetId, byokLeaseId, retry, reasoning, @@ -85,11 +108,7 @@ export const ChatHistorySchema = z title: z.string().nullable(), action: z.string().nullable(), - model: z.string(), - optionalModels: z.array(z.string()), promptName: z.string(), - - tokens: z.number(), messages: z.array(ChatMessageSchema), createdAt: z.date(), updatedAt: z.date(), diff --git a/packages/backend/server/src/plugins/copilot/workspace/resolver.ts b/packages/backend/server/src/plugins/copilot/workspace/resolver.ts index 3f4e6b436d..61eee5b852 100644 --- a/packages/backend/server/src/plugins/copilot/workspace/resolver.ts +++ b/packages/backend/server/src/plugins/copilot/workspace/resolver.ts @@ -26,6 +26,7 @@ import { import { CurrentUser } from '../../../core/auth'; import { PermissionAccess } from '../../../core/permission'; import { WorkspaceType } from '../../../core/workspaces'; +import { CopilotEnabled } from '../feature'; import { COPILOT_LOCKER } from '../resolver'; import { MAX_EMBEDDABLE_SIZE } from '../utils'; import { CopilotWorkspaceService } from './service'; @@ -47,6 +48,7 @@ export class CopilotWorkspaceConfigType { * Public apis rate limit: 10 req/m * Other rate limit: 120 req/m */ +@CopilotEnabled() @Resolver(() => WorkspaceType) export class CopilotWorkspaceEmbeddingResolver { constructor(private readonly ac: PermissionAccess) {} @@ -67,6 +69,7 @@ export class CopilotWorkspaceEmbeddingResolver { } } +@CopilotEnabled() @Resolver(() => CopilotWorkspaceConfigType) export class CopilotWorkspaceEmbeddingConfigResolver { constructor( diff --git a/packages/backend/server/src/schema.gql b/packages/backend/server/src/schema.gql index 27015cb17d..735b6e5af7 100644 --- a/packages/backend/server/src/schema.gql +++ b/packages/backend/server/src/schema.gql @@ -347,17 +347,6 @@ type BlobUploadedPart { partNumber: Int! } -enum ByokKeyStorage { - local - server -} - -enum ByokKeyTestStatus { - failed - passed - untested -} - enum ByokProvider { anthropic fal @@ -553,12 +542,12 @@ type Copilot { contexts(contextId: String, sessionId: String): [CopilotContext!]! histories(docId: String, options: QueryChatHistoriesInput): [CopilotHistories!]! @deprecated(reason: "use `chats` instead") - """List available models for a prompt, with human-readable names""" - models(promptName: String!): CopilotModelsType! - """Get the quota of the user in the workspace""" quota: CopilotQuota! + """List native built-in route choices for a prompt""" + routeOptions(promptName: String!): CopilotRouteOptions + """Get the session by id""" session(sessionId: String!): CopilotSessionType! @@ -664,16 +653,11 @@ type CopilotHistories { createdAt: DateTime! docId: String messages: [ChatMessage!]! - model: String! - optionalModels: [String!]! parentSessionId: String pinned: Boolean! promptName: String! sessionId: String! title: String - - """The number of tokens used in the session""" - tokens: Int! updatedAt: DateTime! workspaceId: String! } @@ -691,17 +675,6 @@ type CopilotMessageNotFoundDataType { messageId: String! } -type CopilotModelType { - id: String! - name: String! -} - -type CopilotModelsType { - defaultModel: String! - optionalModels: [CopilotModelType!]! - proModels: [CopilotModelType!]! -} - type CopilotPromptNotFoundDataType { name: String! } @@ -722,11 +695,22 @@ type CopilotQuota { used: SafeInt! } +type CopilotRouteOptions { + choices: [CopilotRouteTarget!]! + defaultTargetId: String + routeId: String! +} + +type CopilotRouteTarget { + available: Boolean! + displayName: String! + id: String! + minimumTier: String! +} + type CopilotSessionType { docId: String id: ID! - model: String! - optionalModels: [String!]! parentSessionId: ID pinned: Boolean! promptName: String! @@ -821,13 +805,12 @@ input CreateWorkspaceByokLocalLeaseInput { } input CreateWorkspaceByokLocalLeaseProviderInput { - apiKey: String! + credential: String! + definition: WorkspaceByokProfileDefinitionInput! description: String - enabled: Boolean - endpoint: String + enabled: Boolean! name: String! provider: ByokProvider! - sortOrder: SafeInt } type CreateWorkspaceByokLocalLeaseResultType { @@ -835,6 +818,16 @@ type CreateWorkspaceByokLocalLeaseResultType { leaseId: String! } +input CreateWorkspaceByokProfileInput { + credential: String! + definition: WorkspaceByokProfileDefinitionInput! + description: String + enabled: Boolean! + name: String! + provider: ByokProvider! + workspaceId: String! +} + type CredentialsRequirementType { password: PasswordLimitsType! } @@ -1644,7 +1637,6 @@ type Mutation { """Cleanup sessions""" cleanupCopilotSession(options: DeleteSessionInput!): [String!]! - clearWorkspaceByokConfigs(provider: ByokProvider, workspaceId: String!): Boolean! completeBlobUpload(key: String!, parts: [BlobUploadPartInput!], uploadId: String, workspaceId: String!): String! createBlobUpload(key: String!, mime: String!, size: Int!, workspaceId: String!): BlobUploadInit! @@ -1680,6 +1672,7 @@ type Mutation { """Create a new workspace""" createWorkspace(init: Upload): WorkspaceType! createWorkspaceByokLocalLease(input: CreateWorkspaceByokLocalLeaseInput!): CreateWorkspaceByokLocalLeaseResultType! + createWorkspaceByokProfile(input: CreateWorkspaceByokProfileInput!): WorkspaceByokProfileType! deactivateLicense(workspaceId: String!): Boolean! deleteAccount: DeleteAccount! deleteAuthSigningKey(id: String!): [AuthSigningKeyType!]! @@ -1694,7 +1687,7 @@ type Mutation { """Delete a user account""" deleteUser(id: String!): DeleteAccount! deleteWorkspace(id: String!): Boolean! - deleteWorkspaceByokConfig(id: ID!, workspaceId: String!): Boolean! + deleteWorkspaceByokProfile(profileId: ID!, workspaceId: String!): Boolean! """Reenable an banned user""" enableUser(id: String!): UserType! @@ -1717,6 +1710,8 @@ type Mutation { """mention user in a doc""" mentionUser(input: MentionInput!): ID! previewLicense(license: Upload!): AdminLicensePreview! + probeWorkspaceByokDraft(input: ProbeWorkspaceByokDraftInput!): WorkspaceByokProbeResultType! + probeWorkspaceByokProfile(input: ProbeWorkspaceByokProfileInput!): WorkspaceByokProbeResultType! publishDoc(docId: String!, mode: PublicDocMode = Page, workspaceId: String!): DocType! """queue workspace doc embedding""" @@ -1750,7 +1745,8 @@ type Mutation { """Remove workspace embedding files""" removeWorkspaceEmbeddingFiles(fileId: String!, workspaceId: String!): Boolean! - reorderWorkspaceByokConfigs(input: ReorderWorkspaceByokConfigsInput!): [WorkspaceByokKeyConfigType!]! + reorderWorkspaceByokProfiles(input: ReorderWorkspaceByokProfilesInput!): [WorkspaceByokProfileType!]! + replaceWorkspaceByokProfile(input: ReplaceWorkspaceByokProfileInput!): WorkspaceByokProfileType! """Request to apply the subscription in advance""" requestApplySubscription(transactionId: String!): [SubscriptionType!]! @@ -1767,6 +1763,7 @@ type Mutation { revokePublicDoc(docId: String!, workspaceId: String!): DocType! rotateAuthSigningKey(expectedActiveKeyId: String!): [AuthSigningKeyType!]! rotateMcpCredential(expirationDays: Int! = 90, id: ID!, workspaceId: String!): RevealedMcpCredentialType! + rotateWorkspaceByokCredential(input: RotateWorkspaceByokCredentialInput!): WorkspaceByokProfileType! sendChangeEmail(callbackUrl: String!): Boolean! sendChangePasswordEmail(callbackUrl: String!, email: String @deprecated(reason: "fetched from signed in user")): Boolean! sendSetPasswordEmail(callbackUrl: String!, email: String @deprecated(reason: "fetched from signed in user")): Boolean! @@ -1776,7 +1773,6 @@ type Mutation { setBlob(blob: Upload!, workspaceId: String!): String! settleTranscriptTask(taskId: String!, workspaceId: String!): TranscriptionResultType submitTranscriptTask(blob: Upload, blobId: String!, blobs: [Upload!], input: SubmitAudioTranscriptionInput, workspaceId: String!): TranscriptionResultType - testWorkspaceByokConfig(input: TestWorkspaceByokConfigInput!): TestWorkspaceByokConfigResultType! unlinkCalendarAccount(accountId: String!): Boolean! """update app configuration""" @@ -1817,7 +1813,6 @@ type Mutation { """Upload a comment attachment and return the access url""" uploadCommentAttachment(attachment: Upload!, docId: String!, workspaceId: String!): String! - upsertWorkspaceByokConfig(input: UpsertWorkspaceByokConfigInput!): WorkspaceByokKeyConfigType! verifyEmail(token: String!): Boolean! } @@ -2004,6 +1999,22 @@ enum Permission { Owner } +input ProbeWorkspaceByokDraftInput { + checks: [WorkspaceByokProbeCheckInput!]! + credential: String + definition: WorkspaceByokProfileDefinitionInput! + expectedRevision: SafeInt + profileId: ID + provider: ByokProvider! + workspaceId: String! +} + +input ProbeWorkspaceByokProfileInput { + checks: [WorkspaceByokProbeCheckInput!]! + profileId: ID! + workspaceId: String! +} + """The mode which the public doc default in""" enum PublicDocMode { Edgeless @@ -2143,9 +2154,19 @@ input RemoveContextFileInput { fileId: String! } -input ReorderWorkspaceByokConfigsInput { - ids: [ID!]! - storage: ByokKeyStorage! +input ReorderWorkspaceByokProfilesInput { + profiles: [WorkspaceByokProfileOrderInput!]! + workspaceId: String! +} + +input ReplaceWorkspaceByokProfileInput { + credential: String + definition: WorkspaceByokProfileDefinitionInput! + description: String + enabled: Boolean! + expectedRevision: SafeInt! + name: String! + profileId: ID! workspaceId: String! } @@ -2199,6 +2220,13 @@ input RevokeDocUserRoleInput { workspaceId: String! } +input RotateWorkspaceByokCredentialInput { + credential: String! + expectedRevision: SafeInt! + profileId: ID! + workspaceId: String! +} + type RuntimeConfigNotFoundDataType { key: String! } @@ -2386,7 +2414,6 @@ input SubmitAudioTranscriptionInput { quality: TranscriptionQualityInput sliceManifest: [AudioSliceManifestItemInput!] sourceAudio: TranscriptionSourceAudioInput - strategy: String } type SubscriptionAlreadyExistsDataType { @@ -2473,21 +2500,6 @@ enum SubscriptionVariant { Onetime } -input TestWorkspaceByokConfigInput { - apiKey: String - configId: ID - endpoint: String - provider: ByokProvider! - storage: ByokKeyStorage! - workspaceId: String! -} - -type TestWorkspaceByokConfigResultType { - message: String - ok: Boolean! - status: ByokKeyTestStatus! -} - enum TimeBucket { Day Hour @@ -2503,11 +2515,6 @@ type TimeWindow { to: DateTime! } -type TranscriptProviderMetaType { - model: String - provider: String -} - type TranscriptionItemType { end: String! speaker: String! @@ -2530,12 +2537,10 @@ type TranscriptionResultType { id: ID! normalizedSegments: [NormalizedTranscriptSegmentType!] normalizedTranscript: String - providerMeta: TranscriptProviderMetaType quality: TranscriptionQualityType sliceManifest: [AudioSliceManifestItemType!] sourceAudio: TranscriptionSourceAudioType status: AiJobStatus! - strategy: String summary: String summaryJson: MeetingSummaryV2Type title: String @@ -2650,19 +2655,6 @@ input UpdateWorkspaceInput { """The `Upload` scalar type represents a file upload.""" scalar Upload -input UpsertWorkspaceByokConfigInput { - apiKey: String - description: String - enabled: Boolean - endpoint: String - id: ID - name: String! - provider: ByokProvider! - sortOrder: SafeInt - storage: ByokKeyStorage! - workspaceId: String! -} - type UserImportFailedType { email: String! error: String! @@ -2760,45 +2752,128 @@ type VersionRejectedDataType { version: String! } -type WorkspaceByokCapabilityWarningType { - featureKind: String! - reason: String! - requiredProviders: [ByokProvider!]! +input WorkspaceByokCapabilityInput { + attachmentKinds: [String!]! + attachmentSources: [String!]! + features: [String!]! + input: [String!]! + output: [String!]! } -type WorkspaceByokKeyConfigType { - capabilities: [String!]! - configured: Boolean! - description: String - disabledReason: String - enabled: Boolean! - endpoint: String - endpointEditable: Boolean! - id: ID! - lastError: String - lastErrorAt: DateTime - lastTestError: String - lastTestedAt: DateTime - lastUsedAt: DateTime - name: String! +type WorkspaceByokCapabilityType { + attachmentKinds: [String!]! + attachmentSources: [String!]! + features: [String!]! + input: [String!]! + output: [String!]! +} + +type WorkspaceByokCatalogModelType { + capabilities: [WorkspaceByokCapabilityType!]! + displayName: String! + modelId: String! + recommended: Boolean! +} + +type WorkspaceByokCatalogProviderType { + models: [WorkspaceByokCatalogModelType!]! provider: ByokProvider! +} + +type WorkspaceByokCatalogType { + providers: [WorkspaceByokCatalogProviderType!]! + version: String! +} + +input WorkspaceByokEndpointInput { + kind: String! + url: String +} + +type WorkspaceByokEndpointType { + kind: String! + url: String +} + +input WorkspaceByokModelDeclarationInput { + capabilities: [WorkspaceByokCapabilityInput!]! + enabled: Boolean! + modelId: String! +} + +type WorkspaceByokModelDeclarationType { + capabilities: [WorkspaceByokCapabilityType!]! + enabled: Boolean! + modelId: String! +} + +type WorkspaceByokModelProbeCheckType { + operation: String! + status: WorkspaceByokProbeStatusType! +} + +type WorkspaceByokModelProbeType { + checks: [WorkspaceByokModelProbeCheckType!]! + modelId: String! +} + +input WorkspaceByokProbeCheckInput { + modelId: String! + operation: String! +} + +type WorkspaceByokProbeResultType { + connection: WorkspaceByokProbeStatusType! + definitionFingerprint: String! + models: [WorkspaceByokModelProbeType!]! + stale: Boolean! +} + +type WorkspaceByokProbeStatusType { + errorKind: String + kind: String! + testedAt: DateTime +} + +input WorkspaceByokProfileDefinitionInput { + endpoint: WorkspaceByokEndpointInput! + models: [WorkspaceByokModelDeclarationInput!]! + version: SafeInt! +} + +type WorkspaceByokProfileDefinitionType { + endpoint: WorkspaceByokEndpointType! + models: [WorkspaceByokModelDeclarationType!]! + version: SafeInt! +} + +input WorkspaceByokProfileOrderInput { + expectedRevision: SafeInt! + profileId: ID! +} + +type WorkspaceByokProfileType { + definition: WorkspaceByokProfileDefinitionType! + description: String + enabled: Boolean! + name: String! + profileId: ID! + provider: ByokProvider! + revision: SafeInt! sortOrder: SafeInt! - storage: ByokKeyStorage! - testStatus: ByokKeyTestStatus! + validation: WorkspaceByokValidationType + workspaceId: String! } type WorkspaceByokSettingsType { allowedProviders: [ByokProvider!]! + catalog: WorkspaceByokCatalogType! customEndpointSupported: Boolean! entitled: Boolean! - entitlementRequired: [String!]! - hasAiPlan: Boolean! - keys: [WorkspaceByokKeyConfigType!]! localEntitled: Boolean! - localStorageSupported: Boolean! privateEndpointSupported: Boolean! + profiles: [WorkspaceByokProfileType!]! serverEntitled: Boolean! - warnings: [WorkspaceByokCapabilityWarningType!]! workspaceId: String! } @@ -2808,6 +2883,13 @@ type WorkspaceByokUsagePointType { totalTokens: SafeInt! } +type WorkspaceByokValidationType { + connection: WorkspaceByokProbeStatusType! + credentialGeneration: SafeInt! + definitionFingerprint: String! + models: [WorkspaceByokModelProbeType!]! +} + input WorkspaceCalendarItemInput { colorOverride: String sortOrder: Int diff --git a/packages/common/graphql/src/graphql/copilot-models-get.gql b/packages/common/graphql/src/graphql/copilot-models-get.gql deleted file mode 100644 index 4c15fa3bf6..0000000000 --- a/packages/common/graphql/src/graphql/copilot-models-get.gql +++ /dev/null @@ -1,17 +0,0 @@ -query getPromptModels($promptName: String!) { - currentUser { - copilot { - models(promptName: $promptName) { - defaultModel - optionalModels { - id - name - } - proModels { - id - name - } - } - } - } -} diff --git a/packages/common/graphql/src/graphql/copilot-route-options-get.gql b/packages/common/graphql/src/graphql/copilot-route-options-get.gql new file mode 100644 index 0000000000..528ab9b575 --- /dev/null +++ b/packages/common/graphql/src/graphql/copilot-route-options-get.gql @@ -0,0 +1,16 @@ +query getCopilotRouteOptions($promptName: String!) { + currentUser { + copilot { + routeOptions(promptName: $promptName) { + routeId + defaultTargetId + choices { + id + displayName + minimumTier + available + } + } + } + } +} diff --git a/packages/common/graphql/src/graphql/fragments/copilot-chat-history.gql b/packages/common/graphql/src/graphql/fragments/copilot-chat-history.gql index eed7dc06f9..aa524b3a25 100644 --- a/packages/common/graphql/src/graphql/fragments/copilot-chat-history.gql +++ b/packages/common/graphql/src/graphql/fragments/copilot-chat-history.gql @@ -4,12 +4,9 @@ fragment CopilotChatHistory on CopilotHistories { docId parentSessionId promptName - model - optionalModels action pinned title - tokens messages { id role diff --git a/packages/common/graphql/src/graphql/index.ts b/packages/common/graphql/src/graphql/index.ts index 7ff82dbe18..7dc3653c88 100644 --- a/packages/common/graphql/src/graphql/index.ts +++ b/packages/common/graphql/src/graphql/index.ts @@ -12,12 +12,9 @@ export const copilotChatHistoryFragment = `fragment CopilotChatHistory on Copilo docId parentSessionId promptName - model - optionalModels action pinned title - tokens messages { id role @@ -1432,28 +1429,6 @@ export const createCopilotMessageMutation = { file: true, }; -export const getPromptModelsQuery = { - id: 'getPromptModelsQuery' as const, - op: 'getPromptModels', - query: `query getPromptModels($promptName: String!) { - currentUser { - copilot { - models(promptName: $promptName) { - defaultModel - optionalModels { - id - name - } - proModels { - id - name - } - } - } - } -}`, -}; - export const copilotQuotaQuery = { id: 'copilotQuotaQuery' as const, op: 'copilotQuota', @@ -1469,6 +1444,27 @@ export const copilotQuotaQuery = { }`, }; +export const getCopilotRouteOptionsQuery = { + id: 'getCopilotRouteOptionsQuery' as const, + op: 'getCopilotRouteOptions', + query: `query getCopilotRouteOptions($promptName: String!) { + currentUser { + copilot { + routeOptions(promptName: $promptName) { + routeId + defaultTargetId + choices { + id + displayName + minimumTier + available + } + } + } + } +}`, +}; + export const cleanupCopilotSessionMutation = { id: 'cleanupCopilotSessionMutation' as const, op: 'cleanupCopilotSession', @@ -2985,51 +2981,94 @@ export const workspaceBlobQuotaQuery = { }`, }; -export const clearWorkspaceByokConfigsMutation = { - id: 'clearWorkspaceByokConfigsMutation' as const, - op: 'clearWorkspaceByokConfigs', - query: `mutation clearWorkspaceByokConfigs($workspaceId: String!) { - clearWorkspaceByokConfigs(workspaceId: $workspaceId) +export const deleteWorkspaceByokProfileMutation = { + id: 'deleteWorkspaceByokProfileMutation' as const, + op: 'deleteWorkspaceByokProfile', + query: `mutation deleteWorkspaceByokProfile($workspaceId: String!, $profileId: ID!) { + deleteWorkspaceByokProfile(workspaceId: $workspaceId, profileId: $profileId) }`, }; -export const deleteWorkspaceByokConfigMutation = { - id: 'deleteWorkspaceByokConfigMutation' as const, - op: 'deleteWorkspaceByokConfig', - query: `mutation deleteWorkspaceByokConfig($workspaceId: String!, $id: ID!) { - deleteWorkspaceByokConfig(workspaceId: $workspaceId, id: $id) -}`, -}; - -export const reorderWorkspaceByokConfigsMutation = { - id: 'reorderWorkspaceByokConfigsMutation' as const, - op: 'reorderWorkspaceByokConfigs', - query: `mutation reorderWorkspaceByokConfigs($input: ReorderWorkspaceByokConfigsInput!) { - reorderWorkspaceByokConfigs(input: $input) { - id - sortOrder +export const probeWorkspaceByokProfileMutation = { + id: 'probeWorkspaceByokProfileMutation' as const, + op: 'probeWorkspaceByokProfile', + query: `mutation probeWorkspaceByokProfile($input: ProbeWorkspaceByokProfileInput!) { + probeWorkspaceByokProfile(input: $input) { + definitionFingerprint + stale + connection { + kind + testedAt + errorKind + } + models { + modelId + checks { + operation + status { + kind + testedAt + errorKind + } + } + } } }`, }; -export const testWorkspaceByokConfigMutation = { - id: 'testWorkspaceByokConfigMutation' as const, - op: 'testWorkspaceByokConfig', - query: `mutation testWorkspaceByokConfig($input: TestWorkspaceByokConfigInput!) { - testWorkspaceByokConfig(input: $input) { - ok - status - message +export const probeWorkspaceByokDraftMutation = { + id: 'probeWorkspaceByokDraftMutation' as const, + op: 'probeWorkspaceByokDraft', + query: `mutation probeWorkspaceByokDraft($input: ProbeWorkspaceByokDraftInput!) { + probeWorkspaceByokDraft(input: $input) { + definitionFingerprint + stale + connection { + kind + testedAt + errorKind + } + models { + modelId + checks { + operation + status { + kind + testedAt + errorKind + } + } + } } }`, }; -export const upsertWorkspaceByokConfigMutation = { - id: 'upsertWorkspaceByokConfigMutation' as const, - op: 'upsertWorkspaceByokConfig', - query: `mutation upsertWorkspaceByokConfig($input: UpsertWorkspaceByokConfigInput!) { - upsertWorkspaceByokConfig(input: $input) { - id +export const createWorkspaceByokProfileMutation = { + id: 'createWorkspaceByokProfileMutation' as const, + op: 'createWorkspaceByokProfile', + query: `mutation createWorkspaceByokProfile($input: CreateWorkspaceByokProfileInput!) { + createWorkspaceByokProfile(input: $input) { + profileId + } +}`, +}; + +export const replaceWorkspaceByokProfileMutation = { + id: 'replaceWorkspaceByokProfileMutation' as const, + op: 'replaceWorkspaceByokProfile', + query: `mutation replaceWorkspaceByokProfile($input: ReplaceWorkspaceByokProfileInput!) { + replaceWorkspaceByokProfile(input: $input) { + profileId + } +}`, +}; + +export const rotateWorkspaceByokCredentialMutation = { + id: 'rotateWorkspaceByokCredentialMutation' as const, + op: 'rotateWorkspaceByokCredential', + query: `mutation rotateWorkspaceByokCredential($input: RotateWorkspaceByokCredentialInput!) { + rotateWorkspaceByokCredential(input: $input) { + profileId } }`, }; @@ -3045,6 +3084,18 @@ export const createWorkspaceByokLocalLeaseMutation = { }`, }; +export const reorderWorkspaceByokProfilesMutation = { + id: 'reorderWorkspaceByokProfilesMutation' as const, + op: 'reorderWorkspaceByokProfiles', + query: `mutation reorderWorkspaceByokProfiles($input: ReorderWorkspaceByokProfilesInput!) { + reorderWorkspaceByokProfiles(input: $input) { + profileId + sortOrder + revision + } +}`, +}; + export const workspaceByokSettingsQuery = { id: 'workspaceByokSettingsQuery' as const, op: 'workspaceByokSettings', @@ -3056,36 +3107,73 @@ export const workspaceByokSettingsQuery = { entitled serverEntitled localEntitled - entitlementRequired allowedProviders - localStorageSupported customEndpointSupported privateEndpointSupported - hasAiPlan - keys { - id + catalog { + version + providers { + provider + models { + modelId + displayName + recommended + capabilities { + input + output + features + attachmentKinds + attachmentSources + } + } + } + } + profiles { + profileId provider name description - storage - configured enabled - endpoint - endpointEditable sortOrder - capabilities - testStatus - disabledReason - lastTestedAt - lastTestError - lastUsedAt - lastErrorAt - lastError - } - warnings { - featureKind - reason - requiredProviders + revision + definition { + version + endpoint { + kind + url + } + models { + modelId + enabled + capabilities { + input + output + features + attachmentKinds + attachmentSources + } + } + } + validation { + definitionFingerprint + credentialGeneration + connection { + kind + testedAt + errorKind + } + models { + modelId + checks { + operation + status { + kind + testedAt + errorKind + } + } + } + } } } byokUsage(from: $from, to: $to) { diff --git a/packages/common/graphql/src/graphql/workspace-byok-config-clear.gql b/packages/common/graphql/src/graphql/workspace-byok-config-clear.gql deleted file mode 100644 index cc2318163d..0000000000 --- a/packages/common/graphql/src/graphql/workspace-byok-config-clear.gql +++ /dev/null @@ -1,3 +0,0 @@ -mutation clearWorkspaceByokConfigs($workspaceId: String!) { - clearWorkspaceByokConfigs(workspaceId: $workspaceId) -} diff --git a/packages/common/graphql/src/graphql/workspace-byok-config-delete.gql b/packages/common/graphql/src/graphql/workspace-byok-config-delete.gql index 9ef1704e70..54bf9d3642 100644 --- a/packages/common/graphql/src/graphql/workspace-byok-config-delete.gql +++ b/packages/common/graphql/src/graphql/workspace-byok-config-delete.gql @@ -1,3 +1,3 @@ -mutation deleteWorkspaceByokConfig($workspaceId: String!, $id: ID!) { - deleteWorkspaceByokConfig(workspaceId: $workspaceId, id: $id) +mutation deleteWorkspaceByokProfile($workspaceId: String!, $profileId: ID!) { + deleteWorkspaceByokProfile(workspaceId: $workspaceId, profileId: $profileId) } diff --git a/packages/common/graphql/src/graphql/workspace-byok-config-reorder.gql b/packages/common/graphql/src/graphql/workspace-byok-config-reorder.gql deleted file mode 100644 index 9ebeaef8d5..0000000000 --- a/packages/common/graphql/src/graphql/workspace-byok-config-reorder.gql +++ /dev/null @@ -1,8 +0,0 @@ -mutation reorderWorkspaceByokConfigs( - $input: ReorderWorkspaceByokConfigsInput! -) { - reorderWorkspaceByokConfigs(input: $input) { - id - sortOrder - } -} diff --git a/packages/common/graphql/src/graphql/workspace-byok-config-test.gql b/packages/common/graphql/src/graphql/workspace-byok-config-test.gql index a83299a947..afb929c373 100644 --- a/packages/common/graphql/src/graphql/workspace-byok-config-test.gql +++ b/packages/common/graphql/src/graphql/workspace-byok-config-test.gql @@ -1,7 +1,23 @@ -mutation testWorkspaceByokConfig($input: TestWorkspaceByokConfigInput!) { - testWorkspaceByokConfig(input: $input) { - ok - status - message +mutation probeWorkspaceByokProfile($input: ProbeWorkspaceByokProfileInput!) { + probeWorkspaceByokProfile(input: $input) { + definitionFingerprint + stale + connection { kind testedAt errorKind } + models { + modelId + checks { operation status { kind testedAt errorKind } } + } + } +} + +mutation probeWorkspaceByokDraft($input: ProbeWorkspaceByokDraftInput!) { + probeWorkspaceByokDraft(input: $input) { + definitionFingerprint + stale + connection { kind testedAt errorKind } + models { + modelId + checks { operation status { kind testedAt errorKind } } + } } } diff --git a/packages/common/graphql/src/graphql/workspace-byok-config-upsert.gql b/packages/common/graphql/src/graphql/workspace-byok-config-upsert.gql index 83404e27b7..27be33444b 100644 --- a/packages/common/graphql/src/graphql/workspace-byok-config-upsert.gql +++ b/packages/common/graphql/src/graphql/workspace-byok-config-upsert.gql @@ -1,5 +1,11 @@ -mutation upsertWorkspaceByokConfig($input: UpsertWorkspaceByokConfigInput!) { - upsertWorkspaceByokConfig(input: $input) { - id - } +mutation createWorkspaceByokProfile($input: CreateWorkspaceByokProfileInput!) { + createWorkspaceByokProfile(input: $input) { profileId } +} + +mutation replaceWorkspaceByokProfile($input: ReplaceWorkspaceByokProfileInput!) { + replaceWorkspaceByokProfile(input: $input) { profileId } +} + +mutation rotateWorkspaceByokCredential($input: RotateWorkspaceByokCredentialInput!) { + rotateWorkspaceByokCredential(input: $input) { profileId } } diff --git a/packages/common/graphql/src/graphql/workspace-byok-profile-reorder.gql b/packages/common/graphql/src/graphql/workspace-byok-profile-reorder.gql new file mode 100644 index 0000000000..b49773f61a --- /dev/null +++ b/packages/common/graphql/src/graphql/workspace-byok-profile-reorder.gql @@ -0,0 +1,7 @@ +mutation reorderWorkspaceByokProfiles($input: ReorderWorkspaceByokProfilesInput!) { + reorderWorkspaceByokProfiles(input: $input) { + profileId + sortOrder + revision + } +} diff --git a/packages/common/graphql/src/graphql/workspace-byok-settings.gql b/packages/common/graphql/src/graphql/workspace-byok-settings.gql index 68e66f09f8..addb773fe6 100644 --- a/packages/common/graphql/src/graphql/workspace-byok-settings.gql +++ b/packages/common/graphql/src/graphql/workspace-byok-settings.gql @@ -6,36 +6,59 @@ query workspaceByokSettings($id: String!, $from: DateTime!, $to: DateTime!) { entitled serverEntitled localEntitled - entitlementRequired allowedProviders - localStorageSupported customEndpointSupported privateEndpointSupported - hasAiPlan - keys { - id + catalog { + version + providers { + provider + models { + modelId + displayName + recommended + capabilities { + input + output + features + attachmentKinds + attachmentSources + } + } + } + } + profiles { + profileId provider name description - storage - configured enabled - endpoint - endpointEditable sortOrder - capabilities - testStatus - disabledReason - lastTestedAt - lastTestError - lastUsedAt - lastErrorAt - lastError - } - warnings { - featureKind - reason - requiredProviders + revision + definition { + version + endpoint { kind url } + models { + modelId + enabled + capabilities { + input + output + features + attachmentKinds + attachmentSources + } + } + } + validation { + definitionFingerprint + credentialGeneration + connection { kind testedAt errorKind } + models { + modelId + checks { operation status { kind testedAt errorKind } } + } + } } } byokUsage(from: $from, to: $to) { diff --git a/packages/common/graphql/src/schema.ts b/packages/common/graphql/src/schema.ts index 5a1bffae09..8d0bb6fae6 100644 --- a/packages/common/graphql/src/schema.ts +++ b/packages/common/graphql/src/schema.ts @@ -411,17 +411,6 @@ export interface BlobUploadedPart { partNumber: Scalars['Int']['output']; } -export enum ByokKeyStorage { - local = 'local', - server = 'server', -} - -export enum ByokKeyTestStatus { - failed = 'failed', - passed = 'passed', - untested = 'untested', -} - export enum ByokProvider { anthropic = 'anthropic', fal = 'fal', @@ -618,10 +607,10 @@ export interface Copilot { contexts: Array; /** @deprecated use `chats` instead */ histories: Array; - /** List available models for a prompt, with human-readable names */ - models: CopilotModelsType; /** Get the quota of the user in the workspace */ quota: CopilotQuota; + /** List native built-in route choices for a prompt */ + routeOptions: Maybe; /** Get the session by id */ session: CopilotSessionType; /** @@ -650,7 +639,7 @@ export interface CopilotHistoriesArgs { options?: InputMaybe; } -export interface CopilotModelsArgs { +export interface CopilotRouteOptionsArgs { promptName: Scalars['String']['input']; } @@ -785,15 +774,11 @@ export interface CopilotHistories { createdAt: Scalars['DateTime']['output']; docId: Maybe; messages: Array; - model: Scalars['String']['output']; - optionalModels: Array; parentSessionId: Maybe; pinned: Scalars['Boolean']['output']; promptName: Scalars['String']['output']; sessionId: Scalars['String']['output']; title: Maybe; - /** The number of tokens used in the session */ - tokens: Scalars['Int']['output']; updatedAt: Scalars['DateTime']['output']; workspaceId: Scalars['String']['output']; } @@ -814,19 +799,6 @@ export interface CopilotMessageNotFoundDataType { messageId: Scalars['String']['output']; } -export interface CopilotModelType { - __typename?: 'CopilotModelType'; - id: Scalars['String']['output']; - name: Scalars['String']['output']; -} - -export interface CopilotModelsType { - __typename?: 'CopilotModelsType'; - defaultModel: Scalars['String']['output']; - optionalModels: Array; - proModels: Array; -} - export interface CopilotPromptNotFoundDataType { __typename?: 'CopilotPromptNotFoundDataType'; name: Scalars['String']['output']; @@ -851,12 +823,25 @@ export interface CopilotQuota { used: Scalars['SafeInt']['output']; } +export interface CopilotRouteOptions { + __typename?: 'CopilotRouteOptions'; + choices: Array; + defaultTargetId: Maybe; + routeId: Scalars['String']['output']; +} + +export interface CopilotRouteTarget { + __typename?: 'CopilotRouteTarget'; + available: Scalars['Boolean']['output']; + displayName: Scalars['String']['output']; + id: Scalars['String']['output']; + minimumTier: Scalars['String']['output']; +} + export interface CopilotSessionType { __typename?: 'CopilotSessionType'; docId: Maybe; id: Scalars['ID']['output']; - model: Scalars['String']['output']; - optionalModels: Array; parentSessionId: Maybe; pinned: Scalars['Boolean']['output']; promptName: Scalars['String']['output']; @@ -964,13 +949,12 @@ export interface CreateWorkspaceByokLocalLeaseInput { } export interface CreateWorkspaceByokLocalLeaseProviderInput { - apiKey: Scalars['String']['input']; + credential: Scalars['String']['input']; + definition: WorkspaceByokProfileDefinitionInput; description?: InputMaybe; - enabled?: InputMaybe; - endpoint?: InputMaybe; + enabled: Scalars['Boolean']['input']; name: Scalars['String']['input']; provider: ByokProvider; - sortOrder?: InputMaybe; } export interface CreateWorkspaceByokLocalLeaseResultType { @@ -979,6 +963,16 @@ export interface CreateWorkspaceByokLocalLeaseResultType { leaseId: Scalars['String']['output']; } +export interface CreateWorkspaceByokProfileInput { + credential: Scalars['String']['input']; + definition: WorkspaceByokProfileDefinitionInput; + description?: InputMaybe; + enabled: Scalars['Boolean']['input']; + name: Scalars['String']['input']; + provider: ByokProvider; + workspaceId: Scalars['String']['input']; +} + export interface CredentialsRequirementType { __typename?: 'CredentialsRequirementType'; password: PasswordLimitsType; @@ -1861,7 +1855,6 @@ export interface Mutation { changePassword: Scalars['Boolean']['output']; /** Cleanup sessions */ cleanupCopilotSession: Array; - clearWorkspaceByokConfigs: Scalars['Boolean']['output']; completeBlobUpload: Scalars['String']['output']; createBlobUpload: BlobUploadInit; /** Create change password url */ @@ -1891,6 +1884,7 @@ export interface Mutation { /** Create a new workspace */ createWorkspace: WorkspaceType; createWorkspaceByokLocalLease: CreateWorkspaceByokLocalLeaseResultType; + createWorkspaceByokProfile: WorkspaceByokProfileType; deactivateLicense: Scalars['Boolean']['output']; deleteAccount: DeleteAccount; deleteAuthSigningKey: Array; @@ -1902,7 +1896,7 @@ export interface Mutation { /** Delete a user account */ deleteUser: DeleteAccount; deleteWorkspace: Scalars['Boolean']['output']; - deleteWorkspaceByokConfig: Scalars['Boolean']['output']; + deleteWorkspaceByokProfile: Scalars['Boolean']['output']; /** Reenable an banned user */ enableUser: UserType; /** Create a chat session */ @@ -1921,6 +1915,8 @@ export interface Mutation { /** mention user in a doc */ mentionUser: Scalars['ID']['output']; previewLicense: AdminLicensePreview; + probeWorkspaceByokDraft: WorkspaceByokProbeResultType; + probeWorkspaceByokProfile: WorkspaceByokProbeResultType; publishDoc: DocType; /** queue workspace doc embedding */ queueWorkspaceEmbedding: Scalars['Boolean']['output']; @@ -1944,7 +1940,8 @@ export interface Mutation { removeContextFile: Scalars['Boolean']['output']; /** Remove workspace embedding files */ removeWorkspaceEmbeddingFiles: Scalars['Boolean']['output']; - reorderWorkspaceByokConfigs: Array; + reorderWorkspaceByokProfiles: Array; + replaceWorkspaceByokProfile: WorkspaceByokProfileType; /** Request to apply the subscription in advance */ requestApplySubscription: Array; /** Resolve a comment or not */ @@ -1959,6 +1956,7 @@ export interface Mutation { revokePublicDoc: DocType; rotateAuthSigningKey: Array; rotateMcpCredential: RevealedMcpCredentialType; + rotateWorkspaceByokCredential: WorkspaceByokProfileType; sendChangeEmail: Scalars['Boolean']['output']; sendChangePasswordEmail: Scalars['Boolean']['output']; sendSetPasswordEmail: Scalars['Boolean']['output']; @@ -1968,7 +1966,6 @@ export interface Mutation { setBlob: Scalars['String']['output']; settleTranscriptTask: Maybe; submitTranscriptTask: Maybe; - testWorkspaceByokConfig: TestWorkspaceByokConfigResultType; unlinkCalendarAccount: Scalars['Boolean']['output']; /** update app configuration */ updateAppConfig: Scalars['JSONObject']['output']; @@ -1998,7 +1995,6 @@ export interface Mutation { uploadAvatar: UserType; /** Upload a comment attachment and return the access url */ uploadCommentAttachment: Scalars['String']['output']; - upsertWorkspaceByokConfig: WorkspaceByokKeyConfigType; verifyEmail: Scalars['Boolean']['output']; } @@ -2075,11 +2071,6 @@ export interface MutationCleanupCopilotSessionArgs { options: DeleteSessionInput; } -export interface MutationClearWorkspaceByokConfigsArgs { - provider?: InputMaybe; - workspaceId: Scalars['String']['input']; -} - export interface MutationCompleteBlobUploadArgs { key: Scalars['String']['input']; parts?: InputMaybe>; @@ -2153,6 +2144,10 @@ export interface MutationCreateWorkspaceByokLocalLeaseArgs { input: CreateWorkspaceByokLocalLeaseInput; } +export interface MutationCreateWorkspaceByokProfileArgs { + input: CreateWorkspaceByokProfileInput; +} + export interface MutationDeactivateLicenseArgs { workspaceId: Scalars['String']['input']; } @@ -2184,8 +2179,8 @@ export interface MutationDeleteWorkspaceArgs { id: Scalars['String']['input']; } -export interface MutationDeleteWorkspaceByokConfigArgs { - id: Scalars['ID']['input']; +export interface MutationDeleteWorkspaceByokProfileArgs { + profileId: Scalars['ID']['input']; workspaceId: Scalars['String']['input']; } @@ -2254,6 +2249,14 @@ export interface MutationPreviewLicenseArgs { license: Scalars['Upload']['input']; } +export interface MutationProbeWorkspaceByokDraftArgs { + input: ProbeWorkspaceByokDraftInput; +} + +export interface MutationProbeWorkspaceByokProfileArgs { + input: ProbeWorkspaceByokProfileInput; +} + export interface MutationPublishDocArgs { docId: Scalars['String']['input']; mode?: InputMaybe; @@ -2300,8 +2303,12 @@ export interface MutationRemoveWorkspaceEmbeddingFilesArgs { workspaceId: Scalars['String']['input']; } -export interface MutationReorderWorkspaceByokConfigsArgs { - input: ReorderWorkspaceByokConfigsInput; +export interface MutationReorderWorkspaceByokProfilesArgs { + input: ReorderWorkspaceByokProfilesInput; +} + +export interface MutationReplaceWorkspaceByokProfileArgs { + input: ReplaceWorkspaceByokProfileInput; } export interface MutationRequestApplySubscriptionArgs { @@ -2361,6 +2368,10 @@ export interface MutationRotateMcpCredentialArgs { workspaceId: Scalars['String']['input']; } +export interface MutationRotateWorkspaceByokCredentialArgs { + input: RotateWorkspaceByokCredentialInput; +} + export interface MutationSendChangeEmailArgs { callbackUrl: Scalars['String']['input']; } @@ -2407,10 +2418,6 @@ export interface MutationSubmitTranscriptTaskArgs { workspaceId: Scalars['String']['input']; } -export interface MutationTestWorkspaceByokConfigArgs { - input: TestWorkspaceByokConfigInput; -} - export interface MutationUnlinkCalendarAccountArgs { accountId: Scalars['String']['input']; } @@ -2493,10 +2500,6 @@ export interface MutationUploadCommentAttachmentArgs { workspaceId: Scalars['String']['input']; } -export interface MutationUpsertWorkspaceByokConfigArgs { - input: UpsertWorkspaceByokConfigInput; -} - export interface MutationVerifyEmailArgs { token: Scalars['String']['input']; } @@ -2692,6 +2695,22 @@ export enum Permission { Owner = 'Owner', } +export interface ProbeWorkspaceByokDraftInput { + checks: Array; + credential?: InputMaybe; + definition: WorkspaceByokProfileDefinitionInput; + expectedRevision?: InputMaybe; + profileId?: InputMaybe; + provider: ByokProvider; + workspaceId: Scalars['String']['input']; +} + +export interface ProbeWorkspaceByokProfileInput { + checks: Array; + profileId: Scalars['ID']['input']; + workspaceId: Scalars['String']['input']; +} + /** The mode which the public doc default in */ export enum PublicDocMode { Edgeless = 'Edgeless', @@ -2899,9 +2918,19 @@ export interface RemoveContextFileInput { fileId: Scalars['String']['input']; } -export interface ReorderWorkspaceByokConfigsInput { - ids: Array; - storage: ByokKeyStorage; +export interface ReorderWorkspaceByokProfilesInput { + profiles: Array; + workspaceId: Scalars['String']['input']; +} + +export interface ReplaceWorkspaceByokProfileInput { + credential?: InputMaybe; + definition: WorkspaceByokProfileDefinitionInput; + description?: InputMaybe; + enabled: Scalars['Boolean']['input']; + expectedRevision: Scalars['SafeInt']['input']; + name: Scalars['String']['input']; + profileId: Scalars['ID']['input']; workspaceId: Scalars['String']['input']; } @@ -2951,6 +2980,13 @@ export interface RevokeDocUserRoleInput { workspaceId: Scalars['String']['input']; } +export interface RotateWorkspaceByokCredentialInput { + credential: Scalars['String']['input']; + expectedRevision: Scalars['SafeInt']['input']; + profileId: Scalars['ID']['input']; + workspaceId: Scalars['String']['input']; +} + export interface RuntimeConfigNotFoundDataType { __typename?: 'RuntimeConfigNotFoundDataType'; key: Scalars['String']['output']; @@ -3136,7 +3172,6 @@ export interface SubmitAudioTranscriptionInput { quality?: InputMaybe; sliceManifest?: InputMaybe>; sourceAudio?: InputMaybe; - strategy?: InputMaybe; } export interface SubscriptionAlreadyExistsDataType { @@ -3222,22 +3257,6 @@ export enum SubscriptionVariant { Onetime = 'Onetime', } -export interface TestWorkspaceByokConfigInput { - apiKey?: InputMaybe; - configId?: InputMaybe; - endpoint?: InputMaybe; - provider: ByokProvider; - storage: ByokKeyStorage; - workspaceId: Scalars['String']['input']; -} - -export interface TestWorkspaceByokConfigResultType { - __typename?: 'TestWorkspaceByokConfigResultType'; - message: Maybe; - ok: Scalars['Boolean']['output']; - status: ByokKeyTestStatus; -} - export enum TimeBucket { Day = 'Day', Hour = 'Hour', @@ -3254,12 +3273,6 @@ export interface TimeWindow { to: Scalars['DateTime']['output']; } -export interface TranscriptProviderMetaType { - __typename?: 'TranscriptProviderMetaType'; - model: Maybe; - provider: Maybe; -} - export interface TranscriptionItemType { __typename?: 'TranscriptionItemType'; end: Scalars['String']['output']; @@ -3285,12 +3298,10 @@ export interface TranscriptionResultType { id: Scalars['ID']['output']; normalizedSegments: Maybe>; normalizedTranscript: Maybe; - providerMeta: Maybe; quality: Maybe; sliceManifest: Maybe>; sourceAudio: Maybe; status: AiJobStatus; - strategy: Maybe; summary: Maybe; summaryJson: Maybe; title: Maybe; @@ -3406,19 +3417,6 @@ export interface UpdateWorkspaceInput { public?: InputMaybe; } -export interface UpsertWorkspaceByokConfigInput { - apiKey?: InputMaybe; - description?: InputMaybe; - enabled?: InputMaybe; - endpoint?: InputMaybe; - id?: InputMaybe; - name: Scalars['String']['input']; - provider: ByokProvider; - sortOrder?: InputMaybe; - storage: ByokKeyStorage; - workspaceId: Scalars['String']['input']; -} - export interface UserImportFailedType { __typename?: 'UserImportFailedType'; email: Scalars['String']['output']; @@ -3530,48 +3528,141 @@ export interface VersionRejectedDataType { version: Scalars['String']['output']; } -export interface WorkspaceByokCapabilityWarningType { - __typename?: 'WorkspaceByokCapabilityWarningType'; - featureKind: Scalars['String']['output']; - reason: Scalars['String']['output']; - requiredProviders: Array; +export interface WorkspaceByokCapabilityInput { + attachmentKinds: Array; + attachmentSources: Array; + features: Array; + input: Array; + output: Array; } -export interface WorkspaceByokKeyConfigType { - __typename?: 'WorkspaceByokKeyConfigType'; - capabilities: Array; - configured: Scalars['Boolean']['output']; - description: Maybe; - disabledReason: Maybe; - enabled: Scalars['Boolean']['output']; - endpoint: Maybe; - endpointEditable: Scalars['Boolean']['output']; - id: Scalars['ID']['output']; - lastError: Maybe; - lastErrorAt: Maybe; - lastTestError: Maybe; - lastTestedAt: Maybe; - lastUsedAt: Maybe; - name: Scalars['String']['output']; +export interface WorkspaceByokCapabilityType { + __typename?: 'WorkspaceByokCapabilityType'; + attachmentKinds: Array; + attachmentSources: Array; + features: Array; + input: Array; + output: Array; +} + +export interface WorkspaceByokCatalogModelType { + __typename?: 'WorkspaceByokCatalogModelType'; + capabilities: Array; + displayName: Scalars['String']['output']; + modelId: Scalars['String']['output']; + recommended: Scalars['Boolean']['output']; +} + +export interface WorkspaceByokCatalogProviderType { + __typename?: 'WorkspaceByokCatalogProviderType'; + models: Array; provider: ByokProvider; +} + +export interface WorkspaceByokCatalogType { + __typename?: 'WorkspaceByokCatalogType'; + providers: Array; + version: Scalars['String']['output']; +} + +export interface WorkspaceByokEndpointInput { + kind: Scalars['String']['input']; + url?: InputMaybe; +} + +export interface WorkspaceByokEndpointType { + __typename?: 'WorkspaceByokEndpointType'; + kind: Scalars['String']['output']; + url: Maybe; +} + +export interface WorkspaceByokModelDeclarationInput { + capabilities: Array; + enabled: Scalars['Boolean']['input']; + modelId: Scalars['String']['input']; +} + +export interface WorkspaceByokModelDeclarationType { + __typename?: 'WorkspaceByokModelDeclarationType'; + capabilities: Array; + enabled: Scalars['Boolean']['output']; + modelId: Scalars['String']['output']; +} + +export interface WorkspaceByokModelProbeCheckType { + __typename?: 'WorkspaceByokModelProbeCheckType'; + operation: Scalars['String']['output']; + status: WorkspaceByokProbeStatusType; +} + +export interface WorkspaceByokModelProbeType { + __typename?: 'WorkspaceByokModelProbeType'; + checks: Array; + modelId: Scalars['String']['output']; +} + +export interface WorkspaceByokProbeCheckInput { + modelId: Scalars['String']['input']; + operation: Scalars['String']['input']; +} + +export interface WorkspaceByokProbeResultType { + __typename?: 'WorkspaceByokProbeResultType'; + connection: WorkspaceByokProbeStatusType; + definitionFingerprint: Scalars['String']['output']; + models: Array; + stale: Scalars['Boolean']['output']; +} + +export interface WorkspaceByokProbeStatusType { + __typename?: 'WorkspaceByokProbeStatusType'; + errorKind: Maybe; + kind: Scalars['String']['output']; + testedAt: Maybe; +} + +export interface WorkspaceByokProfileDefinitionInput { + endpoint: WorkspaceByokEndpointInput; + models: Array; + version: Scalars['SafeInt']['input']; +} + +export interface WorkspaceByokProfileDefinitionType { + __typename?: 'WorkspaceByokProfileDefinitionType'; + endpoint: WorkspaceByokEndpointType; + models: Array; + version: Scalars['SafeInt']['output']; +} + +export interface WorkspaceByokProfileOrderInput { + expectedRevision: Scalars['SafeInt']['input']; + profileId: Scalars['ID']['input']; +} + +export interface WorkspaceByokProfileType { + __typename?: 'WorkspaceByokProfileType'; + definition: WorkspaceByokProfileDefinitionType; + description: Maybe; + enabled: Scalars['Boolean']['output']; + name: Scalars['String']['output']; + profileId: Scalars['ID']['output']; + provider: ByokProvider; + revision: Scalars['SafeInt']['output']; sortOrder: Scalars['SafeInt']['output']; - storage: ByokKeyStorage; - testStatus: ByokKeyTestStatus; + validation: Maybe; + workspaceId: Scalars['String']['output']; } export interface WorkspaceByokSettingsType { __typename?: 'WorkspaceByokSettingsType'; allowedProviders: Array; + catalog: WorkspaceByokCatalogType; customEndpointSupported: Scalars['Boolean']['output']; entitled: Scalars['Boolean']['output']; - entitlementRequired: Array; - hasAiPlan: Scalars['Boolean']['output']; - keys: Array; localEntitled: Scalars['Boolean']['output']; - localStorageSupported: Scalars['Boolean']['output']; privateEndpointSupported: Scalars['Boolean']['output']; + profiles: Array; serverEntitled: Scalars['Boolean']['output']; - warnings: Array; workspaceId: Scalars['String']['output']; } @@ -3582,6 +3673,14 @@ export interface WorkspaceByokUsagePointType { totalTokens: Scalars['SafeInt']['output']; } +export interface WorkspaceByokValidationType { + __typename?: 'WorkspaceByokValidationType'; + connection: WorkspaceByokProbeStatusType; + credentialGeneration: Scalars['SafeInt']['output']; + definitionFingerprint: Scalars['String']['output']; + models: Array; +} + export interface WorkspaceCalendarItemInput { colorOverride?: InputMaybe; sortOrder?: InputMaybe; @@ -5367,12 +5466,9 @@ export type GetCopilotDocSessionsQuery = { docId: string | null; parentSessionId: string | null; promptName: string; - model: string; - optionalModels: Array; action: string | null; pinned: boolean; title: string | null; - tokens: number; createdAt: string; updatedAt: string; messages: Array<{ @@ -5431,12 +5527,9 @@ export type GetCopilotPinnedSessionsQuery = { docId: string | null; parentSessionId: string | null; promptName: string; - model: string; - optionalModels: Array; action: string | null; pinned: boolean; title: string | null; - tokens: number; createdAt: string; updatedAt: string; messages: Array<{ @@ -5494,12 +5587,9 @@ export type GetCopilotWorkspaceSessionsQuery = { docId: string | null; parentSessionId: string | null; promptName: string; - model: string; - optionalModels: Array; action: string | null; pinned: boolean; title: string | null; - tokens: number; createdAt: string; updatedAt: string; messages: Array<{ @@ -5558,12 +5648,9 @@ export type GetCopilotHistoriesQuery = { docId: string | null; parentSessionId: string | null; promptName: string; - model: string; - optionalModels: Array; action: string | null; pinned: boolean; title: string | null; - tokens: number; createdAt: string; updatedAt: string; messages: Array<{ @@ -5599,34 +5686,6 @@ export type CreateCopilotMessageMutation = { createCopilotMessage: string; }; -export type GetPromptModelsQueryVariables = Exact<{ - promptName: Scalars['String']['input']; -}>; - -export type GetPromptModelsQuery = { - __typename?: 'Query'; - currentUser: { - __typename?: 'UserType'; - copilot: { - __typename?: 'Copilot'; - models: { - __typename?: 'CopilotModelsType'; - defaultModel: string; - optionalModels: Array<{ - __typename?: 'CopilotModelType'; - id: string; - name: string; - }>; - proModels: Array<{ - __typename?: 'CopilotModelType'; - id: string; - name: string; - }>; - }; - }; - } | null; -}; - export type CopilotQuotaQueryVariables = Exact<{ [key: string]: never }>; export type CopilotQuotaQuery = { @@ -5644,6 +5703,32 @@ export type CopilotQuotaQuery = { } | null; }; +export type GetCopilotRouteOptionsQueryVariables = Exact<{ + promptName: Scalars['String']['input']; +}>; + +export type GetCopilotRouteOptionsQuery = { + __typename?: 'Query'; + currentUser: { + __typename?: 'UserType'; + copilot: { + __typename?: 'Copilot'; + routeOptions: { + __typename?: 'CopilotRouteOptions'; + routeId: string; + defaultTargetId: string | null; + choices: Array<{ + __typename?: 'CopilotRouteTarget'; + id: string; + displayName: string; + minimumTier: string; + available: boolean; + }>; + } | null; + }; + } | null; +}; + export type CleanupCopilotSessionMutationVariables = Exact<{ input: DeleteSessionInput; }>; @@ -5666,12 +5751,9 @@ export type CreateCopilotSessionWithHistoryMutation = { docId: string | null; parentSessionId: string | null; promptName: string; - model: string; - optionalModels: Array; action: string | null; pinned: boolean; title: string | null; - tokens: number; createdAt: string; updatedAt: string; messages: Array<{ @@ -5742,12 +5824,9 @@ export type GetCopilotLatestDocSessionQuery = { docId: string | null; parentSessionId: string | null; promptName: string; - model: string; - optionalModels: Array; action: string | null; pinned: boolean; title: string | null; - tokens: number; createdAt: string; updatedAt: string; messages: Array<{ @@ -5804,12 +5883,9 @@ export type GetCopilotSessionQuery = { docId: string | null; parentSessionId: string | null; promptName: string; - model: string; - optionalModels: Array; action: string | null; pinned: boolean; title: string | null; - tokens: number; createdAt: string; updatedAt: string; messages: Array<{ @@ -5867,12 +5943,9 @@ export type GetCopilotRecentSessionsQuery = { docId: string | null; parentSessionId: string | null; promptName: string; - model: string; - optionalModels: Array; action: string | null; pinned: boolean; title: string | null; - tokens: number; createdAt: string; updatedAt: string; messages: Array<{ @@ -5940,12 +6013,9 @@ export type GetCopilotSessionsQuery = { docId: string | null; parentSessionId: string | null; promptName: string; - model: string; - optionalModels: Array; action: string | null; pinned: boolean; title: string | null; - tokens: number; createdAt: string; updatedAt: string; messages: Array<{ @@ -6393,12 +6463,9 @@ export type CopilotChatHistoryFragment = { docId: string | null; parentSessionId: string | null; promptName: string; - model: string; - optionalModels: Array; action: string | null; pinned: boolean; title: string | null; - tokens: number; createdAt: string; updatedAt: string; messages: Array<{ @@ -6448,12 +6515,9 @@ export type PaginatedCopilotChatsFragment = { docId: string | null; parentSessionId: string | null; promptName: string; - model: string; - optionalModels: Array; action: string | null; pinned: boolean; title: string | null; - tokens: number; createdAt: string; updatedAt: string; messages: Array<{ @@ -7693,61 +7757,115 @@ export type WorkspaceBlobQuotaQuery = { }; }; -export type ClearWorkspaceByokConfigsMutationVariables = Exact<{ +export type DeleteWorkspaceByokProfileMutationVariables = Exact<{ workspaceId: Scalars['String']['input']; + profileId: Scalars['ID']['input']; }>; -export type ClearWorkspaceByokConfigsMutation = { +export type DeleteWorkspaceByokProfileMutation = { __typename?: 'Mutation'; - clearWorkspaceByokConfigs: boolean; + deleteWorkspaceByokProfile: boolean; }; -export type DeleteWorkspaceByokConfigMutationVariables = Exact<{ - workspaceId: Scalars['String']['input']; - id: Scalars['ID']['input']; +export type ProbeWorkspaceByokProfileMutationVariables = Exact<{ + input: ProbeWorkspaceByokProfileInput; }>; -export type DeleteWorkspaceByokConfigMutation = { +export type ProbeWorkspaceByokProfileMutation = { __typename?: 'Mutation'; - deleteWorkspaceByokConfig: boolean; -}; - -export type ReorderWorkspaceByokConfigsMutationVariables = Exact<{ - input: ReorderWorkspaceByokConfigsInput; -}>; - -export type ReorderWorkspaceByokConfigsMutation = { - __typename?: 'Mutation'; - reorderWorkspaceByokConfigs: Array<{ - __typename?: 'WorkspaceByokKeyConfigType'; - id: string; - sortOrder: number; - }>; -}; - -export type TestWorkspaceByokConfigMutationVariables = Exact<{ - input: TestWorkspaceByokConfigInput; -}>; - -export type TestWorkspaceByokConfigMutation = { - __typename?: 'Mutation'; - testWorkspaceByokConfig: { - __typename?: 'TestWorkspaceByokConfigResultType'; - ok: boolean; - status: ByokKeyTestStatus; - message: string | null; + probeWorkspaceByokProfile: { + __typename?: 'WorkspaceByokProbeResultType'; + definitionFingerprint: string; + stale: boolean; + connection: { + __typename?: 'WorkspaceByokProbeStatusType'; + kind: string; + testedAt: string | null; + errorKind: string | null; + }; + models: Array<{ + __typename?: 'WorkspaceByokModelProbeType'; + modelId: string; + checks: Array<{ + __typename?: 'WorkspaceByokModelProbeCheckType'; + operation: string; + status: { + __typename?: 'WorkspaceByokProbeStatusType'; + kind: string; + testedAt: string | null; + errorKind: string | null; + }; + }>; + }>; }; }; -export type UpsertWorkspaceByokConfigMutationVariables = Exact<{ - input: UpsertWorkspaceByokConfigInput; +export type ProbeWorkspaceByokDraftMutationVariables = Exact<{ + input: ProbeWorkspaceByokDraftInput; }>; -export type UpsertWorkspaceByokConfigMutation = { +export type ProbeWorkspaceByokDraftMutation = { __typename?: 'Mutation'; - upsertWorkspaceByokConfig: { - __typename?: 'WorkspaceByokKeyConfigType'; - id: string; + probeWorkspaceByokDraft: { + __typename?: 'WorkspaceByokProbeResultType'; + definitionFingerprint: string; + stale: boolean; + connection: { + __typename?: 'WorkspaceByokProbeStatusType'; + kind: string; + testedAt: string | null; + errorKind: string | null; + }; + models: Array<{ + __typename?: 'WorkspaceByokModelProbeType'; + modelId: string; + checks: Array<{ + __typename?: 'WorkspaceByokModelProbeCheckType'; + operation: string; + status: { + __typename?: 'WorkspaceByokProbeStatusType'; + kind: string; + testedAt: string | null; + errorKind: string | null; + }; + }>; + }>; + }; +}; + +export type CreateWorkspaceByokProfileMutationVariables = Exact<{ + input: CreateWorkspaceByokProfileInput; +}>; + +export type CreateWorkspaceByokProfileMutation = { + __typename?: 'Mutation'; + createWorkspaceByokProfile: { + __typename?: 'WorkspaceByokProfileType'; + profileId: string; + }; +}; + +export type ReplaceWorkspaceByokProfileMutationVariables = Exact<{ + input: ReplaceWorkspaceByokProfileInput; +}>; + +export type ReplaceWorkspaceByokProfileMutation = { + __typename?: 'Mutation'; + replaceWorkspaceByokProfile: { + __typename?: 'WorkspaceByokProfileType'; + profileId: string; + }; +}; + +export type RotateWorkspaceByokCredentialMutationVariables = Exact<{ + input: RotateWorkspaceByokCredentialInput; +}>; + +export type RotateWorkspaceByokCredentialMutation = { + __typename?: 'Mutation'; + rotateWorkspaceByokCredential: { + __typename?: 'WorkspaceByokProfileType'; + profileId: string; }; }; @@ -7764,6 +7882,20 @@ export type CreateWorkspaceByokLocalLeaseMutation = { }; }; +export type ReorderWorkspaceByokProfilesMutationVariables = Exact<{ + input: ReorderWorkspaceByokProfilesInput; +}>; + +export type ReorderWorkspaceByokProfilesMutation = { + __typename?: 'Mutation'; + reorderWorkspaceByokProfiles: Array<{ + __typename?: 'WorkspaceByokProfileType'; + profileId: string; + sortOrder: number; + revision: number; + }>; +}; + export type WorkspaceByokSettingsQueryVariables = Exact<{ id: Scalars['String']['input']; from: Scalars['DateTime']['input']; @@ -7781,38 +7913,87 @@ export type WorkspaceByokSettingsQuery = { entitled: boolean; serverEntitled: boolean; localEntitled: boolean; - entitlementRequired: Array; allowedProviders: Array; - localStorageSupported: boolean; customEndpointSupported: boolean; privateEndpointSupported: boolean; - hasAiPlan: boolean; - keys: Array<{ - __typename?: 'WorkspaceByokKeyConfigType'; - id: string; + catalog: { + __typename?: 'WorkspaceByokCatalogType'; + version: string; + providers: Array<{ + __typename?: 'WorkspaceByokCatalogProviderType'; + provider: ByokProvider; + models: Array<{ + __typename?: 'WorkspaceByokCatalogModelType'; + modelId: string; + displayName: string; + recommended: boolean; + capabilities: Array<{ + __typename?: 'WorkspaceByokCapabilityType'; + input: Array; + output: Array; + features: Array; + attachmentKinds: Array; + attachmentSources: Array; + }>; + }>; + }>; + }; + profiles: Array<{ + __typename?: 'WorkspaceByokProfileType'; + profileId: string; provider: ByokProvider; name: string; description: string | null; - storage: ByokKeyStorage; - configured: boolean; enabled: boolean; - endpoint: string | null; - endpointEditable: boolean; sortOrder: number; - capabilities: Array; - testStatus: ByokKeyTestStatus; - disabledReason: string | null; - lastTestedAt: string | null; - lastTestError: string | null; - lastUsedAt: string | null; - lastErrorAt: string | null; - lastError: string | null; - }>; - warnings: Array<{ - __typename?: 'WorkspaceByokCapabilityWarningType'; - featureKind: string; - reason: string; - requiredProviders: Array; + revision: number; + definition: { + __typename?: 'WorkspaceByokProfileDefinitionType'; + version: number; + endpoint: { + __typename?: 'WorkspaceByokEndpointType'; + kind: string; + url: string | null; + }; + models: Array<{ + __typename?: 'WorkspaceByokModelDeclarationType'; + modelId: string; + enabled: boolean; + capabilities: Array<{ + __typename?: 'WorkspaceByokCapabilityType'; + input: Array; + output: Array; + features: Array; + attachmentKinds: Array; + attachmentSources: Array; + }>; + }>; + }; + validation: { + __typename?: 'WorkspaceByokValidationType'; + definitionFingerprint: string; + credentialGeneration: number; + connection: { + __typename?: 'WorkspaceByokProbeStatusType'; + kind: string; + testedAt: string | null; + errorKind: string | null; + }; + models: Array<{ + __typename?: 'WorkspaceByokModelProbeType'; + modelId: string; + checks: Array<{ + __typename?: 'WorkspaceByokModelProbeCheckType'; + operation: string; + status: { + __typename?: 'WorkspaceByokProbeStatusType'; + kind: string; + testedAt: string | null; + errorKind: string | null; + }; + }>; + }>; + } | null; }>; }; byokUsage: Array<{ @@ -8141,16 +8322,16 @@ export type Queries = variables: GetCopilotHistoriesQueryVariables; response: GetCopilotHistoriesQuery; } - | { - name: 'getPromptModelsQuery'; - variables: GetPromptModelsQueryVariables; - response: GetPromptModelsQuery; - } | { name: 'copilotQuotaQuery'; variables: CopilotQuotaQueryVariables; response: CopilotQuotaQuery; } + | { + name: 'getCopilotRouteOptionsQuery'; + variables: GetCopilotRouteOptionsQueryVariables; + response: GetCopilotRouteOptionsQuery; + } | { name: 'getCopilotLatestDocSessionQuery'; variables: GetCopilotLatestDocSessionQueryVariables; @@ -8884,35 +9065,45 @@ export type Mutations = response: VerifyEmailMutation; } | { - name: 'clearWorkspaceByokConfigsMutation'; - variables: ClearWorkspaceByokConfigsMutationVariables; - response: ClearWorkspaceByokConfigsMutation; + name: 'deleteWorkspaceByokProfileMutation'; + variables: DeleteWorkspaceByokProfileMutationVariables; + response: DeleteWorkspaceByokProfileMutation; } | { - name: 'deleteWorkspaceByokConfigMutation'; - variables: DeleteWorkspaceByokConfigMutationVariables; - response: DeleteWorkspaceByokConfigMutation; + name: 'probeWorkspaceByokProfileMutation'; + variables: ProbeWorkspaceByokProfileMutationVariables; + response: ProbeWorkspaceByokProfileMutation; } | { - name: 'reorderWorkspaceByokConfigsMutation'; - variables: ReorderWorkspaceByokConfigsMutationVariables; - response: ReorderWorkspaceByokConfigsMutation; + name: 'probeWorkspaceByokDraftMutation'; + variables: ProbeWorkspaceByokDraftMutationVariables; + response: ProbeWorkspaceByokDraftMutation; } | { - name: 'testWorkspaceByokConfigMutation'; - variables: TestWorkspaceByokConfigMutationVariables; - response: TestWorkspaceByokConfigMutation; + name: 'createWorkspaceByokProfileMutation'; + variables: CreateWorkspaceByokProfileMutationVariables; + response: CreateWorkspaceByokProfileMutation; } | { - name: 'upsertWorkspaceByokConfigMutation'; - variables: UpsertWorkspaceByokConfigMutationVariables; - response: UpsertWorkspaceByokConfigMutation; + name: 'replaceWorkspaceByokProfileMutation'; + variables: ReplaceWorkspaceByokProfileMutationVariables; + response: ReplaceWorkspaceByokProfileMutation; + } + | { + name: 'rotateWorkspaceByokCredentialMutation'; + variables: RotateWorkspaceByokCredentialMutationVariables; + response: RotateWorkspaceByokCredentialMutation; } | { name: 'createWorkspaceByokLocalLeaseMutation'; variables: CreateWorkspaceByokLocalLeaseMutationVariables; response: CreateWorkspaceByokLocalLeaseMutation; } + | { + name: 'reorderWorkspaceByokProfilesMutation'; + variables: ReorderWorkspaceByokProfilesMutationVariables; + response: ReorderWorkspaceByokProfilesMutation; + } | { name: 'setEnableAiMutation'; variables: SetEnableAiMutationVariables; diff --git a/packages/frontend/admin/src/config.json b/packages/frontend/admin/src/config.json index 0f37779635..f8ba0a7061 100644 --- a/packages/frontend/admin/src/config.json +++ b/packages/frontend/admin/src/config.json @@ -389,39 +389,6 @@ "type": "Array", "desc": "The profile list for copilot providers." }, - "providers.defaults": { - "type": "Object", - "desc": "The default provider ids for model output types and global fallback." - }, - "providers.openai": { - "type": "Object", - "desc": "The config for the openai provider.", - "link": "https://github.com/openai/openai-node" - }, - "providers.cloudflareWorkersAi": { - "type": "Object", - "desc": "The config for the Cloudflare Workers AI provider." - }, - "providers.fal": { - "type": "Object", - "desc": "The config for the fal provider." - }, - "providers.gemini": { - "type": "Object", - "desc": "The config for the gemini provider." - }, - "providers.geminiVertex": { - "type": "Object", - "desc": "The config for the gemini provider in Google Vertex AI." - }, - "providers.anthropic": { - "type": "Object", - "desc": "The config for the anthropic provider." - }, - "providers.anthropicVertex": { - "type": "Object", - "desc": "The config for the anthropic provider in Google Vertex AI." - }, "unsplash": { "type": "Object", "desc": "The config for the unsplash key." diff --git a/packages/frontend/apps/electron/src/main/byok-storage/handlers.ts b/packages/frontend/apps/electron/src/main/byok-storage/handlers.ts index be97a7cad3..f6ea3e12af 100644 --- a/packages/frontend/apps/electron/src/main/byok-storage/handlers.ts +++ b/packages/frontend/apps/electron/src/main/byok-storage/handlers.ts @@ -14,20 +14,55 @@ export function disposeWorkspaceByokStorage() { } const allowedProviders = new Set(['openai', 'anthropic', 'gemini', 'fal']); +const allowedInputs = new Set(['text', 'image', 'audio', 'file']); +const allowedOutputs = new Set([ + 'text', + 'object', + 'structured', + 'embedding', + 'rerank', + 'image', +]); +const allowedFeatures = new Set(['tool_calling', 'reasoning', 'web_search']); +const allowedAttachmentKinds = new Set(['image', 'audio', 'file']); +const allowedAttachmentSources = new Set([ + 'url', + 'data', + 'bytes', + 'file_handle', +]); type WorkspaceByokKey = { id: string; provider: 'openai' | 'anthropic' | 'gemini' | 'fal'; name: string; description?: string | null; - apiKey: string; - endpoint?: string | null; + credential: string; + definition: { + version: number; + endpoint: { kind: string; url?: string | null }; + models: Array<{ + modelId: string; + enabled: boolean; + capabilities: Array<{ + input: string[]; + output: string[]; + features: string[]; + attachmentKinds: string[]; + attachmentSources: string[]; + }>; + }>; + }; sortOrder?: number | null; enabled?: boolean | null; }; -type WorkspaceByokKeyInput = Omit & { - apiKey?: string | null; +type WorkspaceByokKeyInput = Omit< + WorkspaceByokKey, + 'credential' | 'definition' +> & { + credential?: string | null; + definition?: WorkspaceByokKey['definition']; }; function assertSupported() { @@ -43,6 +78,73 @@ function hasOwnField( return Object.prototype.hasOwnProperty.call(key, field); } +function isRecord(value: unknown): value is Record { + return typeof value === 'object' && value !== null && !Array.isArray(value); +} + +function isAllowedStringArray( + value: unknown, + allowed: Set +): value is string[] { + return ( + Array.isArray(value) && + value.every(item => typeof item === 'string' && allowed.has(item)) + ); +} + +function isValidEndpoint(value: unknown) { + if (!isRecord(value) || typeof value.kind !== 'string') return false; + if (value.kind === 'provider_default') return value.url == null; + if (value.kind !== 'custom' || typeof value.url !== 'string') return false; + try { + const endpoint = new URL(value.url); + return ( + (endpoint.protocol === 'http:' || endpoint.protocol === 'https:') && + !!endpoint.hostname && + !endpoint.username && + !endpoint.password + ); + } catch { + return false; + } +} + +function isValidCapability(value: unknown) { + return ( + isRecord(value) && + isAllowedStringArray(value.input, allowedInputs) && + value.input.length > 0 && + isAllowedStringArray(value.output, allowedOutputs) && + value.output.length > 0 && + isAllowedStringArray(value.features, allowedFeatures) && + isAllowedStringArray(value.attachmentKinds, allowedAttachmentKinds) && + isAllowedStringArray(value.attachmentSources, allowedAttachmentSources) + ); +} + +function isValidDefinition( + value: unknown +): value is WorkspaceByokKey['definition'] { + return ( + isRecord(value) && + value.version === 1 && + isValidEndpoint(value.endpoint) && + Array.isArray(value.models) && + value.models.length > 0 && + value.models.every( + model => + isRecord(model) && + typeof model.modelId === 'string' && + model.modelId.trim().length > 0 && + model.modelId.length <= 512 && + typeof model.enabled === 'boolean' && + Array.isArray(model.capabilities) && + model.capabilities.length > 0 && + model.capabilities.every(isValidCapability) + ) + ); +} + function normalizeKey( key: WorkspaceByokKeyInput, existing?: WorkspaceByokKey, @@ -51,8 +153,9 @@ function normalizeKey( if (!allowedProviders.has(key.provider)) { throw new Error('Unsupported BYOK provider.'); } - const apiKey = key.apiKey ?? existing?.apiKey; - if (!key.id || !key.name || !apiKey) { + const credential = key.credential ?? existing?.credential; + const definition = key.definition ?? existing?.definition; + if (!key.id || !key.name || !credential || !isValidDefinition(definition)) { throw new Error('Invalid BYOK key.'); } return { @@ -62,10 +165,8 @@ function normalizeKey( description: hasOwnField(key, 'description') ? (key.description ?? null) : (existing?.description ?? null), - apiKey, - endpoint: hasOwnField(key, 'endpoint') - ? (key.endpoint ?? null) - : (existing?.endpoint ?? null), + credential, + definition, sortOrder: hasOwnField(key, 'sortOrder') ? (key.sortOrder ?? defaultSortOrder) : (existing?.sortOrder ?? defaultSortOrder), @@ -111,13 +212,11 @@ function writeWorkspaceKeys(workspaceId: string, keys: WorkspaceByokKey[]) { byokStorage.set(workspaceId, keys.map(encryptKey)); } -function toPublicKey({ apiKey: _, ...key }: WorkspaceByokKey) { +function toPublicKey({ credential: _, ...key }: WorkspaceByokKey) { return { ...key, storage: 'local', configured: true, - endpointEditable: false, - testStatus: 'passed', }; } diff --git a/packages/frontend/apps/electron/test/main/byok-storage.spec.ts b/packages/frontend/apps/electron/test/main/byok-storage.spec.ts index f5b3ac4739..b4275e4b2e 100644 --- a/packages/frontend/apps/electron/test/main/byok-storage.spec.ts +++ b/packages/frontend/apps/electron/test/main/byok-storage.spec.ts @@ -75,6 +75,25 @@ afterEach(async () => { }); describe('byok storage handlers', () => { + const definition = { + version: 1, + endpoint: { kind: 'provider_default' }, + models: [ + { + modelId: 'model-1', + enabled: true, + capabilities: [ + { + input: ['text'], + output: ['text'], + features: [], + attachmentKinds: [], + attachmentSources: [], + }, + ], + }, + ], + }; test('stores encrypted local keys and keeps lease providers sorted', async () => { const { byokStorageHandlers, disposeWorkspaceByokStorage: dispose } = await import('@affine/electron/main/byok-storage/handlers'); @@ -85,14 +104,16 @@ describe('byok storage handlers', () => { id: 'local-openai', provider: 'openai', name: 'OpenAI', - apiKey: 'sk-openai', + credential: 'sk-openai', + definition, sortOrder: 1, }); await byokStorageHandlers.upsertWorkspaceKey(ipcEvent, 'workspace-1', { id: 'local-gemini', provider: 'gemini', name: 'Gemini', - apiKey: 'sk-gemini', + credential: 'sk-gemini', + definition, sortOrder: 0, }); @@ -117,7 +138,7 @@ describe('byok storage handlers', () => { ipcEvent, 'workspace-1' ); - expect(leaseProviders.map(key => key.apiKey)).toEqual([ + expect(leaseProviders.map(key => key.credential)).toEqual([ 'sk-openai', 'sk-gemini', ]); @@ -142,12 +163,60 @@ describe('byok storage handlers', () => { id: 'local-openai', provider: 'openai', name: 'OpenAI', - apiKey: 'sk-openai', + credential: 'sk-openai', + definition, }) ).rejects.toThrow('Secure BYOK key storage is not available.'); expect(electronMock.encryptString).not.toHaveBeenCalled(); }); + test.each([ + [ + 'custom endpoint without URL', + { ...definition, endpoint: { kind: 'custom' } }, + ], + [ + 'unsupported endpoint protocol', + { ...definition, endpoint: { kind: 'custom', url: 'file:///tmp/api' } }, + ], + [ + 'malformed capability object', + { + ...definition, + models: [{ ...definition.models[0], capabilities: [{}] }], + }, + ], + [ + 'unknown capability value', + { + ...definition, + models: [ + { + ...definition.models[0], + capabilities: [ + { ...definition.models[0].capabilities[0], input: ['video'] }, + ], + }, + ], + }, + ], + ])('rejects %s from IPC input', async (_name, malformedDefinition) => { + const { byokStorageHandlers, disposeWorkspaceByokStorage: dispose } = + await import('@affine/electron/main/byok-storage/handlers'); + disposeWorkspaceByokStorage = dispose; + + await expect( + byokStorageHandlers.upsertWorkspaceKey(undefined, 'workspace-1', { + id: 'local-openai', + provider: 'openai', + name: 'OpenAI', + credential: 'sk-openai', + definition: malformedDefinition as typeof definition, + }) + ).rejects.toThrow('Invalid BYOK key.'); + expect(electronMock.encryptString).not.toHaveBeenCalled(); + }); + test('preserves existing local key fields during partial updates', async () => { const { byokStorageHandlers, disposeWorkspaceByokStorage: dispose } = await import('@affine/electron/main/byok-storage/handlers'); @@ -159,8 +228,11 @@ describe('byok storage handlers', () => { provider: 'openai', name: 'OpenAI', description: 'Primary key', - apiKey: 'sk-openai', - endpoint: 'https://api.openai.example/v1', + credential: 'sk-openai', + definition: { + ...definition, + endpoint: { kind: 'custom', url: 'https://api.openai.example/v1' }, + }, sortOrder: 4, enabled: false, }); @@ -169,7 +241,7 @@ describe('byok storage handlers', () => { id: 'local-openai', provider: 'openai', name: 'OpenAI renamed', - apiKey: 'sk-openai-next', + credential: 'sk-openai-next', }); const [publicKey] = await byokStorageHandlers.listWorkspaceKeys( @@ -180,7 +252,10 @@ describe('byok storage handlers', () => { id: 'local-openai', name: 'OpenAI renamed', description: 'Primary key', - endpoint: 'https://api.openai.example/v1', + definition: { + ...definition, + endpoint: { kind: 'custom', url: 'https://api.openai.example/v1' }, + }, sortOrder: 4, enabled: false, }); @@ -206,8 +281,11 @@ describe('byok storage handlers', () => { ); expect(enabledLeaseProvider).toMatchObject({ name: 'OpenAI renamed again', - apiKey: 'sk-openai-next', - endpoint: 'https://api.openai.example/v1', + credential: 'sk-openai-next', + definition: { + ...definition, + endpoint: { kind: 'custom', url: 'https://api.openai.example/v1' }, + }, sortOrder: 4, enabled: true, }); diff --git a/packages/frontend/apps/ios/App/Packages/AffineGraphQL/Package.swift b/packages/frontend/apps/ios/App/Packages/AffineGraphQL/Package.swift index 10e90b965e..119fff171d 100644 --- a/packages/frontend/apps/ios/App/Packages/AffineGraphQL/Package.swift +++ b/packages/frontend/apps/ios/App/Packages/AffineGraphQL/Package.swift @@ -14,7 +14,7 @@ let package = Package( .library(name: "AffineGraphQL", targets: ["AffineGraphQL"]), ], dependencies: [ - .package(url: "https://github.com/apollographql/apollo-ios", exact: "1.25.7"), + .package(url: "https://github.com/apollographql/apollo-ios", exact: "1.25.4"), ], targets: [ .target( diff --git a/packages/frontend/apps/ios/App/Packages/AffineGraphQL/Sources/Fragments/CopilotChatHistory.graphql.swift b/packages/frontend/apps/ios/App/Packages/AffineGraphQL/Sources/Fragments/CopilotChatHistory.graphql.swift index e30d7d1040..0af6f123b1 100644 --- a/packages/frontend/apps/ios/App/Packages/AffineGraphQL/Sources/Fragments/CopilotChatHistory.graphql.swift +++ b/packages/frontend/apps/ios/App/Packages/AffineGraphQL/Sources/Fragments/CopilotChatHistory.graphql.swift @@ -5,7 +5,7 @@ public struct CopilotChatHistory: AffineGraphQL.SelectionSet, Fragment { public static var fragmentDefinition: StaticString { - #"fragment CopilotChatHistory on CopilotHistories { __typename sessionId workspaceId docId parentSessionId promptName model optionalModels action pinned title tokens messages { __typename id role content attachments streamObjects { __typename type textDelta toolCallId toolName args result } createdAt } createdAt updatedAt }"# + #"fragment CopilotChatHistory on CopilotHistories { __typename sessionId workspaceId docId parentSessionId promptName action pinned title messages { __typename id role content attachments streamObjects { __typename type textDelta toolCallId toolName args result } createdAt } createdAt updatedAt }"# } public let __data: DataDict @@ -19,12 +19,9 @@ public struct CopilotChatHistory: AffineGraphQL.SelectionSet, Fragment { .field("docId", String?.self), .field("parentSessionId", String?.self), .field("promptName", String.self), - .field("model", String.self), - .field("optionalModels", [String].self), .field("action", String?.self), .field("pinned", Bool.self), .field("title", String?.self), - .field("tokens", Int.self), .field("messages", [Message].self), .field("createdAt", AffineGraphQL.DateTime.self), .field("updatedAt", AffineGraphQL.DateTime.self), @@ -38,14 +35,10 @@ public struct CopilotChatHistory: AffineGraphQL.SelectionSet, Fragment { public var docId: String? { __data["docId"] } public var parentSessionId: String? { __data["parentSessionId"] } public var promptName: String { __data["promptName"] } - public var model: String { __data["model"] } - public var optionalModels: [String] { __data["optionalModels"] } /// An mark identifying which view to use to display the session public var action: String? { __data["action"] } public var pinned: Bool { __data["pinned"] } public var title: String? { __data["title"] } - /// The number of tokens used in the session - public var tokens: Int { __data["tokens"] } public var messages: [Message] { __data["messages"] } public var createdAt: AffineGraphQL.DateTime { __data["createdAt"] } public var updatedAt: AffineGraphQL.DateTime { __data["updatedAt"] } diff --git a/packages/frontend/apps/ios/App/Packages/AffineGraphQL/Sources/Fragments/PaginatedCopilotChats.graphql.swift b/packages/frontend/apps/ios/App/Packages/AffineGraphQL/Sources/Fragments/PaginatedCopilotChats.graphql.swift index 3a03a5db16..f8fa7d7801 100644 --- a/packages/frontend/apps/ios/App/Packages/AffineGraphQL/Sources/Fragments/PaginatedCopilotChats.graphql.swift +++ b/packages/frontend/apps/ios/App/Packages/AffineGraphQL/Sources/Fragments/PaginatedCopilotChats.graphql.swift @@ -91,14 +91,10 @@ public struct PaginatedCopilotChats: AffineGraphQL.SelectionSet, Fragment { public var docId: String? { __data["docId"] } public var parentSessionId: String? { __data["parentSessionId"] } public var promptName: String { __data["promptName"] } - public var model: String { __data["model"] } - public var optionalModels: [String] { __data["optionalModels"] } /// An mark identifying which view to use to display the session public var action: String? { __data["action"] } public var pinned: Bool { __data["pinned"] } public var title: String? { __data["title"] } - /// The number of tokens used in the session - public var tokens: Int { __data["tokens"] } public var messages: [Message] { __data["messages"] } public var createdAt: AffineGraphQL.DateTime { __data["createdAt"] } public var updatedAt: AffineGraphQL.DateTime { __data["updatedAt"] } diff --git a/packages/frontend/apps/ios/App/Packages/AffineGraphQL/Sources/Operations/Mutations/AdminUpdateWorkspaceMutation.graphql.swift b/packages/frontend/apps/ios/App/Packages/AffineGraphQL/Sources/Operations/Mutations/AdminUpdateWorkspaceMutation.graphql.swift index 888361e25e..82a14d75a0 100644 --- a/packages/frontend/apps/ios/App/Packages/AffineGraphQL/Sources/Operations/Mutations/AdminUpdateWorkspaceMutation.graphql.swift +++ b/packages/frontend/apps/ios/App/Packages/AffineGraphQL/Sources/Operations/Mutations/AdminUpdateWorkspaceMutation.graphql.swift @@ -7,7 +7,7 @@ public class AdminUpdateWorkspaceMutation: GraphQLMutation { public static let operationName: String = "adminUpdateWorkspace" public static let operationDocument: ApolloAPI.OperationDocument = .init( definition: .init( - #"mutation adminUpdateWorkspace($input: AdminUpdateWorkspaceInput!) { adminUpdateWorkspace(input: $input) { __typename id public createdAt name avatarKey enableAi enableSharing enableUrlPreview enableDocEmbedding features owner { __typename id name email avatarUrl } memberCount publicPageCount snapshotCount snapshotSize blobCount blobSize } }"# + #"mutation adminUpdateWorkspace($input: AdminUpdateWorkspaceInput!) { adminUpdateWorkspace(input: $input) { __typename id public createdAt name avatarKey enableAi enableSharing enableUrlPreview enableDocEmbedding owner { __typename id name email avatarUrl } memberCount publicPageCount snapshotCount snapshotSize blobCount blobSize } }"# )) public var input: AdminUpdateWorkspaceInput @@ -52,7 +52,6 @@ public class AdminUpdateWorkspaceMutation: GraphQLMutation { .field("enableSharing", Bool.self), .field("enableUrlPreview", Bool.self), .field("enableDocEmbedding", Bool.self), - .field("features", [GraphQLEnum].self), .field("owner", Owner?.self), .field("memberCount", Int.self), .field("publicPageCount", Int.self), @@ -74,7 +73,6 @@ public class AdminUpdateWorkspaceMutation: GraphQLMutation { public var enableSharing: Bool { __data["enableSharing"] } public var enableUrlPreview: Bool { __data["enableUrlPreview"] } public var enableDocEmbedding: Bool { __data["enableDocEmbedding"] } - public var features: [GraphQLEnum] { __data["features"] } public var owner: Owner? { __data["owner"] } public var memberCount: Int { __data["memberCount"] } public var publicPageCount: Int { __data["publicPageCount"] } diff --git a/packages/frontend/apps/ios/App/Packages/AffineGraphQL/Sources/Operations/Mutations/ClearWorkspaceByokConfigsMutation.graphql.swift b/packages/frontend/apps/ios/App/Packages/AffineGraphQL/Sources/Operations/Mutations/ClearWorkspaceByokConfigsMutation.graphql.swift deleted file mode 100644 index 2408bbc579..0000000000 --- a/packages/frontend/apps/ios/App/Packages/AffineGraphQL/Sources/Operations/Mutations/ClearWorkspaceByokConfigsMutation.graphql.swift +++ /dev/null @@ -1,35 +0,0 @@ -// @generated -// This file was automatically generated and should not be edited. - -@_exported import ApolloAPI - -public class ClearWorkspaceByokConfigsMutation: GraphQLMutation { - public static let operationName: String = "clearWorkspaceByokConfigs" - public static let operationDocument: ApolloAPI.OperationDocument = .init( - definition: .init( - #"mutation clearWorkspaceByokConfigs($workspaceId: String!) { clearWorkspaceByokConfigs(workspaceId: $workspaceId) }"# - )) - - public var workspaceId: String - - public init(workspaceId: String) { - self.workspaceId = workspaceId - } - - public var __variables: Variables? { ["workspaceId": workspaceId] } - - public struct Data: AffineGraphQL.SelectionSet { - public let __data: DataDict - public init(_dataDict: DataDict) { __data = _dataDict } - - public static var __parentType: any ApolloAPI.ParentType { AffineGraphQL.Objects.Mutation } - public static var __selections: [ApolloAPI.Selection] { [ - .field("clearWorkspaceByokConfigs", Bool.self, arguments: ["workspaceId": .variable("workspaceId")]), - ] } - public static var __fulfilledFragments: [any ApolloAPI.SelectionSet.Type] { [ - ClearWorkspaceByokConfigsMutation.Data.self - ] } - - public var clearWorkspaceByokConfigs: Bool { __data["clearWorkspaceByokConfigs"] } - } -} diff --git a/packages/frontend/apps/ios/App/Packages/AffineGraphQL/Sources/Operations/Mutations/CreateCopilotSessionWithHistoryMutation.graphql.swift b/packages/frontend/apps/ios/App/Packages/AffineGraphQL/Sources/Operations/Mutations/CreateCopilotSessionWithHistoryMutation.graphql.swift index 15fba7665b..e9406f3ccb 100644 --- a/packages/frontend/apps/ios/App/Packages/AffineGraphQL/Sources/Operations/Mutations/CreateCopilotSessionWithHistoryMutation.graphql.swift +++ b/packages/frontend/apps/ios/App/Packages/AffineGraphQL/Sources/Operations/Mutations/CreateCopilotSessionWithHistoryMutation.graphql.swift @@ -56,14 +56,10 @@ public class CreateCopilotSessionWithHistoryMutation: GraphQLMutation { public var docId: String? { __data["docId"] } public var parentSessionId: String? { __data["parentSessionId"] } public var promptName: String { __data["promptName"] } - public var model: String { __data["model"] } - public var optionalModels: [String] { __data["optionalModels"] } /// An mark identifying which view to use to display the session public var action: String? { __data["action"] } public var pinned: Bool { __data["pinned"] } public var title: String? { __data["title"] } - /// The number of tokens used in the session - public var tokens: Int { __data["tokens"] } public var messages: [Message] { __data["messages"] } public var createdAt: AffineGraphQL.DateTime { __data["createdAt"] } public var updatedAt: AffineGraphQL.DateTime { __data["updatedAt"] } diff --git a/packages/frontend/apps/ios/App/Packages/AffineGraphQL/Sources/Operations/Mutations/CreateMcpCredentialMutation.graphql.swift b/packages/frontend/apps/ios/App/Packages/AffineGraphQL/Sources/Operations/Mutations/CreateMcpCredentialMutation.graphql.swift new file mode 100644 index 0000000000..7dc7492940 --- /dev/null +++ b/packages/frontend/apps/ios/App/Packages/AffineGraphQL/Sources/Operations/Mutations/CreateMcpCredentialMutation.graphql.swift @@ -0,0 +1,95 @@ +// @generated +// This file was automatically generated and should not be edited. + +@_exported import ApolloAPI + +public class CreateMcpCredentialMutation: GraphQLMutation { + public static let operationName: String = "createMcpCredential" + public static let operationDocument: ApolloAPI.OperationDocument = .init( + definition: .init( + #"mutation createMcpCredential($input: CreateMcpCredentialInput!) { createMcpCredential(input: $input) { __typename credential { __typename id name workspaceId accessMode fingerprint createdAt expiresAt lastUsedAt revokedAt graceEndsAt status } token } }"# + )) + + public var input: CreateMcpCredentialInput + + public init(input: CreateMcpCredentialInput) { + self.input = input + } + + public var __variables: Variables? { ["input": input] } + + public struct Data: AffineGraphQL.SelectionSet { + public let __data: DataDict + public init(_dataDict: DataDict) { __data = _dataDict } + + public static var __parentType: any ApolloAPI.ParentType { AffineGraphQL.Objects.Mutation } + public static var __selections: [ApolloAPI.Selection] { [ + .field("createMcpCredential", CreateMcpCredential.self, arguments: ["input": .variable("input")]), + ] } + public static var __fulfilledFragments: [any ApolloAPI.SelectionSet.Type] { [ + CreateMcpCredentialMutation.Data.self + ] } + + public var createMcpCredential: CreateMcpCredential { __data["createMcpCredential"] } + + /// CreateMcpCredential + /// + /// Parent Type: `RevealedMcpCredentialType` + public struct CreateMcpCredential: AffineGraphQL.SelectionSet { + public let __data: DataDict + public init(_dataDict: DataDict) { __data = _dataDict } + + public static var __parentType: any ApolloAPI.ParentType { AffineGraphQL.Objects.RevealedMcpCredentialType } + public static var __selections: [ApolloAPI.Selection] { [ + .field("__typename", String.self), + .field("credential", Credential.self), + .field("token", String.self), + ] } + public static var __fulfilledFragments: [any ApolloAPI.SelectionSet.Type] { [ + CreateMcpCredentialMutation.Data.CreateMcpCredential.self + ] } + + public var credential: Credential { __data["credential"] } + public var token: String { __data["token"] } + + /// CreateMcpCredential.Credential + /// + /// Parent Type: `McpCredentialType` + public struct Credential: AffineGraphQL.SelectionSet { + public let __data: DataDict + public init(_dataDict: DataDict) { __data = _dataDict } + + public static var __parentType: any ApolloAPI.ParentType { AffineGraphQL.Objects.McpCredentialType } + public static var __selections: [ApolloAPI.Selection] { [ + .field("__typename", String.self), + .field("id", AffineGraphQL.ID.self), + .field("name", String.self), + .field("workspaceId", String.self), + .field("accessMode", GraphQLEnum.self), + .field("fingerprint", String.self), + .field("createdAt", AffineGraphQL.DateTime.self), + .field("expiresAt", AffineGraphQL.DateTime.self), + .field("lastUsedAt", AffineGraphQL.DateTime?.self), + .field("revokedAt", AffineGraphQL.DateTime?.self), + .field("graceEndsAt", AffineGraphQL.DateTime?.self), + .field("status", GraphQLEnum.self), + ] } + public static var __fulfilledFragments: [any ApolloAPI.SelectionSet.Type] { [ + CreateMcpCredentialMutation.Data.CreateMcpCredential.Credential.self + ] } + + public var id: AffineGraphQL.ID { __data["id"] } + public var name: String { __data["name"] } + public var workspaceId: String { __data["workspaceId"] } + public var accessMode: GraphQLEnum { __data["accessMode"] } + public var fingerprint: String { __data["fingerprint"] } + public var createdAt: AffineGraphQL.DateTime { __data["createdAt"] } + public var expiresAt: AffineGraphQL.DateTime { __data["expiresAt"] } + public var lastUsedAt: AffineGraphQL.DateTime? { __data["lastUsedAt"] } + public var revokedAt: AffineGraphQL.DateTime? { __data["revokedAt"] } + public var graceEndsAt: AffineGraphQL.DateTime? { __data["graceEndsAt"] } + public var status: GraphQLEnum { __data["status"] } + } + } + } +} diff --git a/packages/frontend/apps/ios/App/Packages/AffineGraphQL/Sources/Operations/Mutations/UpsertWorkspaceByokConfigMutation.graphql.swift b/packages/frontend/apps/ios/App/Packages/AffineGraphQL/Sources/Operations/Mutations/CreateWorkspaceByokProfileMutation.graphql.swift similarity index 50% rename from packages/frontend/apps/ios/App/Packages/AffineGraphQL/Sources/Operations/Mutations/UpsertWorkspaceByokConfigMutation.graphql.swift rename to packages/frontend/apps/ios/App/Packages/AffineGraphQL/Sources/Operations/Mutations/CreateWorkspaceByokProfileMutation.graphql.swift index 34be8b91a2..d845e909d2 100644 --- a/packages/frontend/apps/ios/App/Packages/AffineGraphQL/Sources/Operations/Mutations/UpsertWorkspaceByokConfigMutation.graphql.swift +++ b/packages/frontend/apps/ios/App/Packages/AffineGraphQL/Sources/Operations/Mutations/CreateWorkspaceByokProfileMutation.graphql.swift @@ -3,16 +3,16 @@ @_exported import ApolloAPI -public class UpsertWorkspaceByokConfigMutation: GraphQLMutation { - public static let operationName: String = "upsertWorkspaceByokConfig" +public class CreateWorkspaceByokProfileMutation: GraphQLMutation { + public static let operationName: String = "createWorkspaceByokProfile" public static let operationDocument: ApolloAPI.OperationDocument = .init( definition: .init( - #"mutation upsertWorkspaceByokConfig($input: UpsertWorkspaceByokConfigInput!) { upsertWorkspaceByokConfig(input: $input) { __typename id } }"# + #"mutation createWorkspaceByokProfile($input: CreateWorkspaceByokProfileInput!) { createWorkspaceByokProfile(input: $input) { __typename profileId } }"# )) - public var input: UpsertWorkspaceByokConfigInput + public var input: CreateWorkspaceByokProfileInput - public init(input: UpsertWorkspaceByokConfigInput) { + public init(input: CreateWorkspaceByokProfileInput) { self.input = input } @@ -24,31 +24,31 @@ public class UpsertWorkspaceByokConfigMutation: GraphQLMutation { public static var __parentType: any ApolloAPI.ParentType { AffineGraphQL.Objects.Mutation } public static var __selections: [ApolloAPI.Selection] { [ - .field("upsertWorkspaceByokConfig", UpsertWorkspaceByokConfig.self, arguments: ["input": .variable("input")]), + .field("createWorkspaceByokProfile", CreateWorkspaceByokProfile.self, arguments: ["input": .variable("input")]), ] } public static var __fulfilledFragments: [any ApolloAPI.SelectionSet.Type] { [ - UpsertWorkspaceByokConfigMutation.Data.self + CreateWorkspaceByokProfileMutation.Data.self ] } - public var upsertWorkspaceByokConfig: UpsertWorkspaceByokConfig { __data["upsertWorkspaceByokConfig"] } + public var createWorkspaceByokProfile: CreateWorkspaceByokProfile { __data["createWorkspaceByokProfile"] } - /// UpsertWorkspaceByokConfig + /// CreateWorkspaceByokProfile /// - /// Parent Type: `WorkspaceByokKeyConfigType` - public struct UpsertWorkspaceByokConfig: AffineGraphQL.SelectionSet { + /// Parent Type: `WorkspaceByokProfileType` + public struct CreateWorkspaceByokProfile: AffineGraphQL.SelectionSet { public let __data: DataDict public init(_dataDict: DataDict) { __data = _dataDict } - public static var __parentType: any ApolloAPI.ParentType { AffineGraphQL.Objects.WorkspaceByokKeyConfigType } + public static var __parentType: any ApolloAPI.ParentType { AffineGraphQL.Objects.WorkspaceByokProfileType } public static var __selections: [ApolloAPI.Selection] { [ .field("__typename", String.self), - .field("id", AffineGraphQL.ID.self), + .field("profileId", AffineGraphQL.ID.self), ] } public static var __fulfilledFragments: [any ApolloAPI.SelectionSet.Type] { [ - UpsertWorkspaceByokConfigMutation.Data.UpsertWorkspaceByokConfig.self + CreateWorkspaceByokProfileMutation.Data.CreateWorkspaceByokProfile.self ] } - public var id: AffineGraphQL.ID { __data["id"] } + public var profileId: AffineGraphQL.ID { __data["profileId"] } } } } diff --git a/packages/frontend/apps/ios/App/Packages/AffineGraphQL/Sources/Operations/Mutations/DeleteAuthSigningKeyMutation.graphql.swift b/packages/frontend/apps/ios/App/Packages/AffineGraphQL/Sources/Operations/Mutations/DeleteAuthSigningKeyMutation.graphql.swift new file mode 100644 index 0000000000..90c8f781dd --- /dev/null +++ b/packages/frontend/apps/ios/App/Packages/AffineGraphQL/Sources/Operations/Mutations/DeleteAuthSigningKeyMutation.graphql.swift @@ -0,0 +1,66 @@ +// @generated +// This file was automatically generated and should not be edited. + +@_exported import ApolloAPI + +public class DeleteAuthSigningKeyMutation: GraphQLMutation { + public static let operationName: String = "deleteAuthSigningKey" + public static let operationDocument: ApolloAPI.OperationDocument = .init( + definition: .init( + #"mutation deleteAuthSigningKey($id: String!) { deleteAuthSigningKey(id: $id) { __typename id status source createdAt retiredAt verifyUntil canDelete } }"# + )) + + public var id: String + + public init(id: String) { + self.id = id + } + + public var __variables: Variables? { ["id": id] } + + public struct Data: AffineGraphQL.SelectionSet { + public let __data: DataDict + public init(_dataDict: DataDict) { __data = _dataDict } + + public static var __parentType: any ApolloAPI.ParentType { AffineGraphQL.Objects.Mutation } + public static var __selections: [ApolloAPI.Selection] { [ + .field("deleteAuthSigningKey", [DeleteAuthSigningKey].self, arguments: ["id": .variable("id")]), + ] } + public static var __fulfilledFragments: [any ApolloAPI.SelectionSet.Type] { [ + DeleteAuthSigningKeyMutation.Data.self + ] } + + public var deleteAuthSigningKey: [DeleteAuthSigningKey] { __data["deleteAuthSigningKey"] } + + /// DeleteAuthSigningKey + /// + /// Parent Type: `AuthSigningKeyType` + public struct DeleteAuthSigningKey: AffineGraphQL.SelectionSet { + public let __data: DataDict + public init(_dataDict: DataDict) { __data = _dataDict } + + public static var __parentType: any ApolloAPI.ParentType { AffineGraphQL.Objects.AuthSigningKeyType } + public static var __selections: [ApolloAPI.Selection] { [ + .field("__typename", String.self), + .field("id", String.self), + .field("status", String.self), + .field("source", String.self), + .field("createdAt", AffineGraphQL.DateTime?.self), + .field("retiredAt", AffineGraphQL.DateTime?.self), + .field("verifyUntil", AffineGraphQL.DateTime?.self), + .field("canDelete", Bool.self), + ] } + public static var __fulfilledFragments: [any ApolloAPI.SelectionSet.Type] { [ + DeleteAuthSigningKeyMutation.Data.DeleteAuthSigningKey.self + ] } + + public var id: String { __data["id"] } + public var status: String { __data["status"] } + public var source: String { __data["source"] } + public var createdAt: AffineGraphQL.DateTime? { __data["createdAt"] } + public var retiredAt: AffineGraphQL.DateTime? { __data["retiredAt"] } + public var verifyUntil: AffineGraphQL.DateTime? { __data["verifyUntil"] } + public var canDelete: Bool { __data["canDelete"] } + } + } +} diff --git a/packages/frontend/apps/ios/App/Packages/AffineGraphQL/Sources/Operations/Mutations/DeleteWorkspaceByokConfigMutation.graphql.swift b/packages/frontend/apps/ios/App/Packages/AffineGraphQL/Sources/Operations/Mutations/DeleteWorkspaceByokProfileMutation.graphql.swift similarity index 61% rename from packages/frontend/apps/ios/App/Packages/AffineGraphQL/Sources/Operations/Mutations/DeleteWorkspaceByokConfigMutation.graphql.swift rename to packages/frontend/apps/ios/App/Packages/AffineGraphQL/Sources/Operations/Mutations/DeleteWorkspaceByokProfileMutation.graphql.swift index 7b5e1b8573..d0d947ae66 100644 --- a/packages/frontend/apps/ios/App/Packages/AffineGraphQL/Sources/Operations/Mutations/DeleteWorkspaceByokConfigMutation.graphql.swift +++ b/packages/frontend/apps/ios/App/Packages/AffineGraphQL/Sources/Operations/Mutations/DeleteWorkspaceByokProfileMutation.graphql.swift @@ -3,27 +3,27 @@ @_exported import ApolloAPI -public class DeleteWorkspaceByokConfigMutation: GraphQLMutation { - public static let operationName: String = "deleteWorkspaceByokConfig" +public class DeleteWorkspaceByokProfileMutation: GraphQLMutation { + public static let operationName: String = "deleteWorkspaceByokProfile" public static let operationDocument: ApolloAPI.OperationDocument = .init( definition: .init( - #"mutation deleteWorkspaceByokConfig($workspaceId: String!, $id: ID!) { deleteWorkspaceByokConfig(workspaceId: $workspaceId, id: $id) }"# + #"mutation deleteWorkspaceByokProfile($workspaceId: String!, $profileId: ID!) { deleteWorkspaceByokProfile(workspaceId: $workspaceId, profileId: $profileId) }"# )) public var workspaceId: String - public var id: ID + public var profileId: ID public init( workspaceId: String, - id: ID + profileId: ID ) { self.workspaceId = workspaceId - self.id = id + self.profileId = profileId } public var __variables: Variables? { [ "workspaceId": workspaceId, - "id": id + "profileId": profileId ] } public struct Data: AffineGraphQL.SelectionSet { @@ -32,15 +32,15 @@ public class DeleteWorkspaceByokConfigMutation: GraphQLMutation { public static var __parentType: any ApolloAPI.ParentType { AffineGraphQL.Objects.Mutation } public static var __selections: [ApolloAPI.Selection] { [ - .field("deleteWorkspaceByokConfig", Bool.self, arguments: [ + .field("deleteWorkspaceByokProfile", Bool.self, arguments: [ "workspaceId": .variable("workspaceId"), - "id": .variable("id") + "profileId": .variable("profileId") ]), ] } public static var __fulfilledFragments: [any ApolloAPI.SelectionSet.Type] { [ - DeleteWorkspaceByokConfigMutation.Data.self + DeleteWorkspaceByokProfileMutation.Data.self ] } - public var deleteWorkspaceByokConfig: Bool { __data["deleteWorkspaceByokConfig"] } + public var deleteWorkspaceByokProfile: Bool { __data["deleteWorkspaceByokProfile"] } } } diff --git a/packages/frontend/apps/ios/App/Packages/AffineGraphQL/Sources/Operations/Mutations/GenerateUserAccessTokenMutation.graphql.swift b/packages/frontend/apps/ios/App/Packages/AffineGraphQL/Sources/Operations/Mutations/GenerateUserAccessTokenMutation.graphql.swift deleted file mode 100644 index 58c53e313a..0000000000 --- a/packages/frontend/apps/ios/App/Packages/AffineGraphQL/Sources/Operations/Mutations/GenerateUserAccessTokenMutation.graphql.swift +++ /dev/null @@ -1,62 +0,0 @@ -// @generated -// This file was automatically generated and should not be edited. - -@_exported import ApolloAPI - -public class GenerateUserAccessTokenMutation: GraphQLMutation { - public static let operationName: String = "generateUserAccessToken" - public static let operationDocument: ApolloAPI.OperationDocument = .init( - definition: .init( - #"mutation generateUserAccessToken($input: GenerateAccessTokenInput!) { generateUserAccessToken(input: $input) { __typename id name token createdAt expiresAt } }"# - )) - - public var input: GenerateAccessTokenInput - - public init(input: GenerateAccessTokenInput) { - self.input = input - } - - public var __variables: Variables? { ["input": input] } - - public struct Data: AffineGraphQL.SelectionSet { - public let __data: DataDict - public init(_dataDict: DataDict) { __data = _dataDict } - - public static var __parentType: any ApolloAPI.ParentType { AffineGraphQL.Objects.Mutation } - public static var __selections: [ApolloAPI.Selection] { [ - .field("generateUserAccessToken", GenerateUserAccessToken.self, arguments: ["input": .variable("input")]), - ] } - public static var __fulfilledFragments: [any ApolloAPI.SelectionSet.Type] { [ - GenerateUserAccessTokenMutation.Data.self - ] } - - public var generateUserAccessToken: GenerateUserAccessToken { __data["generateUserAccessToken"] } - - /// GenerateUserAccessToken - /// - /// Parent Type: `RevealedAccessToken` - public struct GenerateUserAccessToken: AffineGraphQL.SelectionSet { - public let __data: DataDict - public init(_dataDict: DataDict) { __data = _dataDict } - - public static var __parentType: any ApolloAPI.ParentType { AffineGraphQL.Objects.RevealedAccessToken } - public static var __selections: [ApolloAPI.Selection] { [ - .field("__typename", String.self), - .field("id", String.self), - .field("name", String.self), - .field("token", String.self), - .field("createdAt", AffineGraphQL.DateTime.self), - .field("expiresAt", AffineGraphQL.DateTime?.self), - ] } - public static var __fulfilledFragments: [any ApolloAPI.SelectionSet.Type] { [ - GenerateUserAccessTokenMutation.Data.GenerateUserAccessToken.self - ] } - - public var id: String { __data["id"] } - public var name: String { __data["name"] } - public var token: String { __data["token"] } - public var createdAt: AffineGraphQL.DateTime { __data["createdAt"] } - public var expiresAt: AffineGraphQL.DateTime? { __data["expiresAt"] } - } - } -} diff --git a/packages/frontend/apps/ios/App/Packages/AffineGraphQL/Sources/Operations/Mutations/ProbeWorkspaceByokProfileMutation.graphql.swift b/packages/frontend/apps/ios/App/Packages/AffineGraphQL/Sources/Operations/Mutations/ProbeWorkspaceByokProfileMutation.graphql.swift new file mode 100644 index 0000000000..dd44b92f03 --- /dev/null +++ b/packages/frontend/apps/ios/App/Packages/AffineGraphQL/Sources/Operations/Mutations/ProbeWorkspaceByokProfileMutation.graphql.swift @@ -0,0 +1,90 @@ +// @generated +// This file was automatically generated and should not be edited. + +@_exported import ApolloAPI + +public class ProbeWorkspaceByokProfileMutation: GraphQLMutation { + public static let operationName: String = "probeWorkspaceByokProfile" + public static let operationDocument: ApolloAPI.OperationDocument = .init( + definition: .init( + #"mutation probeWorkspaceByokProfile($workspaceId: String!, $profileId: ID!) { probeWorkspaceByokProfile(workspaceId: $workspaceId, profileId: $profileId) { __typename profileId probe { __typename kind testedAt errorKind } } }"# + )) + + public var workspaceId: String + public var profileId: ID + + public init( + workspaceId: String, + profileId: ID + ) { + self.workspaceId = workspaceId + self.profileId = profileId + } + + public var __variables: Variables? { [ + "workspaceId": workspaceId, + "profileId": profileId + ] } + + public struct Data: AffineGraphQL.SelectionSet { + public let __data: DataDict + public init(_dataDict: DataDict) { __data = _dataDict } + + public static var __parentType: any ApolloAPI.ParentType { AffineGraphQL.Objects.Mutation } + public static var __selections: [ApolloAPI.Selection] { [ + .field("probeWorkspaceByokProfile", ProbeWorkspaceByokProfile.self, arguments: [ + "workspaceId": .variable("workspaceId"), + "profileId": .variable("profileId") + ]), + ] } + public static var __fulfilledFragments: [any ApolloAPI.SelectionSet.Type] { [ + ProbeWorkspaceByokProfileMutation.Data.self + ] } + + public var probeWorkspaceByokProfile: ProbeWorkspaceByokProfile { __data["probeWorkspaceByokProfile"] } + + /// ProbeWorkspaceByokProfile + /// + /// Parent Type: `WorkspaceByokProfileType` + public struct ProbeWorkspaceByokProfile: AffineGraphQL.SelectionSet { + public let __data: DataDict + public init(_dataDict: DataDict) { __data = _dataDict } + + public static var __parentType: any ApolloAPI.ParentType { AffineGraphQL.Objects.WorkspaceByokProfileType } + public static var __selections: [ApolloAPI.Selection] { [ + .field("__typename", String.self), + .field("profileId", AffineGraphQL.ID.self), + .field("probe", Probe.self), + ] } + public static var __fulfilledFragments: [any ApolloAPI.SelectionSet.Type] { [ + ProbeWorkspaceByokProfileMutation.Data.ProbeWorkspaceByokProfile.self + ] } + + public var profileId: AffineGraphQL.ID { __data["profileId"] } + public var probe: Probe { __data["probe"] } + + /// ProbeWorkspaceByokProfile.Probe + /// + /// Parent Type: `WorkspaceByokProbeType` + public struct Probe: AffineGraphQL.SelectionSet { + public let __data: DataDict + public init(_dataDict: DataDict) { __data = _dataDict } + + public static var __parentType: any ApolloAPI.ParentType { AffineGraphQL.Objects.WorkspaceByokProbeType } + public static var __selections: [ApolloAPI.Selection] { [ + .field("__typename", String.self), + .field("kind", String.self), + .field("testedAt", AffineGraphQL.DateTime?.self), + .field("errorKind", String?.self), + ] } + public static var __fulfilledFragments: [any ApolloAPI.SelectionSet.Type] { [ + ProbeWorkspaceByokProfileMutation.Data.ProbeWorkspaceByokProfile.Probe.self + ] } + + public var kind: String { __data["kind"] } + public var testedAt: AffineGraphQL.DateTime? { __data["testedAt"] } + public var errorKind: String? { __data["errorKind"] } + } + } + } +} diff --git a/packages/frontend/apps/ios/App/Packages/AffineGraphQL/Sources/Operations/Mutations/ReorderWorkspaceByokConfigsMutation.graphql.swift b/packages/frontend/apps/ios/App/Packages/AffineGraphQL/Sources/Operations/Mutations/ReorderWorkspaceByokConfigsMutation.graphql.swift deleted file mode 100644 index a37e504385..0000000000 --- a/packages/frontend/apps/ios/App/Packages/AffineGraphQL/Sources/Operations/Mutations/ReorderWorkspaceByokConfigsMutation.graphql.swift +++ /dev/null @@ -1,56 +0,0 @@ -// @generated -// This file was automatically generated and should not be edited. - -@_exported import ApolloAPI - -public class ReorderWorkspaceByokConfigsMutation: GraphQLMutation { - public static let operationName: String = "reorderWorkspaceByokConfigs" - public static let operationDocument: ApolloAPI.OperationDocument = .init( - definition: .init( - #"mutation reorderWorkspaceByokConfigs($input: ReorderWorkspaceByokConfigsInput!) { reorderWorkspaceByokConfigs(input: $input) { __typename id sortOrder } }"# - )) - - public var input: ReorderWorkspaceByokConfigsInput - - public init(input: ReorderWorkspaceByokConfigsInput) { - self.input = input - } - - public var __variables: Variables? { ["input": input] } - - public struct Data: AffineGraphQL.SelectionSet { - public let __data: DataDict - public init(_dataDict: DataDict) { __data = _dataDict } - - public static var __parentType: any ApolloAPI.ParentType { AffineGraphQL.Objects.Mutation } - public static var __selections: [ApolloAPI.Selection] { [ - .field("reorderWorkspaceByokConfigs", [ReorderWorkspaceByokConfig].self, arguments: ["input": .variable("input")]), - ] } - public static var __fulfilledFragments: [any ApolloAPI.SelectionSet.Type] { [ - ReorderWorkspaceByokConfigsMutation.Data.self - ] } - - public var reorderWorkspaceByokConfigs: [ReorderWorkspaceByokConfig] { __data["reorderWorkspaceByokConfigs"] } - - /// ReorderWorkspaceByokConfig - /// - /// Parent Type: `WorkspaceByokKeyConfigType` - public struct ReorderWorkspaceByokConfig: AffineGraphQL.SelectionSet { - public let __data: DataDict - public init(_dataDict: DataDict) { __data = _dataDict } - - public static var __parentType: any ApolloAPI.ParentType { AffineGraphQL.Objects.WorkspaceByokKeyConfigType } - public static var __selections: [ApolloAPI.Selection] { [ - .field("__typename", String.self), - .field("id", AffineGraphQL.ID.self), - .field("sortOrder", AffineGraphQL.SafeInt.self), - ] } - public static var __fulfilledFragments: [any ApolloAPI.SelectionSet.Type] { [ - ReorderWorkspaceByokConfigsMutation.Data.ReorderWorkspaceByokConfig.self - ] } - - public var id: AffineGraphQL.ID { __data["id"] } - public var sortOrder: AffineGraphQL.SafeInt { __data["sortOrder"] } - } - } -} diff --git a/packages/frontend/apps/ios/App/Packages/AffineGraphQL/Sources/Operations/Mutations/ReplaceWorkspaceByokProfileMutation.graphql.swift b/packages/frontend/apps/ios/App/Packages/AffineGraphQL/Sources/Operations/Mutations/ReplaceWorkspaceByokProfileMutation.graphql.swift new file mode 100644 index 0000000000..61869e5985 --- /dev/null +++ b/packages/frontend/apps/ios/App/Packages/AffineGraphQL/Sources/Operations/Mutations/ReplaceWorkspaceByokProfileMutation.graphql.swift @@ -0,0 +1,54 @@ +// @generated +// This file was automatically generated and should not be edited. + +@_exported import ApolloAPI + +public class ReplaceWorkspaceByokProfileMutation: GraphQLMutation { + public static let operationName: String = "replaceWorkspaceByokProfile" + public static let operationDocument: ApolloAPI.OperationDocument = .init( + definition: .init( + #"mutation replaceWorkspaceByokProfile($input: ReplaceWorkspaceByokProfileInput!) { replaceWorkspaceByokProfile(input: $input) { __typename profileId } }"# + )) + + public var input: ReplaceWorkspaceByokProfileInput + + public init(input: ReplaceWorkspaceByokProfileInput) { + self.input = input + } + + public var __variables: Variables? { ["input": input] } + + public struct Data: AffineGraphQL.SelectionSet { + public let __data: DataDict + public init(_dataDict: DataDict) { __data = _dataDict } + + public static var __parentType: any ApolloAPI.ParentType { AffineGraphQL.Objects.Mutation } + public static var __selections: [ApolloAPI.Selection] { [ + .field("replaceWorkspaceByokProfile", ReplaceWorkspaceByokProfile.self, arguments: ["input": .variable("input")]), + ] } + public static var __fulfilledFragments: [any ApolloAPI.SelectionSet.Type] { [ + ReplaceWorkspaceByokProfileMutation.Data.self + ] } + + public var replaceWorkspaceByokProfile: ReplaceWorkspaceByokProfile { __data["replaceWorkspaceByokProfile"] } + + /// ReplaceWorkspaceByokProfile + /// + /// Parent Type: `WorkspaceByokProfileType` + public struct ReplaceWorkspaceByokProfile: AffineGraphQL.SelectionSet { + public let __data: DataDict + public init(_dataDict: DataDict) { __data = _dataDict } + + public static var __parentType: any ApolloAPI.ParentType { AffineGraphQL.Objects.WorkspaceByokProfileType } + public static var __selections: [ApolloAPI.Selection] { [ + .field("__typename", String.self), + .field("profileId", AffineGraphQL.ID.self), + ] } + public static var __fulfilledFragments: [any ApolloAPI.SelectionSet.Type] { [ + ReplaceWorkspaceByokProfileMutation.Data.ReplaceWorkspaceByokProfile.self + ] } + + public var profileId: AffineGraphQL.ID { __data["profileId"] } + } + } +} diff --git a/packages/frontend/apps/ios/App/Packages/AffineGraphQL/Sources/Operations/Mutations/RevokeMcpCredentialMutation.graphql.swift b/packages/frontend/apps/ios/App/Packages/AffineGraphQL/Sources/Operations/Mutations/RevokeMcpCredentialMutation.graphql.swift new file mode 100644 index 0000000000..bc2fdff83f --- /dev/null +++ b/packages/frontend/apps/ios/App/Packages/AffineGraphQL/Sources/Operations/Mutations/RevokeMcpCredentialMutation.graphql.swift @@ -0,0 +1,46 @@ +// @generated +// This file was automatically generated and should not be edited. + +@_exported import ApolloAPI + +public class RevokeMcpCredentialMutation: GraphQLMutation { + public static let operationName: String = "revokeMcpCredential" + public static let operationDocument: ApolloAPI.OperationDocument = .init( + definition: .init( + #"mutation revokeMcpCredential($id: ID!, $workspaceId: String!) { revokeMcpCredential(id: $id, workspaceId: $workspaceId) }"# + )) + + public var id: ID + public var workspaceId: String + + public init( + id: ID, + workspaceId: String + ) { + self.id = id + self.workspaceId = workspaceId + } + + public var __variables: Variables? { [ + "id": id, + "workspaceId": workspaceId + ] } + + public struct Data: AffineGraphQL.SelectionSet { + public let __data: DataDict + public init(_dataDict: DataDict) { __data = _dataDict } + + public static var __parentType: any ApolloAPI.ParentType { AffineGraphQL.Objects.Mutation } + public static var __selections: [ApolloAPI.Selection] { [ + .field("revokeMcpCredential", Bool.self, arguments: [ + "id": .variable("id"), + "workspaceId": .variable("workspaceId") + ]), + ] } + public static var __fulfilledFragments: [any ApolloAPI.SelectionSet.Type] { [ + RevokeMcpCredentialMutation.Data.self + ] } + + public var revokeMcpCredential: Bool { __data["revokeMcpCredential"] } + } +} diff --git a/packages/frontend/apps/ios/App/Packages/AffineGraphQL/Sources/Operations/Mutations/RevokeUserAccessTokenMutation.graphql.swift b/packages/frontend/apps/ios/App/Packages/AffineGraphQL/Sources/Operations/Mutations/RevokeUserAccessTokenMutation.graphql.swift deleted file mode 100644 index 619c065ad0..0000000000 --- a/packages/frontend/apps/ios/App/Packages/AffineGraphQL/Sources/Operations/Mutations/RevokeUserAccessTokenMutation.graphql.swift +++ /dev/null @@ -1,35 +0,0 @@ -// @generated -// This file was automatically generated and should not be edited. - -@_exported import ApolloAPI - -public class RevokeUserAccessTokenMutation: GraphQLMutation { - public static let operationName: String = "revokeUserAccessToken" - public static let operationDocument: ApolloAPI.OperationDocument = .init( - definition: .init( - #"mutation revokeUserAccessToken($id: String!) { revokeUserAccessToken(id: $id) }"# - )) - - public var id: String - - public init(id: String) { - self.id = id - } - - public var __variables: Variables? { ["id": id] } - - public struct Data: AffineGraphQL.SelectionSet { - public let __data: DataDict - public init(_dataDict: DataDict) { __data = _dataDict } - - public static var __parentType: any ApolloAPI.ParentType { AffineGraphQL.Objects.Mutation } - public static var __selections: [ApolloAPI.Selection] { [ - .field("revokeUserAccessToken", Bool.self, arguments: ["id": .variable("id")]), - ] } - public static var __fulfilledFragments: [any ApolloAPI.SelectionSet.Type] { [ - RevokeUserAccessTokenMutation.Data.self - ] } - - public var revokeUserAccessToken: Bool { __data["revokeUserAccessToken"] } - } -} diff --git a/packages/frontend/apps/ios/App/Packages/AffineGraphQL/Sources/Operations/Mutations/RotateAuthSigningKeyMutation.graphql.swift b/packages/frontend/apps/ios/App/Packages/AffineGraphQL/Sources/Operations/Mutations/RotateAuthSigningKeyMutation.graphql.swift new file mode 100644 index 0000000000..fe2d3357af --- /dev/null +++ b/packages/frontend/apps/ios/App/Packages/AffineGraphQL/Sources/Operations/Mutations/RotateAuthSigningKeyMutation.graphql.swift @@ -0,0 +1,66 @@ +// @generated +// This file was automatically generated and should not be edited. + +@_exported import ApolloAPI + +public class RotateAuthSigningKeyMutation: GraphQLMutation { + public static let operationName: String = "rotateAuthSigningKey" + public static let operationDocument: ApolloAPI.OperationDocument = .init( + definition: .init( + #"mutation rotateAuthSigningKey($expectedActiveKeyId: String!) { rotateAuthSigningKey(expectedActiveKeyId: $expectedActiveKeyId) { __typename id status source createdAt retiredAt verifyUntil canDelete } }"# + )) + + public var expectedActiveKeyId: String + + public init(expectedActiveKeyId: String) { + self.expectedActiveKeyId = expectedActiveKeyId + } + + public var __variables: Variables? { ["expectedActiveKeyId": expectedActiveKeyId] } + + public struct Data: AffineGraphQL.SelectionSet { + public let __data: DataDict + public init(_dataDict: DataDict) { __data = _dataDict } + + public static var __parentType: any ApolloAPI.ParentType { AffineGraphQL.Objects.Mutation } + public static var __selections: [ApolloAPI.Selection] { [ + .field("rotateAuthSigningKey", [RotateAuthSigningKey].self, arguments: ["expectedActiveKeyId": .variable("expectedActiveKeyId")]), + ] } + public static var __fulfilledFragments: [any ApolloAPI.SelectionSet.Type] { [ + RotateAuthSigningKeyMutation.Data.self + ] } + + public var rotateAuthSigningKey: [RotateAuthSigningKey] { __data["rotateAuthSigningKey"] } + + /// RotateAuthSigningKey + /// + /// Parent Type: `AuthSigningKeyType` + public struct RotateAuthSigningKey: AffineGraphQL.SelectionSet { + public let __data: DataDict + public init(_dataDict: DataDict) { __data = _dataDict } + + public static var __parentType: any ApolloAPI.ParentType { AffineGraphQL.Objects.AuthSigningKeyType } + public static var __selections: [ApolloAPI.Selection] { [ + .field("__typename", String.self), + .field("id", String.self), + .field("status", String.self), + .field("source", String.self), + .field("createdAt", AffineGraphQL.DateTime?.self), + .field("retiredAt", AffineGraphQL.DateTime?.self), + .field("verifyUntil", AffineGraphQL.DateTime?.self), + .field("canDelete", Bool.self), + ] } + public static var __fulfilledFragments: [any ApolloAPI.SelectionSet.Type] { [ + RotateAuthSigningKeyMutation.Data.RotateAuthSigningKey.self + ] } + + public var id: String { __data["id"] } + public var status: String { __data["status"] } + public var source: String { __data["source"] } + public var createdAt: AffineGraphQL.DateTime? { __data["createdAt"] } + public var retiredAt: AffineGraphQL.DateTime? { __data["retiredAt"] } + public var verifyUntil: AffineGraphQL.DateTime? { __data["verifyUntil"] } + public var canDelete: Bool { __data["canDelete"] } + } + } +} diff --git a/packages/frontend/apps/ios/App/Packages/AffineGraphQL/Sources/Operations/Mutations/RotateMcpCredentialMutation.graphql.swift b/packages/frontend/apps/ios/App/Packages/AffineGraphQL/Sources/Operations/Mutations/RotateMcpCredentialMutation.graphql.swift new file mode 100644 index 0000000000..6171ac8c6b --- /dev/null +++ b/packages/frontend/apps/ios/App/Packages/AffineGraphQL/Sources/Operations/Mutations/RotateMcpCredentialMutation.graphql.swift @@ -0,0 +1,111 @@ +// @generated +// This file was automatically generated and should not be edited. + +@_exported import ApolloAPI + +public class RotateMcpCredentialMutation: GraphQLMutation { + public static let operationName: String = "rotateMcpCredential" + public static let operationDocument: ApolloAPI.OperationDocument = .init( + definition: .init( + #"mutation rotateMcpCredential($id: ID!, $workspaceId: String!, $expirationDays: Int!) { rotateMcpCredential( id: $id workspaceId: $workspaceId expirationDays: $expirationDays ) { __typename credential { __typename id name workspaceId accessMode fingerprint createdAt expiresAt lastUsedAt revokedAt graceEndsAt status } token } }"# + )) + + public var id: ID + public var workspaceId: String + public var expirationDays: Int + + public init( + id: ID, + workspaceId: String, + expirationDays: Int + ) { + self.id = id + self.workspaceId = workspaceId + self.expirationDays = expirationDays + } + + public var __variables: Variables? { [ + "id": id, + "workspaceId": workspaceId, + "expirationDays": expirationDays + ] } + + public struct Data: AffineGraphQL.SelectionSet { + public let __data: DataDict + public init(_dataDict: DataDict) { __data = _dataDict } + + public static var __parentType: any ApolloAPI.ParentType { AffineGraphQL.Objects.Mutation } + public static var __selections: [ApolloAPI.Selection] { [ + .field("rotateMcpCredential", RotateMcpCredential.self, arguments: [ + "id": .variable("id"), + "workspaceId": .variable("workspaceId"), + "expirationDays": .variable("expirationDays") + ]), + ] } + public static var __fulfilledFragments: [any ApolloAPI.SelectionSet.Type] { [ + RotateMcpCredentialMutation.Data.self + ] } + + public var rotateMcpCredential: RotateMcpCredential { __data["rotateMcpCredential"] } + + /// RotateMcpCredential + /// + /// Parent Type: `RevealedMcpCredentialType` + public struct RotateMcpCredential: AffineGraphQL.SelectionSet { + public let __data: DataDict + public init(_dataDict: DataDict) { __data = _dataDict } + + public static var __parentType: any ApolloAPI.ParentType { AffineGraphQL.Objects.RevealedMcpCredentialType } + public static var __selections: [ApolloAPI.Selection] { [ + .field("__typename", String.self), + .field("credential", Credential.self), + .field("token", String.self), + ] } + public static var __fulfilledFragments: [any ApolloAPI.SelectionSet.Type] { [ + RotateMcpCredentialMutation.Data.RotateMcpCredential.self + ] } + + public var credential: Credential { __data["credential"] } + public var token: String { __data["token"] } + + /// RotateMcpCredential.Credential + /// + /// Parent Type: `McpCredentialType` + public struct Credential: AffineGraphQL.SelectionSet { + public let __data: DataDict + public init(_dataDict: DataDict) { __data = _dataDict } + + public static var __parentType: any ApolloAPI.ParentType { AffineGraphQL.Objects.McpCredentialType } + public static var __selections: [ApolloAPI.Selection] { [ + .field("__typename", String.self), + .field("id", AffineGraphQL.ID.self), + .field("name", String.self), + .field("workspaceId", String.self), + .field("accessMode", GraphQLEnum.self), + .field("fingerprint", String.self), + .field("createdAt", AffineGraphQL.DateTime.self), + .field("expiresAt", AffineGraphQL.DateTime.self), + .field("lastUsedAt", AffineGraphQL.DateTime?.self), + .field("revokedAt", AffineGraphQL.DateTime?.self), + .field("graceEndsAt", AffineGraphQL.DateTime?.self), + .field("status", GraphQLEnum.self), + ] } + public static var __fulfilledFragments: [any ApolloAPI.SelectionSet.Type] { [ + RotateMcpCredentialMutation.Data.RotateMcpCredential.Credential.self + ] } + + public var id: AffineGraphQL.ID { __data["id"] } + public var name: String { __data["name"] } + public var workspaceId: String { __data["workspaceId"] } + public var accessMode: GraphQLEnum { __data["accessMode"] } + public var fingerprint: String { __data["fingerprint"] } + public var createdAt: AffineGraphQL.DateTime { __data["createdAt"] } + public var expiresAt: AffineGraphQL.DateTime { __data["expiresAt"] } + public var lastUsedAt: AffineGraphQL.DateTime? { __data["lastUsedAt"] } + public var revokedAt: AffineGraphQL.DateTime? { __data["revokedAt"] } + public var graceEndsAt: AffineGraphQL.DateTime? { __data["graceEndsAt"] } + public var status: GraphQLEnum { __data["status"] } + } + } + } +} diff --git a/packages/frontend/apps/ios/App/Packages/AffineGraphQL/Sources/Operations/Mutations/RotateWorkspaceByokCredentialMutation.graphql.swift b/packages/frontend/apps/ios/App/Packages/AffineGraphQL/Sources/Operations/Mutations/RotateWorkspaceByokCredentialMutation.graphql.swift new file mode 100644 index 0000000000..5ba0d779c8 --- /dev/null +++ b/packages/frontend/apps/ios/App/Packages/AffineGraphQL/Sources/Operations/Mutations/RotateWorkspaceByokCredentialMutation.graphql.swift @@ -0,0 +1,54 @@ +// @generated +// This file was automatically generated and should not be edited. + +@_exported import ApolloAPI + +public class RotateWorkspaceByokCredentialMutation: GraphQLMutation { + public static let operationName: String = "rotateWorkspaceByokCredential" + public static let operationDocument: ApolloAPI.OperationDocument = .init( + definition: .init( + #"mutation rotateWorkspaceByokCredential($input: RotateWorkspaceByokCredentialInput!) { rotateWorkspaceByokCredential(input: $input) { __typename profileId } }"# + )) + + public var input: RotateWorkspaceByokCredentialInput + + public init(input: RotateWorkspaceByokCredentialInput) { + self.input = input + } + + public var __variables: Variables? { ["input": input] } + + public struct Data: AffineGraphQL.SelectionSet { + public let __data: DataDict + public init(_dataDict: DataDict) { __data = _dataDict } + + public static var __parentType: any ApolloAPI.ParentType { AffineGraphQL.Objects.Mutation } + public static var __selections: [ApolloAPI.Selection] { [ + .field("rotateWorkspaceByokCredential", RotateWorkspaceByokCredential.self, arguments: ["input": .variable("input")]), + ] } + public static var __fulfilledFragments: [any ApolloAPI.SelectionSet.Type] { [ + RotateWorkspaceByokCredentialMutation.Data.self + ] } + + public var rotateWorkspaceByokCredential: RotateWorkspaceByokCredential { __data["rotateWorkspaceByokCredential"] } + + /// RotateWorkspaceByokCredential + /// + /// Parent Type: `WorkspaceByokProfileType` + public struct RotateWorkspaceByokCredential: AffineGraphQL.SelectionSet { + public let __data: DataDict + public init(_dataDict: DataDict) { __data = _dataDict } + + public static var __parentType: any ApolloAPI.ParentType { AffineGraphQL.Objects.WorkspaceByokProfileType } + public static var __selections: [ApolloAPI.Selection] { [ + .field("__typename", String.self), + .field("profileId", AffineGraphQL.ID.self), + ] } + public static var __fulfilledFragments: [any ApolloAPI.SelectionSet.Type] { [ + RotateWorkspaceByokCredentialMutation.Data.RotateWorkspaceByokCredential.self + ] } + + public var profileId: AffineGraphQL.ID { __data["profileId"] } + } + } +} diff --git a/packages/frontend/apps/ios/App/Packages/AffineGraphQL/Sources/Operations/Mutations/TestWorkspaceByokConfigMutation.graphql.swift b/packages/frontend/apps/ios/App/Packages/AffineGraphQL/Sources/Operations/Mutations/TestWorkspaceByokConfigMutation.graphql.swift deleted file mode 100644 index 857cc4cf7c..0000000000 --- a/packages/frontend/apps/ios/App/Packages/AffineGraphQL/Sources/Operations/Mutations/TestWorkspaceByokConfigMutation.graphql.swift +++ /dev/null @@ -1,58 +0,0 @@ -// @generated -// This file was automatically generated and should not be edited. - -@_exported import ApolloAPI - -public class TestWorkspaceByokConfigMutation: GraphQLMutation { - public static let operationName: String = "testWorkspaceByokConfig" - public static let operationDocument: ApolloAPI.OperationDocument = .init( - definition: .init( - #"mutation testWorkspaceByokConfig($input: TestWorkspaceByokConfigInput!) { testWorkspaceByokConfig(input: $input) { __typename ok status message } }"# - )) - - public var input: TestWorkspaceByokConfigInput - - public init(input: TestWorkspaceByokConfigInput) { - self.input = input - } - - public var __variables: Variables? { ["input": input] } - - public struct Data: AffineGraphQL.SelectionSet { - public let __data: DataDict - public init(_dataDict: DataDict) { __data = _dataDict } - - public static var __parentType: any ApolloAPI.ParentType { AffineGraphQL.Objects.Mutation } - public static var __selections: [ApolloAPI.Selection] { [ - .field("testWorkspaceByokConfig", TestWorkspaceByokConfig.self, arguments: ["input": .variable("input")]), - ] } - public static var __fulfilledFragments: [any ApolloAPI.SelectionSet.Type] { [ - TestWorkspaceByokConfigMutation.Data.self - ] } - - public var testWorkspaceByokConfig: TestWorkspaceByokConfig { __data["testWorkspaceByokConfig"] } - - /// TestWorkspaceByokConfig - /// - /// Parent Type: `TestWorkspaceByokConfigResultType` - public struct TestWorkspaceByokConfig: AffineGraphQL.SelectionSet { - public let __data: DataDict - public init(_dataDict: DataDict) { __data = _dataDict } - - public static var __parentType: any ApolloAPI.ParentType { AffineGraphQL.Objects.TestWorkspaceByokConfigResultType } - public static var __selections: [ApolloAPI.Selection] { [ - .field("__typename", String.self), - .field("ok", Bool.self), - .field("status", GraphQLEnum.self), - .field("message", String?.self), - ] } - public static var __fulfilledFragments: [any ApolloAPI.SelectionSet.Type] { [ - TestWorkspaceByokConfigMutation.Data.TestWorkspaceByokConfig.self - ] } - - public var ok: Bool { __data["ok"] } - public var status: GraphQLEnum { __data["status"] } - public var message: String? { __data["message"] } - } - } -} diff --git a/packages/frontend/apps/ios/App/Packages/AffineGraphQL/Sources/Operations/Queries/AdminDashboardQuery.graphql.swift b/packages/frontend/apps/ios/App/Packages/AffineGraphQL/Sources/Operations/Queries/AdminDashboardQuery.graphql.swift index 051d652010..0efb5b20a6 100644 --- a/packages/frontend/apps/ios/App/Packages/AffineGraphQL/Sources/Operations/Queries/AdminDashboardQuery.graphql.swift +++ b/packages/frontend/apps/ios/App/Packages/AffineGraphQL/Sources/Operations/Queries/AdminDashboardQuery.graphql.swift @@ -7,7 +7,7 @@ public class AdminDashboardQuery: GraphQLQuery { public static let operationName: String = "adminDashboard" public static let operationDocument: ApolloAPI.OperationDocument = .init( definition: .init( - #"query adminDashboard($input: AdminDashboardInput) { adminDashboard(input: $input) { __typename syncActiveUsers syncActiveUsersTimeline { __typename minute activeUsers } syncWindow { __typename from to timezone bucket requestedSize effectiveSize } copilotConversations workspaceStorageBytes blobStorageBytes workspaceStorageHistory { __typename date value } blobStorageHistory { __typename date value } storageWindow { __typename from to timezone bucket requestedSize effectiveSize } topSharedLinks { __typename workspaceId docId title shareUrl publishedAt views uniqueViews guestViews lastAccessedAt } topSharedLinksWindow { __typename from to timezone bucket requestedSize effectiveSize } generatedAt } }"# + #"query adminDashboard($input: AdminDashboardInput) { adminDashboard(input: $input) { __typename syncActiveUsers syncActiveUsersTimeline { __typename minute activeUsers } syncWindow { __typename from to timezone bucket requestedSize effectiveSize } copilotConversations copilotWindow { __typename from to timezone bucket requestedSize effectiveSize } workspaceStorageBytes blobStorageBytes workspaceStorageHistory { __typename date value } blobStorageHistory { __typename date value } storageWindow { __typename from to timezone bucket requestedSize effectiveSize } topSharedLinks { __typename workspaceId docId title shareUrl publishedAt views uniqueViews guestViews lastAccessedAt } topSharedLinksWindow { __typename from to timezone bucket requestedSize effectiveSize } generatedAt } }"# )) public var input: GraphQLNullable @@ -47,6 +47,7 @@ public class AdminDashboardQuery: GraphQLQuery { .field("syncActiveUsersTimeline", [SyncActiveUsersTimeline].self), .field("syncWindow", SyncWindow.self), .field("copilotConversations", AffineGraphQL.SafeInt.self), + .field("copilotWindow", CopilotWindow.self), .field("workspaceStorageBytes", AffineGraphQL.SafeInt.self), .field("blobStorageBytes", AffineGraphQL.SafeInt.self), .field("workspaceStorageHistory", [WorkspaceStorageHistory].self), @@ -64,6 +65,7 @@ public class AdminDashboardQuery: GraphQLQuery { public var syncActiveUsersTimeline: [SyncActiveUsersTimeline] { __data["syncActiveUsersTimeline"] } public var syncWindow: SyncWindow { __data["syncWindow"] } public var copilotConversations: AffineGraphQL.SafeInt { __data["copilotConversations"] } + public var copilotWindow: CopilotWindow { __data["copilotWindow"] } public var workspaceStorageBytes: AffineGraphQL.SafeInt { __data["workspaceStorageBytes"] } public var blobStorageBytes: AffineGraphQL.SafeInt { __data["blobStorageBytes"] } public var workspaceStorageHistory: [WorkspaceStorageHistory] { __data["workspaceStorageHistory"] } @@ -123,6 +125,35 @@ public class AdminDashboardQuery: GraphQLQuery { public var effectiveSize: Int { __data["effectiveSize"] } } + /// AdminDashboard.CopilotWindow + /// + /// Parent Type: `TimeWindow` + public struct CopilotWindow: AffineGraphQL.SelectionSet { + public let __data: DataDict + public init(_dataDict: DataDict) { __data = _dataDict } + + public static var __parentType: any ApolloAPI.ParentType { AffineGraphQL.Objects.TimeWindow } + public static var __selections: [ApolloAPI.Selection] { [ + .field("__typename", String.self), + .field("from", AffineGraphQL.DateTime.self), + .field("to", AffineGraphQL.DateTime.self), + .field("timezone", String.self), + .field("bucket", GraphQLEnum.self), + .field("requestedSize", Int.self), + .field("effectiveSize", Int.self), + ] } + public static var __fulfilledFragments: [any ApolloAPI.SelectionSet.Type] { [ + AdminDashboardQuery.Data.AdminDashboard.CopilotWindow.self + ] } + + public var from: AffineGraphQL.DateTime { __data["from"] } + public var to: AffineGraphQL.DateTime { __data["to"] } + public var timezone: String { __data["timezone"] } + public var bucket: GraphQLEnum { __data["bucket"] } + public var requestedSize: Int { __data["requestedSize"] } + public var effectiveSize: Int { __data["effectiveSize"] } + } + /// AdminDashboard.WorkspaceStorageHistory /// /// Parent Type: `AdminDashboardValueDayPoint` diff --git a/packages/frontend/apps/ios/App/Packages/AffineGraphQL/Sources/Operations/Queries/AdminMailDeliveriesQuery.graphql.swift b/packages/frontend/apps/ios/App/Packages/AffineGraphQL/Sources/Operations/Queries/AdminMailDeliveriesQuery.graphql.swift new file mode 100644 index 0000000000..d0b584cfd5 --- /dev/null +++ b/packages/frontend/apps/ios/App/Packages/AffineGraphQL/Sources/Operations/Queries/AdminMailDeliveriesQuery.graphql.swift @@ -0,0 +1,265 @@ +// @generated +// This file was automatically generated and should not be edited. + +@_exported import ApolloAPI + +public class AdminMailDeliveriesQuery: GraphQLQuery { + public static let operationName: String = "adminMailDeliveries" + public static let operationDocument: ApolloAPI.OperationDocument = .init( + definition: .init( + #"query adminMailDeliveries($input: AdminMailDeliveriesInput) { adminMailDeliveries(input: $input) { __typename window { __typename from to timezone bucket requestedSize effectiveSize } summary { __typename total sent failed skipped canceled queued sending retryWait successRate } byStatus { __typename key label total points { __typename bucket count } } byType { __typename key label total points { __typename bucket count } } byOutcome { __typename key label total points { __typename bucket count } } } }"# + )) + + public var input: GraphQLNullable + + public init(input: GraphQLNullable) { + self.input = input + } + + public var __variables: Variables? { ["input": input] } + + public struct Data: AffineGraphQL.SelectionSet { + public let __data: DataDict + public init(_dataDict: DataDict) { __data = _dataDict } + + public static var __parentType: any ApolloAPI.ParentType { AffineGraphQL.Objects.Query } + public static var __selections: [ApolloAPI.Selection] { [ + .field("adminMailDeliveries", AdminMailDeliveries.self, arguments: ["input": .variable("input")]), + ] } + public static var __fulfilledFragments: [any ApolloAPI.SelectionSet.Type] { [ + AdminMailDeliveriesQuery.Data.self + ] } + + /// Aggregate mail delivery timeline facts for admin panel + public var adminMailDeliveries: AdminMailDeliveries { __data["adminMailDeliveries"] } + + /// AdminMailDeliveries + /// + /// Parent Type: `AdminMailDeliveryAnalytics` + public struct AdminMailDeliveries: AffineGraphQL.SelectionSet { + public let __data: DataDict + public init(_dataDict: DataDict) { __data = _dataDict } + + public static var __parentType: any ApolloAPI.ParentType { AffineGraphQL.Objects.AdminMailDeliveryAnalytics } + public static var __selections: [ApolloAPI.Selection] { [ + .field("__typename", String.self), + .field("window", Window.self), + .field("summary", Summary.self), + .field("byStatus", [ByStatus].self), + .field("byType", [ByType].self), + .field("byOutcome", [ByOutcome].self), + ] } + public static var __fulfilledFragments: [any ApolloAPI.SelectionSet.Type] { [ + AdminMailDeliveriesQuery.Data.AdminMailDeliveries.self + ] } + + public var window: Window { __data["window"] } + public var summary: Summary { __data["summary"] } + public var byStatus: [ByStatus] { __data["byStatus"] } + public var byType: [ByType] { __data["byType"] } + public var byOutcome: [ByOutcome] { __data["byOutcome"] } + + /// AdminMailDeliveries.Window + /// + /// Parent Type: `TimeWindow` + public struct Window: AffineGraphQL.SelectionSet { + public let __data: DataDict + public init(_dataDict: DataDict) { __data = _dataDict } + + public static var __parentType: any ApolloAPI.ParentType { AffineGraphQL.Objects.TimeWindow } + public static var __selections: [ApolloAPI.Selection] { [ + .field("__typename", String.self), + .field("from", AffineGraphQL.DateTime.self), + .field("to", AffineGraphQL.DateTime.self), + .field("timezone", String.self), + .field("bucket", GraphQLEnum.self), + .field("requestedSize", Int.self), + .field("effectiveSize", Int.self), + ] } + public static var __fulfilledFragments: [any ApolloAPI.SelectionSet.Type] { [ + AdminMailDeliveriesQuery.Data.AdminMailDeliveries.Window.self + ] } + + public var from: AffineGraphQL.DateTime { __data["from"] } + public var to: AffineGraphQL.DateTime { __data["to"] } + public var timezone: String { __data["timezone"] } + public var bucket: GraphQLEnum { __data["bucket"] } + public var requestedSize: Int { __data["requestedSize"] } + public var effectiveSize: Int { __data["effectiveSize"] } + } + + /// AdminMailDeliveries.Summary + /// + /// Parent Type: `AdminMailDeliverySummary` + public struct Summary: AffineGraphQL.SelectionSet { + public let __data: DataDict + public init(_dataDict: DataDict) { __data = _dataDict } + + public static var __parentType: any ApolloAPI.ParentType { AffineGraphQL.Objects.AdminMailDeliverySummary } + public static var __selections: [ApolloAPI.Selection] { [ + .field("__typename", String.self), + .field("total", Int.self), + .field("sent", Int.self), + .field("failed", Int.self), + .field("skipped", Int.self), + .field("canceled", Int.self), + .field("queued", Int.self), + .field("sending", Int.self), + .field("retryWait", Int.self), + .field("successRate", Double.self), + ] } + public static var __fulfilledFragments: [any ApolloAPI.SelectionSet.Type] { [ + AdminMailDeliveriesQuery.Data.AdminMailDeliveries.Summary.self + ] } + + public var total: Int { __data["total"] } + public var sent: Int { __data["sent"] } + public var failed: Int { __data["failed"] } + public var skipped: Int { __data["skipped"] } + public var canceled: Int { __data["canceled"] } + public var queued: Int { __data["queued"] } + public var sending: Int { __data["sending"] } + public var retryWait: Int { __data["retryWait"] } + public var successRate: Double { __data["successRate"] } + } + + /// AdminMailDeliveries.ByStatus + /// + /// Parent Type: `AdminMailDeliverySeries` + public struct ByStatus: AffineGraphQL.SelectionSet { + public let __data: DataDict + public init(_dataDict: DataDict) { __data = _dataDict } + + public static var __parentType: any ApolloAPI.ParentType { AffineGraphQL.Objects.AdminMailDeliverySeries } + public static var __selections: [ApolloAPI.Selection] { [ + .field("__typename", String.self), + .field("key", String.self), + .field("label", String.self), + .field("total", Int.self), + .field("points", [Point].self), + ] } + public static var __fulfilledFragments: [any ApolloAPI.SelectionSet.Type] { [ + AdminMailDeliveriesQuery.Data.AdminMailDeliveries.ByStatus.self + ] } + + public var key: String { __data["key"] } + public var label: String { __data["label"] } + public var total: Int { __data["total"] } + public var points: [Point] { __data["points"] } + + /// AdminMailDeliveries.ByStatus.Point + /// + /// Parent Type: `AdminMailDeliveryPoint` + public struct Point: AffineGraphQL.SelectionSet { + public let __data: DataDict + public init(_dataDict: DataDict) { __data = _dataDict } + + public static var __parentType: any ApolloAPI.ParentType { AffineGraphQL.Objects.AdminMailDeliveryPoint } + public static var __selections: [ApolloAPI.Selection] { [ + .field("__typename", String.self), + .field("bucket", AffineGraphQL.DateTime.self), + .field("count", Int.self), + ] } + public static var __fulfilledFragments: [any ApolloAPI.SelectionSet.Type] { [ + AdminMailDeliveriesQuery.Data.AdminMailDeliveries.ByStatus.Point.self + ] } + + public var bucket: AffineGraphQL.DateTime { __data["bucket"] } + public var count: Int { __data["count"] } + } + } + + /// AdminMailDeliveries.ByType + /// + /// Parent Type: `AdminMailDeliverySeries` + public struct ByType: AffineGraphQL.SelectionSet { + public let __data: DataDict + public init(_dataDict: DataDict) { __data = _dataDict } + + public static var __parentType: any ApolloAPI.ParentType { AffineGraphQL.Objects.AdminMailDeliverySeries } + public static var __selections: [ApolloAPI.Selection] { [ + .field("__typename", String.self), + .field("key", String.self), + .field("label", String.self), + .field("total", Int.self), + .field("points", [Point].self), + ] } + public static var __fulfilledFragments: [any ApolloAPI.SelectionSet.Type] { [ + AdminMailDeliveriesQuery.Data.AdminMailDeliveries.ByType.self + ] } + + public var key: String { __data["key"] } + public var label: String { __data["label"] } + public var total: Int { __data["total"] } + public var points: [Point] { __data["points"] } + + /// AdminMailDeliveries.ByType.Point + /// + /// Parent Type: `AdminMailDeliveryPoint` + public struct Point: AffineGraphQL.SelectionSet { + public let __data: DataDict + public init(_dataDict: DataDict) { __data = _dataDict } + + public static var __parentType: any ApolloAPI.ParentType { AffineGraphQL.Objects.AdminMailDeliveryPoint } + public static var __selections: [ApolloAPI.Selection] { [ + .field("__typename", String.self), + .field("bucket", AffineGraphQL.DateTime.self), + .field("count", Int.self), + ] } + public static var __fulfilledFragments: [any ApolloAPI.SelectionSet.Type] { [ + AdminMailDeliveriesQuery.Data.AdminMailDeliveries.ByType.Point.self + ] } + + public var bucket: AffineGraphQL.DateTime { __data["bucket"] } + public var count: Int { __data["count"] } + } + } + + /// AdminMailDeliveries.ByOutcome + /// + /// Parent Type: `AdminMailDeliverySeries` + public struct ByOutcome: AffineGraphQL.SelectionSet { + public let __data: DataDict + public init(_dataDict: DataDict) { __data = _dataDict } + + public static var __parentType: any ApolloAPI.ParentType { AffineGraphQL.Objects.AdminMailDeliverySeries } + public static var __selections: [ApolloAPI.Selection] { [ + .field("__typename", String.self), + .field("key", String.self), + .field("label", String.self), + .field("total", Int.self), + .field("points", [Point].self), + ] } + public static var __fulfilledFragments: [any ApolloAPI.SelectionSet.Type] { [ + AdminMailDeliveriesQuery.Data.AdminMailDeliveries.ByOutcome.self + ] } + + public var key: String { __data["key"] } + public var label: String { __data["label"] } + public var total: Int { __data["total"] } + public var points: [Point] { __data["points"] } + + /// AdminMailDeliveries.ByOutcome.Point + /// + /// Parent Type: `AdminMailDeliveryPoint` + public struct Point: AffineGraphQL.SelectionSet { + public let __data: DataDict + public init(_dataDict: DataDict) { __data = _dataDict } + + public static var __parentType: any ApolloAPI.ParentType { AffineGraphQL.Objects.AdminMailDeliveryPoint } + public static var __selections: [ApolloAPI.Selection] { [ + .field("__typename", String.self), + .field("bucket", AffineGraphQL.DateTime.self), + .field("count", Int.self), + ] } + public static var __fulfilledFragments: [any ApolloAPI.SelectionSet.Type] { [ + AdminMailDeliveriesQuery.Data.AdminMailDeliveries.ByOutcome.Point.self + ] } + + public var bucket: AffineGraphQL.DateTime { __data["bucket"] } + public var count: Int { __data["count"] } + } + } + } + } +} diff --git a/packages/frontend/apps/ios/App/Packages/AffineGraphQL/Sources/Operations/Queries/AdminServerConfigQuery.graphql.swift b/packages/frontend/apps/ios/App/Packages/AffineGraphQL/Sources/Operations/Queries/AdminServerConfigQuery.graphql.swift index ca75e4db40..93d4ccf8c1 100644 --- a/packages/frontend/apps/ios/App/Packages/AffineGraphQL/Sources/Operations/Queries/AdminServerConfigQuery.graphql.swift +++ b/packages/frontend/apps/ios/App/Packages/AffineGraphQL/Sources/Operations/Queries/AdminServerConfigQuery.graphql.swift @@ -7,7 +7,7 @@ public class AdminServerConfigQuery: GraphQLQuery { public static let operationName: String = "adminServerConfig" public static let operationDocument: ApolloAPI.OperationDocument = .init( definition: .init( - #"query adminServerConfig { serverConfig { __typename version baseUrl name features type initialized credentialsRequirement { __typename ...CredentialsRequirements } availableUpgrade { __typename changelog version publishedAt url } availableUserFeatures availableWorkspaceFeatures } }"#, + #"query adminServerConfig { serverConfig { __typename version baseUrl name features type initialized credentialsRequirement { __typename ...CredentialsRequirements } availableUpgrade { __typename changelog version publishedAt url } availableUserFeatures } }"#, fragments: [CredentialsRequirements.self, PasswordLimits.self] )) @@ -47,7 +47,6 @@ public class AdminServerConfigQuery: GraphQLQuery { .field("credentialsRequirement", CredentialsRequirement.self), .field("availableUpgrade", AvailableUpgrade?.self), .field("availableUserFeatures", [GraphQLEnum].self), - .field("availableWorkspaceFeatures", [GraphQLEnum].self), ] } public static var __fulfilledFragments: [any ApolloAPI.SelectionSet.Type] { [ AdminServerConfigQuery.Data.ServerConfig.self @@ -71,8 +70,6 @@ public class AdminServerConfigQuery: GraphQLQuery { public var availableUpgrade: AvailableUpgrade? { __data["availableUpgrade"] } /// Features for user that can be configured public var availableUserFeatures: [GraphQLEnum] { __data["availableUserFeatures"] } - /// Workspace features available for admin configuration - public var availableWorkspaceFeatures: [GraphQLEnum] { __data["availableWorkspaceFeatures"] } /// ServerConfig.CredentialsRequirement /// diff --git a/packages/frontend/apps/ios/App/Packages/AffineGraphQL/Sources/Operations/Queries/AdminWorkspaceQuery.graphql.swift b/packages/frontend/apps/ios/App/Packages/AffineGraphQL/Sources/Operations/Queries/AdminWorkspaceQuery.graphql.swift index a273a6f8d9..dd7bcd7c3d 100644 --- a/packages/frontend/apps/ios/App/Packages/AffineGraphQL/Sources/Operations/Queries/AdminWorkspaceQuery.graphql.swift +++ b/packages/frontend/apps/ios/App/Packages/AffineGraphQL/Sources/Operations/Queries/AdminWorkspaceQuery.graphql.swift @@ -7,7 +7,7 @@ public class AdminWorkspaceQuery: GraphQLQuery { public static let operationName: String = "adminWorkspace" public static let operationDocument: ApolloAPI.OperationDocument = .init( definition: .init( - #"query adminWorkspace($id: String!, $memberSkip: Int, $memberTake: Int, $memberQuery: String) { adminWorkspace(id: $id) { __typename id public createdAt name avatarKey enableAi enableSharing enableUrlPreview enableDocEmbedding features owner { __typename id name email avatarUrl } memberCount publicPageCount snapshotCount snapshotSize blobCount blobSize sharedLinks { __typename docId title publishedAt } members(skip: $memberSkip, take: $memberTake, query: $memberQuery) { __typename id name email avatarUrl role status } } }"# + #"query adminWorkspace($id: String!, $memberSkip: Int, $memberTake: Int, $memberQuery: String) { adminWorkspace(id: $id) { __typename id public createdAt name avatarKey enableAi enableSharing enableUrlPreview enableDocEmbedding owner { __typename id name email avatarUrl } memberCount publicPageCount snapshotCount snapshotSize blobCount blobSize sharedLinks { __typename docId title publishedAt } members(skip: $memberSkip, take: $memberTake, query: $memberQuery) { __typename id name email avatarUrl role status } } }"# )) public var id: String @@ -68,7 +68,6 @@ public class AdminWorkspaceQuery: GraphQLQuery { .field("enableSharing", Bool.self), .field("enableUrlPreview", Bool.self), .field("enableDocEmbedding", Bool.self), - .field("features", [GraphQLEnum].self), .field("owner", Owner?.self), .field("memberCount", Int.self), .field("publicPageCount", Int.self), @@ -96,7 +95,6 @@ public class AdminWorkspaceQuery: GraphQLQuery { public var enableSharing: Bool { __data["enableSharing"] } public var enableUrlPreview: Bool { __data["enableUrlPreview"] } public var enableDocEmbedding: Bool { __data["enableDocEmbedding"] } - public var features: [GraphQLEnum] { __data["features"] } public var owner: Owner? { __data["owner"] } public var memberCount: Int { __data["memberCount"] } public var publicPageCount: Int { __data["publicPageCount"] } @@ -170,7 +168,7 @@ public class AdminWorkspaceQuery: GraphQLQuery { .field("name", String.self), .field("email", String.self), .field("avatarUrl", String?.self), - .field("role", GraphQLEnum.self), + .field("role", GraphQLEnum.self), .field("status", GraphQLEnum.self), ] } public static var __fulfilledFragments: [any ApolloAPI.SelectionSet.Type] { [ @@ -181,7 +179,7 @@ public class AdminWorkspaceQuery: GraphQLQuery { public var name: String { __data["name"] } public var email: String { __data["email"] } public var avatarUrl: String? { __data["avatarUrl"] } - public var role: GraphQLEnum { __data["role"] } + public var role: GraphQLEnum { __data["role"] } public var status: GraphQLEnum { __data["status"] } } } diff --git a/packages/frontend/apps/ios/App/Packages/AffineGraphQL/Sources/Operations/Queries/AdminWorkspacesQuery.graphql.swift b/packages/frontend/apps/ios/App/Packages/AffineGraphQL/Sources/Operations/Queries/AdminWorkspacesQuery.graphql.swift index 257c7a5726..292545840b 100644 --- a/packages/frontend/apps/ios/App/Packages/AffineGraphQL/Sources/Operations/Queries/AdminWorkspacesQuery.graphql.swift +++ b/packages/frontend/apps/ios/App/Packages/AffineGraphQL/Sources/Operations/Queries/AdminWorkspacesQuery.graphql.swift @@ -7,7 +7,7 @@ public class AdminWorkspacesQuery: GraphQLQuery { public static let operationName: String = "adminWorkspaces" public static let operationDocument: ApolloAPI.OperationDocument = .init( definition: .init( - #"query adminWorkspaces($filter: ListWorkspaceInput!) { adminWorkspaces(filter: $filter) { __typename id public createdAt name avatarKey enableAi enableSharing enableUrlPreview enableDocEmbedding features owner { __typename id name email avatarUrl } memberCount publicPageCount snapshotCount snapshotSize blobCount blobSize } }"# + #"query adminWorkspaces($filter: ListWorkspaceInput!) { adminWorkspaces(filter: $filter) { __typename id public createdAt name avatarKey enableAi enableSharing enableUrlPreview enableDocEmbedding owner { __typename id name email avatarUrl } memberCount publicPageCount snapshotCount snapshotSize blobCount blobSize } }"# )) public var filter: ListWorkspaceInput @@ -52,7 +52,6 @@ public class AdminWorkspacesQuery: GraphQLQuery { .field("enableSharing", Bool.self), .field("enableUrlPreview", Bool.self), .field("enableDocEmbedding", Bool.self), - .field("features", [GraphQLEnum].self), .field("owner", Owner?.self), .field("memberCount", Int.self), .field("publicPageCount", Int.self), @@ -74,7 +73,6 @@ public class AdminWorkspacesQuery: GraphQLQuery { public var enableSharing: Bool { __data["enableSharing"] } public var enableUrlPreview: Bool { __data["enableUrlPreview"] } public var enableDocEmbedding: Bool { __data["enableDocEmbedding"] } - public var features: [GraphQLEnum] { __data["features"] } public var owner: Owner? { __data["owner"] } public var memberCount: Int { __data["memberCount"] } public var publicPageCount: Int { __data["publicPageCount"] } diff --git a/packages/frontend/apps/ios/App/Packages/AffineGraphQL/Sources/Operations/Queries/AuthSigningKeysQuery.graphql.swift b/packages/frontend/apps/ios/App/Packages/AffineGraphQL/Sources/Operations/Queries/AuthSigningKeysQuery.graphql.swift new file mode 100644 index 0000000000..2b9445dc81 --- /dev/null +++ b/packages/frontend/apps/ios/App/Packages/AffineGraphQL/Sources/Operations/Queries/AuthSigningKeysQuery.graphql.swift @@ -0,0 +1,60 @@ +// @generated +// This file was automatically generated and should not be edited. + +@_exported import ApolloAPI + +public class AuthSigningKeysQuery: GraphQLQuery { + public static let operationName: String = "authSigningKeys" + public static let operationDocument: ApolloAPI.OperationDocument = .init( + definition: .init( + #"query authSigningKeys { authSigningKeys { __typename id status source createdAt retiredAt verifyUntil canDelete } }"# + )) + + public init() {} + + public struct Data: AffineGraphQL.SelectionSet { + public let __data: DataDict + public init(_dataDict: DataDict) { __data = _dataDict } + + public static var __parentType: any ApolloAPI.ParentType { AffineGraphQL.Objects.Query } + public static var __selections: [ApolloAPI.Selection] { [ + .field("authSigningKeys", [AuthSigningKey].self), + ] } + public static var __fulfilledFragments: [any ApolloAPI.SelectionSet.Type] { [ + AuthSigningKeysQuery.Data.self + ] } + + public var authSigningKeys: [AuthSigningKey] { __data["authSigningKeys"] } + + /// AuthSigningKey + /// + /// Parent Type: `AuthSigningKeyType` + public struct AuthSigningKey: AffineGraphQL.SelectionSet { + public let __data: DataDict + public init(_dataDict: DataDict) { __data = _dataDict } + + public static var __parentType: any ApolloAPI.ParentType { AffineGraphQL.Objects.AuthSigningKeyType } + public static var __selections: [ApolloAPI.Selection] { [ + .field("__typename", String.self), + .field("id", String.self), + .field("status", String.self), + .field("source", String.self), + .field("createdAt", AffineGraphQL.DateTime?.self), + .field("retiredAt", AffineGraphQL.DateTime?.self), + .field("verifyUntil", AffineGraphQL.DateTime?.self), + .field("canDelete", Bool.self), + ] } + public static var __fulfilledFragments: [any ApolloAPI.SelectionSet.Type] { [ + AuthSigningKeysQuery.Data.AuthSigningKey.self + ] } + + public var id: String { __data["id"] } + public var status: String { __data["status"] } + public var source: String { __data["source"] } + public var createdAt: AffineGraphQL.DateTime? { __data["createdAt"] } + public var retiredAt: AffineGraphQL.DateTime? { __data["retiredAt"] } + public var verifyUntil: AffineGraphQL.DateTime? { __data["verifyUntil"] } + public var canDelete: Bool { __data["canDelete"] } + } + } +} diff --git a/packages/frontend/apps/ios/App/Packages/AffineGraphQL/Sources/Operations/Queries/GetPromptModelsQuery.graphql.swift b/packages/frontend/apps/ios/App/Packages/AffineGraphQL/Sources/Operations/Queries/GetCopilotRouteOptionsQuery.graphql.swift similarity index 53% rename from packages/frontend/apps/ios/App/Packages/AffineGraphQL/Sources/Operations/Queries/GetPromptModelsQuery.graphql.swift rename to packages/frontend/apps/ios/App/Packages/AffineGraphQL/Sources/Operations/Queries/GetCopilotRouteOptionsQuery.graphql.swift index 99efb02bd0..74f9369d85 100644 --- a/packages/frontend/apps/ios/App/Packages/AffineGraphQL/Sources/Operations/Queries/GetPromptModelsQuery.graphql.swift +++ b/packages/frontend/apps/ios/App/Packages/AffineGraphQL/Sources/Operations/Queries/GetCopilotRouteOptionsQuery.graphql.swift @@ -3,11 +3,11 @@ @_exported import ApolloAPI -public class GetPromptModelsQuery: GraphQLQuery { - public static let operationName: String = "getPromptModels" +public class GetCopilotRouteOptionsQuery: GraphQLQuery { + public static let operationName: String = "getCopilotRouteOptions" public static let operationDocument: ApolloAPI.OperationDocument = .init( definition: .init( - #"query getPromptModels($promptName: String!) { currentUser { __typename copilot { __typename models(promptName: $promptName) { __typename defaultModel optionalModels { __typename id name } proModels { __typename id name } } } } }"# + #"query getCopilotRouteOptions($promptName: String!) { currentUser { __typename copilot { __typename routeOptions(promptName: $promptName) { __typename routeId defaultTargetId choices { __typename id displayName minimumTier available } } } } }"# )) public var promptName: String @@ -27,7 +27,7 @@ public class GetPromptModelsQuery: GraphQLQuery { .field("currentUser", CurrentUser?.self), ] } public static var __fulfilledFragments: [any ApolloAPI.SelectionSet.Type] { [ - GetPromptModelsQuery.Data.self + GetCopilotRouteOptionsQuery.Data.self ] } /// Get current user @@ -46,7 +46,7 @@ public class GetPromptModelsQuery: GraphQLQuery { .field("copilot", Copilot.self), ] } public static var __fulfilledFragments: [any ApolloAPI.SelectionSet.Type] { [ - GetPromptModelsQuery.Data.CurrentUser.self + GetCopilotRouteOptionsQuery.Data.CurrentUser.self ] } public var copilot: Copilot { __data["copilot"] } @@ -61,77 +61,60 @@ public class GetPromptModelsQuery: GraphQLQuery { public static var __parentType: any ApolloAPI.ParentType { AffineGraphQL.Objects.Copilot } public static var __selections: [ApolloAPI.Selection] { [ .field("__typename", String.self), - .field("models", Models.self, arguments: ["promptName": .variable("promptName")]), + .field("routeOptions", RouteOptions?.self, arguments: ["promptName": .variable("promptName")]), ] } public static var __fulfilledFragments: [any ApolloAPI.SelectionSet.Type] { [ - GetPromptModelsQuery.Data.CurrentUser.Copilot.self + GetCopilotRouteOptionsQuery.Data.CurrentUser.Copilot.self ] } - /// List available models for a prompt, with human-readable names - public var models: Models { __data["models"] } + /// List native built-in route choices for a prompt + public var routeOptions: RouteOptions? { __data["routeOptions"] } - /// CurrentUser.Copilot.Models + /// CurrentUser.Copilot.RouteOptions /// - /// Parent Type: `CopilotModelsType` - public struct Models: AffineGraphQL.SelectionSet { + /// Parent Type: `CopilotRouteOptions` + public struct RouteOptions: AffineGraphQL.SelectionSet { public let __data: DataDict public init(_dataDict: DataDict) { __data = _dataDict } - public static var __parentType: any ApolloAPI.ParentType { AffineGraphQL.Objects.CopilotModelsType } + public static var __parentType: any ApolloAPI.ParentType { AffineGraphQL.Objects.CopilotRouteOptions } public static var __selections: [ApolloAPI.Selection] { [ .field("__typename", String.self), - .field("defaultModel", String.self), - .field("optionalModels", [OptionalModel].self), - .field("proModels", [ProModel].self), + .field("routeId", String.self), + .field("defaultTargetId", String?.self), + .field("choices", [Choice].self), ] } public static var __fulfilledFragments: [any ApolloAPI.SelectionSet.Type] { [ - GetPromptModelsQuery.Data.CurrentUser.Copilot.Models.self + GetCopilotRouteOptionsQuery.Data.CurrentUser.Copilot.RouteOptions.self ] } - public var defaultModel: String { __data["defaultModel"] } - public var optionalModels: [OptionalModel] { __data["optionalModels"] } - public var proModels: [ProModel] { __data["proModels"] } + public var routeId: String { __data["routeId"] } + public var defaultTargetId: String? { __data["defaultTargetId"] } + public var choices: [Choice] { __data["choices"] } - /// CurrentUser.Copilot.Models.OptionalModel + /// CurrentUser.Copilot.RouteOptions.Choice /// - /// Parent Type: `CopilotModelType` - public struct OptionalModel: AffineGraphQL.SelectionSet { + /// Parent Type: `CopilotRouteTarget` + public struct Choice: AffineGraphQL.SelectionSet { public let __data: DataDict public init(_dataDict: DataDict) { __data = _dataDict } - public static var __parentType: any ApolloAPI.ParentType { AffineGraphQL.Objects.CopilotModelType } + public static var __parentType: any ApolloAPI.ParentType { AffineGraphQL.Objects.CopilotRouteTarget } public static var __selections: [ApolloAPI.Selection] { [ .field("__typename", String.self), .field("id", String.self), - .field("name", String.self), + .field("displayName", String.self), + .field("minimumTier", String.self), + .field("available", Bool.self), ] } public static var __fulfilledFragments: [any ApolloAPI.SelectionSet.Type] { [ - GetPromptModelsQuery.Data.CurrentUser.Copilot.Models.OptionalModel.self + GetCopilotRouteOptionsQuery.Data.CurrentUser.Copilot.RouteOptions.Choice.self ] } public var id: String { __data["id"] } - public var name: String { __data["name"] } - } - - /// CurrentUser.Copilot.Models.ProModel - /// - /// Parent Type: `CopilotModelType` - public struct ProModel: AffineGraphQL.SelectionSet { - public let __data: DataDict - public init(_dataDict: DataDict) { __data = _dataDict } - - public static var __parentType: any ApolloAPI.ParentType { AffineGraphQL.Objects.CopilotModelType } - public static var __selections: [ApolloAPI.Selection] { [ - .field("__typename", String.self), - .field("id", String.self), - .field("name", String.self), - ] } - public static var __fulfilledFragments: [any ApolloAPI.SelectionSet.Type] { [ - GetPromptModelsQuery.Data.CurrentUser.Copilot.Models.ProModel.self - ] } - - public var id: String { __data["id"] } - public var name: String { __data["name"] } + public var displayName: String { __data["displayName"] } + public var minimumTier: String { __data["minimumTier"] } + public var available: Bool { __data["available"] } } } } diff --git a/packages/frontend/apps/ios/App/Packages/AffineGraphQL/Sources/Operations/Queries/GetCurrentUserQuery.graphql.swift b/packages/frontend/apps/ios/App/Packages/AffineGraphQL/Sources/Operations/Queries/GetCurrentUserQuery.graphql.swift index 96633a2124..99fbaa4e37 100644 --- a/packages/frontend/apps/ios/App/Packages/AffineGraphQL/Sources/Operations/Queries/GetCurrentUserQuery.graphql.swift +++ b/packages/frontend/apps/ios/App/Packages/AffineGraphQL/Sources/Operations/Queries/GetCurrentUserQuery.graphql.swift @@ -7,7 +7,7 @@ public class GetCurrentUserQuery: GraphQLQuery { public static let operationName: String = "getCurrentUser" public static let operationDocument: ApolloAPI.OperationDocument = .init( definition: .init( - #"query getCurrentUser { currentUser { __typename id name email emailVerified avatarUrl token { __typename sessionToken } } }"# + #"query getCurrentUser { currentUser { __typename id name email emailVerified avatarUrl hasPassword features } }"# )) public init() {} @@ -42,7 +42,8 @@ public class GetCurrentUserQuery: GraphQLQuery { .field("email", String.self), .field("emailVerified", Bool.self), .field("avatarUrl", String?.self), - .field("token", Token.self), + .field("hasPassword", Bool?.self), + .field("features", [GraphQLEnum].self), ] } public static var __fulfilledFragments: [any ApolloAPI.SelectionSet.Type] { [ GetCurrentUserQuery.Data.CurrentUser.self @@ -57,27 +58,10 @@ public class GetCurrentUserQuery: GraphQLQuery { public var emailVerified: Bool { __data["emailVerified"] } /// User avatar url public var avatarUrl: String? { __data["avatarUrl"] } - @available(*, deprecated, message: "use native session exchange instead") - public var token: Token { __data["token"] } - - /// CurrentUser.Token - /// - /// Parent Type: `TokenType` - public struct Token: AffineGraphQL.SelectionSet { - public let __data: DataDict - public init(_dataDict: DataDict) { __data = _dataDict } - - public static var __parentType: any ApolloAPI.ParentType { AffineGraphQL.Objects.TokenType } - public static var __selections: [ApolloAPI.Selection] { [ - .field("__typename", String.self), - .field("sessionToken", String?.self), - ] } - public static var __fulfilledFragments: [any ApolloAPI.SelectionSet.Type] { [ - GetCurrentUserQuery.Data.CurrentUser.Token.self - ] } - - public var sessionToken: String? { __data["sessionToken"] } - } + /// User password has been set + public var hasPassword: Bool? { __data["hasPassword"] } + /// Enabled features of a user + public var features: [GraphQLEnum] { __data["features"] } } } } diff --git a/packages/frontend/apps/ios/App/Packages/AffineGraphQL/Sources/Operations/Queries/McpCredentialsQuery.graphql.swift b/packages/frontend/apps/ios/App/Packages/AffineGraphQL/Sources/Operations/Queries/McpCredentialsQuery.graphql.swift new file mode 100644 index 0000000000..715d62be26 --- /dev/null +++ b/packages/frontend/apps/ios/App/Packages/AffineGraphQL/Sources/Operations/Queries/McpCredentialsQuery.graphql.swift @@ -0,0 +1,76 @@ +// @generated +// This file was automatically generated and should not be edited. + +@_exported import ApolloAPI + +public class McpCredentialsQuery: GraphQLQuery { + public static let operationName: String = "mcpCredentials" + public static let operationDocument: ApolloAPI.OperationDocument = .init( + definition: .init( + #"query mcpCredentials($workspaceId: String!) { mcpCredentialReadWriteAvailable mcpCredentials(workspaceId: $workspaceId) { __typename id name workspaceId accessMode fingerprint createdAt expiresAt lastUsedAt revokedAt graceEndsAt status } }"# + )) + + public var workspaceId: String + + public init(workspaceId: String) { + self.workspaceId = workspaceId + } + + public var __variables: Variables? { ["workspaceId": workspaceId] } + + public struct Data: AffineGraphQL.SelectionSet { + public let __data: DataDict + public init(_dataDict: DataDict) { __data = _dataDict } + + public static var __parentType: any ApolloAPI.ParentType { AffineGraphQL.Objects.Query } + public static var __selections: [ApolloAPI.Selection] { [ + .field("mcpCredentialReadWriteAvailable", Bool.self), + .field("mcpCredentials", [McpCredential].self, arguments: ["workspaceId": .variable("workspaceId")]), + ] } + public static var __fulfilledFragments: [any ApolloAPI.SelectionSet.Type] { [ + McpCredentialsQuery.Data.self + ] } + + public var mcpCredentialReadWriteAvailable: Bool { __data["mcpCredentialReadWriteAvailable"] } + public var mcpCredentials: [McpCredential] { __data["mcpCredentials"] } + + /// McpCredential + /// + /// Parent Type: `McpCredentialType` + public struct McpCredential: AffineGraphQL.SelectionSet { + public let __data: DataDict + public init(_dataDict: DataDict) { __data = _dataDict } + + public static var __parentType: any ApolloAPI.ParentType { AffineGraphQL.Objects.McpCredentialType } + public static var __selections: [ApolloAPI.Selection] { [ + .field("__typename", String.self), + .field("id", AffineGraphQL.ID.self), + .field("name", String.self), + .field("workspaceId", String.self), + .field("accessMode", GraphQLEnum.self), + .field("fingerprint", String.self), + .field("createdAt", AffineGraphQL.DateTime.self), + .field("expiresAt", AffineGraphQL.DateTime.self), + .field("lastUsedAt", AffineGraphQL.DateTime?.self), + .field("revokedAt", AffineGraphQL.DateTime?.self), + .field("graceEndsAt", AffineGraphQL.DateTime?.self), + .field("status", GraphQLEnum.self), + ] } + public static var __fulfilledFragments: [any ApolloAPI.SelectionSet.Type] { [ + McpCredentialsQuery.Data.McpCredential.self + ] } + + public var id: AffineGraphQL.ID { __data["id"] } + public var name: String { __data["name"] } + public var workspaceId: String { __data["workspaceId"] } + public var accessMode: GraphQLEnum { __data["accessMode"] } + public var fingerprint: String { __data["fingerprint"] } + public var createdAt: AffineGraphQL.DateTime { __data["createdAt"] } + public var expiresAt: AffineGraphQL.DateTime { __data["expiresAt"] } + public var lastUsedAt: AffineGraphQL.DateTime? { __data["lastUsedAt"] } + public var revokedAt: AffineGraphQL.DateTime? { __data["revokedAt"] } + public var graceEndsAt: AffineGraphQL.DateTime? { __data["graceEndsAt"] } + public var status: GraphQLEnum { __data["status"] } + } + } +} diff --git a/packages/frontend/apps/ios/App/Packages/AffineGraphQL/Sources/Operations/Queries/WorkspaceByokSettingsQuery.graphql.swift b/packages/frontend/apps/ios/App/Packages/AffineGraphQL/Sources/Operations/Queries/WorkspaceByokSettingsQuery.graphql.swift index f05ae847ab..ceab96c70b 100644 --- a/packages/frontend/apps/ios/App/Packages/AffineGraphQL/Sources/Operations/Queries/WorkspaceByokSettingsQuery.graphql.swift +++ b/packages/frontend/apps/ios/App/Packages/AffineGraphQL/Sources/Operations/Queries/WorkspaceByokSettingsQuery.graphql.swift @@ -7,7 +7,7 @@ public class WorkspaceByokSettingsQuery: GraphQLQuery { public static let operationName: String = "workspaceByokSettings" public static let operationDocument: ApolloAPI.OperationDocument = .init( definition: .init( - #"query workspaceByokSettings($id: String!, $from: DateTime!, $to: DateTime!) { workspace(id: $id) { __typename id byokSettings { __typename workspaceId entitled serverEntitled localEntitled entitlementRequired allowedProviders localStorageSupported customEndpointSupported hasAiPlan keys { __typename id provider name description storage configured enabled endpoint endpointEditable sortOrder capabilities testStatus disabledReason lastTestedAt lastTestError lastUsedAt lastErrorAt lastError } warnings { __typename featureKind reason requiredProviders } } byokUsage(from: $from, to: $to) { __typename date featureKind totalTokens } } }"# + #"query workspaceByokSettings($id: String!, $from: DateTime!, $to: DateTime!) { workspace(id: $id) { __typename id byokSettings { __typename workspaceId entitled serverEntitled localEntitled allowedProviders customEndpointSupported privateEndpointSupported profiles { __typename profileId provider name description enabled sortOrder definition { __typename version endpoint { __typename kind url } models { __typename modelId capabilities { __typename input output features attachmentKinds attachmentSources } } } probe { __typename kind testedAt errorKind } } } byokUsage(from: $from, to: $to) { __typename date featureKind totalTokens } } }"# )) public var id: String @@ -84,13 +84,10 @@ public class WorkspaceByokSettingsQuery: GraphQLQuery { .field("entitled", Bool.self), .field("serverEntitled", Bool.self), .field("localEntitled", Bool.self), - .field("entitlementRequired", [String].self), .field("allowedProviders", [GraphQLEnum].self), - .field("localStorageSupported", Bool.self), .field("customEndpointSupported", Bool.self), - .field("hasAiPlan", Bool.self), - .field("keys", [Key].self), - .field("warnings", [Warning].self), + .field("privateEndpointSupported", Bool.self), + .field("profiles", [Profile].self), ] } public static var __fulfilledFragments: [any ApolloAPI.SelectionSet.Type] { [ WorkspaceByokSettingsQuery.Data.Workspace.ByokSettings.self @@ -100,88 +97,157 @@ public class WorkspaceByokSettingsQuery: GraphQLQuery { public var entitled: Bool { __data["entitled"] } public var serverEntitled: Bool { __data["serverEntitled"] } public var localEntitled: Bool { __data["localEntitled"] } - public var entitlementRequired: [String] { __data["entitlementRequired"] } public var allowedProviders: [GraphQLEnum] { __data["allowedProviders"] } - public var localStorageSupported: Bool { __data["localStorageSupported"] } public var customEndpointSupported: Bool { __data["customEndpointSupported"] } - public var hasAiPlan: Bool { __data["hasAiPlan"] } - public var keys: [Key] { __data["keys"] } - public var warnings: [Warning] { __data["warnings"] } + public var privateEndpointSupported: Bool { __data["privateEndpointSupported"] } + public var profiles: [Profile] { __data["profiles"] } - /// Workspace.ByokSettings.Key + /// Workspace.ByokSettings.Profile /// - /// Parent Type: `WorkspaceByokKeyConfigType` - public struct Key: AffineGraphQL.SelectionSet { + /// Parent Type: `WorkspaceByokProfileType` + public struct Profile: AffineGraphQL.SelectionSet { public let __data: DataDict public init(_dataDict: DataDict) { __data = _dataDict } - public static var __parentType: any ApolloAPI.ParentType { AffineGraphQL.Objects.WorkspaceByokKeyConfigType } + public static var __parentType: any ApolloAPI.ParentType { AffineGraphQL.Objects.WorkspaceByokProfileType } public static var __selections: [ApolloAPI.Selection] { [ .field("__typename", String.self), - .field("id", AffineGraphQL.ID.self), + .field("profileId", AffineGraphQL.ID.self), .field("provider", GraphQLEnum.self), .field("name", String.self), .field("description", String?.self), - .field("storage", GraphQLEnum.self), - .field("configured", Bool.self), .field("enabled", Bool.self), - .field("endpoint", String?.self), - .field("endpointEditable", Bool.self), .field("sortOrder", AffineGraphQL.SafeInt.self), - .field("capabilities", [String].self), - .field("testStatus", GraphQLEnum.self), - .field("disabledReason", String?.self), - .field("lastTestedAt", AffineGraphQL.DateTime?.self), - .field("lastTestError", String?.self), - .field("lastUsedAt", AffineGraphQL.DateTime?.self), - .field("lastErrorAt", AffineGraphQL.DateTime?.self), - .field("lastError", String?.self), + .field("definition", Definition.self), + .field("probe", Probe.self), ] } public static var __fulfilledFragments: [any ApolloAPI.SelectionSet.Type] { [ - WorkspaceByokSettingsQuery.Data.Workspace.ByokSettings.Key.self + WorkspaceByokSettingsQuery.Data.Workspace.ByokSettings.Profile.self ] } - public var id: AffineGraphQL.ID { __data["id"] } + public var profileId: AffineGraphQL.ID { __data["profileId"] } public var provider: GraphQLEnum { __data["provider"] } public var name: String { __data["name"] } public var description: String? { __data["description"] } - public var storage: GraphQLEnum { __data["storage"] } - public var configured: Bool { __data["configured"] } public var enabled: Bool { __data["enabled"] } - public var endpoint: String? { __data["endpoint"] } - public var endpointEditable: Bool { __data["endpointEditable"] } public var sortOrder: AffineGraphQL.SafeInt { __data["sortOrder"] } - public var capabilities: [String] { __data["capabilities"] } - public var testStatus: GraphQLEnum { __data["testStatus"] } - public var disabledReason: String? { __data["disabledReason"] } - public var lastTestedAt: AffineGraphQL.DateTime? { __data["lastTestedAt"] } - public var lastTestError: String? { __data["lastTestError"] } - public var lastUsedAt: AffineGraphQL.DateTime? { __data["lastUsedAt"] } - public var lastErrorAt: AffineGraphQL.DateTime? { __data["lastErrorAt"] } - public var lastError: String? { __data["lastError"] } - } + public var definition: Definition { __data["definition"] } + public var probe: Probe { __data["probe"] } - /// Workspace.ByokSettings.Warning - /// - /// Parent Type: `WorkspaceByokCapabilityWarningType` - public struct Warning: AffineGraphQL.SelectionSet { - public let __data: DataDict - public init(_dataDict: DataDict) { __data = _dataDict } + /// Workspace.ByokSettings.Profile.Definition + /// + /// Parent Type: `WorkspaceByokProfileDefinitionType` + public struct Definition: AffineGraphQL.SelectionSet { + public let __data: DataDict + public init(_dataDict: DataDict) { __data = _dataDict } - public static var __parentType: any ApolloAPI.ParentType { AffineGraphQL.Objects.WorkspaceByokCapabilityWarningType } - public static var __selections: [ApolloAPI.Selection] { [ - .field("__typename", String.self), - .field("featureKind", String.self), - .field("reason", String.self), - .field("requiredProviders", [GraphQLEnum].self), - ] } - public static var __fulfilledFragments: [any ApolloAPI.SelectionSet.Type] { [ - WorkspaceByokSettingsQuery.Data.Workspace.ByokSettings.Warning.self - ] } + public static var __parentType: any ApolloAPI.ParentType { AffineGraphQL.Objects.WorkspaceByokProfileDefinitionType } + public static var __selections: [ApolloAPI.Selection] { [ + .field("__typename", String.self), + .field("version", AffineGraphQL.SafeInt.self), + .field("endpoint", Endpoint.self), + .field("models", [Model].self), + ] } + public static var __fulfilledFragments: [any ApolloAPI.SelectionSet.Type] { [ + WorkspaceByokSettingsQuery.Data.Workspace.ByokSettings.Profile.Definition.self + ] } - public var featureKind: String { __data["featureKind"] } - public var reason: String { __data["reason"] } - public var requiredProviders: [GraphQLEnum] { __data["requiredProviders"] } + public var version: AffineGraphQL.SafeInt { __data["version"] } + public var endpoint: Endpoint { __data["endpoint"] } + public var models: [Model] { __data["models"] } + + /// Workspace.ByokSettings.Profile.Definition.Endpoint + /// + /// Parent Type: `WorkspaceByokEndpointType` + public struct Endpoint: AffineGraphQL.SelectionSet { + public let __data: DataDict + public init(_dataDict: DataDict) { __data = _dataDict } + + public static var __parentType: any ApolloAPI.ParentType { AffineGraphQL.Objects.WorkspaceByokEndpointType } + public static var __selections: [ApolloAPI.Selection] { [ + .field("__typename", String.self), + .field("kind", String.self), + .field("url", String?.self), + ] } + public static var __fulfilledFragments: [any ApolloAPI.SelectionSet.Type] { [ + WorkspaceByokSettingsQuery.Data.Workspace.ByokSettings.Profile.Definition.Endpoint.self + ] } + + public var kind: String { __data["kind"] } + public var url: String? { __data["url"] } + } + + /// Workspace.ByokSettings.Profile.Definition.Model + /// + /// Parent Type: `WorkspaceByokModelDeclarationType` + public struct Model: AffineGraphQL.SelectionSet { + public let __data: DataDict + public init(_dataDict: DataDict) { __data = _dataDict } + + public static var __parentType: any ApolloAPI.ParentType { AffineGraphQL.Objects.WorkspaceByokModelDeclarationType } + public static var __selections: [ApolloAPI.Selection] { [ + .field("__typename", String.self), + .field("modelId", String.self), + .field("capabilities", [Capability].self), + ] } + public static var __fulfilledFragments: [any ApolloAPI.SelectionSet.Type] { [ + WorkspaceByokSettingsQuery.Data.Workspace.ByokSettings.Profile.Definition.Model.self + ] } + + public var modelId: String { __data["modelId"] } + public var capabilities: [Capability] { __data["capabilities"] } + + /// Workspace.ByokSettings.Profile.Definition.Model.Capability + /// + /// Parent Type: `WorkspaceByokCapabilityType` + public struct Capability: AffineGraphQL.SelectionSet { + public let __data: DataDict + public init(_dataDict: DataDict) { __data = _dataDict } + + public static var __parentType: any ApolloAPI.ParentType { AffineGraphQL.Objects.WorkspaceByokCapabilityType } + public static var __selections: [ApolloAPI.Selection] { [ + .field("__typename", String.self), + .field("input", [String].self), + .field("output", [String].self), + .field("features", [String].self), + .field("attachmentKinds", [String].self), + .field("attachmentSources", [String].self), + ] } + public static var __fulfilledFragments: [any ApolloAPI.SelectionSet.Type] { [ + WorkspaceByokSettingsQuery.Data.Workspace.ByokSettings.Profile.Definition.Model.Capability.self + ] } + + public var input: [String] { __data["input"] } + public var output: [String] { __data["output"] } + public var features: [String] { __data["features"] } + public var attachmentKinds: [String] { __data["attachmentKinds"] } + public var attachmentSources: [String] { __data["attachmentSources"] } + } + } + } + + /// Workspace.ByokSettings.Profile.Probe + /// + /// Parent Type: `WorkspaceByokProbeType` + public struct Probe: AffineGraphQL.SelectionSet { + public let __data: DataDict + public init(_dataDict: DataDict) { __data = _dataDict } + + public static var __parentType: any ApolloAPI.ParentType { AffineGraphQL.Objects.WorkspaceByokProbeType } + public static var __selections: [ApolloAPI.Selection] { [ + .field("__typename", String.self), + .field("kind", String.self), + .field("testedAt", AffineGraphQL.DateTime?.self), + .field("errorKind", String?.self), + ] } + public static var __fulfilledFragments: [any ApolloAPI.SelectionSet.Type] { [ + WorkspaceByokSettingsQuery.Data.Workspace.ByokSettings.Profile.Probe.self + ] } + + public var kind: String { __data["kind"] } + public var testedAt: AffineGraphQL.DateTime? { __data["testedAt"] } + public var errorKind: String? { __data["errorKind"] } + } } } diff --git a/packages/frontend/apps/ios/App/Packages/AffineGraphQL/Sources/Schema/Enums/AdminWorkspaceMemberRole.graphql.swift b/packages/frontend/apps/ios/App/Packages/AffineGraphQL/Sources/Schema/Enums/AdminWorkspaceMemberRole.graphql.swift new file mode 100644 index 0000000000..61cdcd047e --- /dev/null +++ b/packages/frontend/apps/ios/App/Packages/AffineGraphQL/Sources/Schema/Enums/AdminWorkspaceMemberRole.graphql.swift @@ -0,0 +1,10 @@ +// @generated +// This file was automatically generated and should not be edited. + +import ApolloAPI + +public enum AdminWorkspaceMemberRole: String, EnumType { + case admin = "Admin" + case collaborator = "Collaborator" + case owner = "Owner" +} diff --git a/packages/frontend/apps/ios/App/Packages/AffineGraphQL/Sources/Schema/Enums/ByokKeyStorage.graphql.swift b/packages/frontend/apps/ios/App/Packages/AffineGraphQL/Sources/Schema/Enums/ByokKeyStorage.graphql.swift deleted file mode 100644 index a67918aa74..0000000000 --- a/packages/frontend/apps/ios/App/Packages/AffineGraphQL/Sources/Schema/Enums/ByokKeyStorage.graphql.swift +++ /dev/null @@ -1,9 +0,0 @@ -// @generated -// This file was automatically generated and should not be edited. - -import ApolloAPI - -public enum ByokKeyStorage: String, EnumType { - case local = "local" - case server = "server" -} diff --git a/packages/frontend/apps/ios/App/Packages/AffineGraphQL/Sources/Schema/Enums/ByokKeyTestStatus.graphql.swift b/packages/frontend/apps/ios/App/Packages/AffineGraphQL/Sources/Schema/Enums/ByokKeyTestStatus.graphql.swift deleted file mode 100644 index 10e6ab6700..0000000000 --- a/packages/frontend/apps/ios/App/Packages/AffineGraphQL/Sources/Schema/Enums/ByokKeyTestStatus.graphql.swift +++ /dev/null @@ -1,10 +0,0 @@ -// @generated -// This file was automatically generated and should not be edited. - -import ApolloAPI - -public enum ByokKeyTestStatus: String, EnumType { - case failed = "failed" - case passed = "passed" - case untested = "untested" -} diff --git a/packages/frontend/apps/ios/App/Packages/AffineGraphQL/Sources/Schema/Enums/FeatureType.graphql.swift b/packages/frontend/apps/ios/App/Packages/AffineGraphQL/Sources/Schema/Enums/FeatureType.graphql.swift index 68b7a58ee5..87704b325e 100644 --- a/packages/frontend/apps/ios/App/Packages/AffineGraphQL/Sources/Schema/Enums/FeatureType.graphql.swift +++ b/packages/frontend/apps/ios/App/Packages/AffineGraphQL/Sources/Schema/Enums/FeatureType.graphql.swift @@ -5,11 +5,4 @@ import ApolloAPI public enum FeatureType: String, EnumType { case admin = "Admin" - case freePlan = "FreePlan" - case lifetimeProPlan = "LifetimeProPlan" - case proPlan = "ProPlan" - case quotaExceededReadonlyWorkspace = "QuotaExceededReadonlyWorkspace" - case teamPlan = "TeamPlan" - case unlimitedCopilot = "UnlimitedCopilot" - case unlimitedWorkspace = "UnlimitedWorkspace" } diff --git a/packages/frontend/apps/ios/App/Packages/AffineGraphQL/Sources/Schema/Enums/McpAccessMode.graphql.swift b/packages/frontend/apps/ios/App/Packages/AffineGraphQL/Sources/Schema/Enums/McpAccessMode.graphql.swift new file mode 100644 index 0000000000..22e0b8e0ee --- /dev/null +++ b/packages/frontend/apps/ios/App/Packages/AffineGraphQL/Sources/Schema/Enums/McpAccessMode.graphql.swift @@ -0,0 +1,9 @@ +// @generated +// This file was automatically generated and should not be edited. + +import ApolloAPI + +public enum McpAccessMode: String, EnumType { + case readOnly = "READ_ONLY" + case readWrite = "READ_WRITE" +} diff --git a/packages/frontend/apps/ios/App/Packages/AffineGraphQL/Sources/Schema/Enums/McpCredentialStatus.graphql.swift b/packages/frontend/apps/ios/App/Packages/AffineGraphQL/Sources/Schema/Enums/McpCredentialStatus.graphql.swift new file mode 100644 index 0000000000..a813c2f35d --- /dev/null +++ b/packages/frontend/apps/ios/App/Packages/AffineGraphQL/Sources/Schema/Enums/McpCredentialStatus.graphql.swift @@ -0,0 +1,12 @@ +// @generated +// This file was automatically generated and should not be edited. + +import ApolloAPI + +public enum McpCredentialStatus: String, EnumType { + case active = "ACTIVE" + case expired = "EXPIRED" + case expiring = "EXPIRING" + case revoked = "REVOKED" + case rotating = "ROTATING" +} diff --git a/packages/frontend/apps/ios/App/Packages/AffineGraphQL/Sources/Schema/Enums/TimeBucket.graphql.swift b/packages/frontend/apps/ios/App/Packages/AffineGraphQL/Sources/Schema/Enums/TimeBucket.graphql.swift index 3022c170d3..64016eda40 100644 --- a/packages/frontend/apps/ios/App/Packages/AffineGraphQL/Sources/Schema/Enums/TimeBucket.graphql.swift +++ b/packages/frontend/apps/ios/App/Packages/AffineGraphQL/Sources/Schema/Enums/TimeBucket.graphql.swift @@ -5,5 +5,6 @@ import ApolloAPI public enum TimeBucket: String, EnumType { case day = "Day" + case hour = "Hour" case minute = "Minute" } diff --git a/packages/frontend/apps/ios/App/Packages/AffineGraphQL/Sources/Schema/InputObjects/AdminDashboardInput.graphql.swift b/packages/frontend/apps/ios/App/Packages/AffineGraphQL/Sources/Schema/InputObjects/AdminDashboardInput.graphql.swift index e927c22b72..83593fdc1a 100644 --- a/packages/frontend/apps/ios/App/Packages/AffineGraphQL/Sources/Schema/InputObjects/AdminDashboardInput.graphql.swift +++ b/packages/frontend/apps/ios/App/Packages/AffineGraphQL/Sources/Schema/InputObjects/AdminDashboardInput.graphql.swift @@ -11,12 +11,14 @@ public struct AdminDashboardInput: InputObject { } public init( + copilotWindowDays: GraphQLNullable = nil, sharedLinkWindowDays: GraphQLNullable = nil, storageHistoryDays: GraphQLNullable = nil, syncHistoryHours: GraphQLNullable = nil, timezone: GraphQLNullable = nil ) { __data = InputDict([ + "copilotWindowDays": copilotWindowDays, "sharedLinkWindowDays": sharedLinkWindowDays, "storageHistoryDays": storageHistoryDays, "syncHistoryHours": syncHistoryHours, @@ -24,6 +26,11 @@ public struct AdminDashboardInput: InputObject { ]) } + public var copilotWindowDays: GraphQLNullable { + get { __data["copilotWindowDays"] } + set { __data["copilotWindowDays"] = newValue } + } + public var sharedLinkWindowDays: GraphQLNullable { get { __data["sharedLinkWindowDays"] } set { __data["sharedLinkWindowDays"] = newValue } diff --git a/packages/frontend/apps/ios/App/Packages/AffineGraphQL/Sources/Schema/InputObjects/AdminMailDeliveriesInput.graphql.swift b/packages/frontend/apps/ios/App/Packages/AffineGraphQL/Sources/Schema/InputObjects/AdminMailDeliveriesInput.graphql.swift new file mode 100644 index 0000000000..6383e12831 --- /dev/null +++ b/packages/frontend/apps/ios/App/Packages/AffineGraphQL/Sources/Schema/InputObjects/AdminMailDeliveriesInput.graphql.swift @@ -0,0 +1,25 @@ +// @generated +// This file was automatically generated and should not be edited. + +import ApolloAPI + +public struct AdminMailDeliveriesInput: InputObject { + public private(set) var __data: InputDict + + public init(_ data: InputDict) { + __data = data + } + + public init( + hours: Int? = nil + ) { + __data = InputDict([ + "hours": hours + ]) + } + + public var hours: Int? { + get { __data["hours"] } + set { __data["hours"] = newValue } + } +} diff --git a/packages/frontend/apps/ios/App/Packages/AffineGraphQL/Sources/Schema/InputObjects/CreateMcpCredentialInput.graphql.swift b/packages/frontend/apps/ios/App/Packages/AffineGraphQL/Sources/Schema/InputObjects/CreateMcpCredentialInput.graphql.swift new file mode 100644 index 0000000000..c59fec840e --- /dev/null +++ b/packages/frontend/apps/ios/App/Packages/AffineGraphQL/Sources/Schema/InputObjects/CreateMcpCredentialInput.graphql.swift @@ -0,0 +1,46 @@ +// @generated +// This file was automatically generated and should not be edited. + +import ApolloAPI + +public struct CreateMcpCredentialInput: InputObject { + public private(set) var __data: InputDict + + public init(_ data: InputDict) { + __data = data + } + + public init( + accessMode: GraphQLEnum? = nil, + expirationDays: Int? = nil, + name: String, + workspaceId: String + ) { + __data = InputDict([ + "accessMode": accessMode, + "expirationDays": expirationDays, + "name": name, + "workspaceId": workspaceId + ]) + } + + public var accessMode: GraphQLEnum? { + get { __data["accessMode"] } + set { __data["accessMode"] = newValue } + } + + public var expirationDays: Int? { + get { __data["expirationDays"] } + set { __data["expirationDays"] = newValue } + } + + public var name: String { + get { __data["name"] } + set { __data["name"] = newValue } + } + + public var workspaceId: String { + get { __data["workspaceId"] } + set { __data["workspaceId"] = newValue } + } +} diff --git a/packages/frontend/apps/ios/App/Packages/AffineGraphQL/Sources/Schema/InputObjects/CreateWorkspaceByokLocalLeaseProviderInput.graphql.swift b/packages/frontend/apps/ios/App/Packages/AffineGraphQL/Sources/Schema/InputObjects/CreateWorkspaceByokLocalLeaseProviderInput.graphql.swift index d049d82256..f72c6df715 100644 --- a/packages/frontend/apps/ios/App/Packages/AffineGraphQL/Sources/Schema/InputObjects/CreateWorkspaceByokLocalLeaseProviderInput.graphql.swift +++ b/packages/frontend/apps/ios/App/Packages/AffineGraphQL/Sources/Schema/InputObjects/CreateWorkspaceByokLocalLeaseProviderInput.graphql.swift @@ -11,28 +11,31 @@ public struct CreateWorkspaceByokLocalLeaseProviderInput: InputObject { } public init( - apiKey: String, + credential: String, + definition: WorkspaceByokProfileDefinitionInput, description: GraphQLNullable = nil, - enabled: GraphQLNullable = nil, - endpoint: GraphQLNullable = nil, + enabled: Bool, name: String, - provider: GraphQLEnum, - sortOrder: GraphQLNullable = nil + provider: GraphQLEnum ) { __data = InputDict([ - "apiKey": apiKey, + "credential": credential, + "definition": definition, "description": description, "enabled": enabled, - "endpoint": endpoint, "name": name, - "provider": provider, - "sortOrder": sortOrder + "provider": provider ]) } - public var apiKey: String { - get { __data["apiKey"] } - set { __data["apiKey"] = newValue } + public var credential: String { + get { __data["credential"] } + set { __data["credential"] = newValue } + } + + public var definition: WorkspaceByokProfileDefinitionInput { + get { __data["definition"] } + set { __data["definition"] = newValue } } public var description: GraphQLNullable { @@ -40,16 +43,11 @@ public struct CreateWorkspaceByokLocalLeaseProviderInput: InputObject { set { __data["description"] = newValue } } - public var enabled: GraphQLNullable { + public var enabled: Bool { get { __data["enabled"] } set { __data["enabled"] = newValue } } - public var endpoint: GraphQLNullable { - get { __data["endpoint"] } - set { __data["endpoint"] = newValue } - } - public var name: String { get { __data["name"] } set { __data["name"] = newValue } @@ -59,9 +57,4 @@ public struct CreateWorkspaceByokLocalLeaseProviderInput: InputObject { get { __data["provider"] } set { __data["provider"] = newValue } } - - public var sortOrder: GraphQLNullable { - get { __data["sortOrder"] } - set { __data["sortOrder"] = newValue } - } } diff --git a/packages/frontend/apps/ios/App/Packages/AffineGraphQL/Sources/Schema/InputObjects/UpsertWorkspaceByokConfigInput.graphql.swift b/packages/frontend/apps/ios/App/Packages/AffineGraphQL/Sources/Schema/InputObjects/CreateWorkspaceByokProfileInput.graphql.swift similarity index 50% rename from packages/frontend/apps/ios/App/Packages/AffineGraphQL/Sources/Schema/InputObjects/UpsertWorkspaceByokConfigInput.graphql.swift rename to packages/frontend/apps/ios/App/Packages/AffineGraphQL/Sources/Schema/InputObjects/CreateWorkspaceByokProfileInput.graphql.swift index 6ed2bee3df..6845be1d5e 100644 --- a/packages/frontend/apps/ios/App/Packages/AffineGraphQL/Sources/Schema/InputObjects/UpsertWorkspaceByokConfigInput.graphql.swift +++ b/packages/frontend/apps/ios/App/Packages/AffineGraphQL/Sources/Schema/InputObjects/CreateWorkspaceByokProfileInput.graphql.swift @@ -3,7 +3,7 @@ import ApolloAPI -public struct UpsertWorkspaceByokConfigInput: InputObject { +public struct CreateWorkspaceByokProfileInput: InputObject { public private(set) var __data: InputDict public init(_ data: InputDict) { @@ -11,34 +11,33 @@ public struct UpsertWorkspaceByokConfigInput: InputObject { } public init( - apiKey: GraphQLNullable = nil, + credential: String, + definition: WorkspaceByokProfileDefinitionInput, description: GraphQLNullable = nil, - enabled: GraphQLNullable = nil, - endpoint: GraphQLNullable = nil, - id: GraphQLNullable = nil, + enabled: Bool, name: String, provider: GraphQLEnum, - sortOrder: GraphQLNullable = nil, - storage: GraphQLEnum, workspaceId: String ) { __data = InputDict([ - "apiKey": apiKey, + "credential": credential, + "definition": definition, "description": description, "enabled": enabled, - "endpoint": endpoint, - "id": id, "name": name, "provider": provider, - "sortOrder": sortOrder, - "storage": storage, "workspaceId": workspaceId ]) } - public var apiKey: GraphQLNullable { - get { __data["apiKey"] } - set { __data["apiKey"] = newValue } + public var credential: String { + get { __data["credential"] } + set { __data["credential"] = newValue } + } + + public var definition: WorkspaceByokProfileDefinitionInput { + get { __data["definition"] } + set { __data["definition"] = newValue } } public var description: GraphQLNullable { @@ -46,21 +45,11 @@ public struct UpsertWorkspaceByokConfigInput: InputObject { set { __data["description"] = newValue } } - public var enabled: GraphQLNullable { + public var enabled: Bool { get { __data["enabled"] } set { __data["enabled"] = newValue } } - public var endpoint: GraphQLNullable { - get { __data["endpoint"] } - set { __data["endpoint"] = newValue } - } - - public var id: GraphQLNullable { - get { __data["id"] } - set { __data["id"] = newValue } - } - public var name: String { get { __data["name"] } set { __data["name"] = newValue } @@ -71,16 +60,6 @@ public struct UpsertWorkspaceByokConfigInput: InputObject { set { __data["provider"] = newValue } } - public var sortOrder: GraphQLNullable { - get { __data["sortOrder"] } - set { __data["sortOrder"] = newValue } - } - - public var storage: GraphQLEnum { - get { __data["storage"] } - set { __data["storage"] = newValue } - } - public var workspaceId: String { get { __data["workspaceId"] } set { __data["workspaceId"] = newValue } diff --git a/packages/frontend/apps/ios/App/Packages/AffineGraphQL/Sources/Schema/InputObjects/GenerateAccessTokenInput.graphql.swift b/packages/frontend/apps/ios/App/Packages/AffineGraphQL/Sources/Schema/InputObjects/GenerateAccessTokenInput.graphql.swift deleted file mode 100644 index 141892611b..0000000000 --- a/packages/frontend/apps/ios/App/Packages/AffineGraphQL/Sources/Schema/InputObjects/GenerateAccessTokenInput.graphql.swift +++ /dev/null @@ -1,32 +0,0 @@ -// @generated -// This file was automatically generated and should not be edited. - -import ApolloAPI - -public struct GenerateAccessTokenInput: InputObject { - public private(set) var __data: InputDict - - public init(_ data: InputDict) { - __data = data - } - - public init( - expiresAt: GraphQLNullable = nil, - name: String - ) { - __data = InputDict([ - "expiresAt": expiresAt, - "name": name - ]) - } - - public var expiresAt: GraphQLNullable { - get { __data["expiresAt"] } - set { __data["expiresAt"] = newValue } - } - - public var name: String { - get { __data["name"] } - set { __data["name"] = newValue } - } -} diff --git a/packages/frontend/apps/ios/App/Packages/AffineGraphQL/Sources/Schema/InputObjects/ListWorkspaceInput.graphql.swift b/packages/frontend/apps/ios/App/Packages/AffineGraphQL/Sources/Schema/InputObjects/ListWorkspaceInput.graphql.swift index 4fe70f6b36..20df8c8680 100644 --- a/packages/frontend/apps/ios/App/Packages/AffineGraphQL/Sources/Schema/InputObjects/ListWorkspaceInput.graphql.swift +++ b/packages/frontend/apps/ios/App/Packages/AffineGraphQL/Sources/Schema/InputObjects/ListWorkspaceInput.graphql.swift @@ -15,7 +15,6 @@ public struct ListWorkspaceInput: InputObject { enableDocEmbedding: GraphQLNullable = nil, enableSharing: GraphQLNullable = nil, enableUrlPreview: GraphQLNullable = nil, - features: GraphQLNullable<[GraphQLEnum]> = nil, first: Int? = nil, keyword: GraphQLNullable = nil, orderBy: GraphQLNullable> = nil, @@ -27,7 +26,6 @@ public struct ListWorkspaceInput: InputObject { "enableDocEmbedding": enableDocEmbedding, "enableSharing": enableSharing, "enableUrlPreview": enableUrlPreview, - "features": features, "first": first, "keyword": keyword, "orderBy": orderBy, @@ -56,11 +54,6 @@ public struct ListWorkspaceInput: InputObject { set { __data["enableUrlPreview"] = newValue } } - public var features: GraphQLNullable<[GraphQLEnum]> { - get { __data["features"] } - set { __data["features"] = newValue } - } - public var first: Int? { get { __data["first"] } set { __data["first"] = newValue } diff --git a/packages/frontend/apps/ios/App/Packages/AffineGraphQL/Sources/Schema/InputObjects/ReplaceWorkspaceByokProfileInput.graphql.swift b/packages/frontend/apps/ios/App/Packages/AffineGraphQL/Sources/Schema/InputObjects/ReplaceWorkspaceByokProfileInput.graphql.swift new file mode 100644 index 0000000000..5992c407f8 --- /dev/null +++ b/packages/frontend/apps/ios/App/Packages/AffineGraphQL/Sources/Schema/InputObjects/ReplaceWorkspaceByokProfileInput.graphql.swift @@ -0,0 +1,60 @@ +// @generated +// This file was automatically generated and should not be edited. + +import ApolloAPI + +public struct ReplaceWorkspaceByokProfileInput: InputObject { + public private(set) var __data: InputDict + + public init(_ data: InputDict) { + __data = data + } + + public init( + definition: WorkspaceByokProfileDefinitionInput, + description: GraphQLNullable = nil, + enabled: Bool, + name: String, + profileId: ID, + workspaceId: String + ) { + __data = InputDict([ + "definition": definition, + "description": description, + "enabled": enabled, + "name": name, + "profileId": profileId, + "workspaceId": workspaceId + ]) + } + + public var definition: WorkspaceByokProfileDefinitionInput { + get { __data["definition"] } + set { __data["definition"] = newValue } + } + + public var description: GraphQLNullable { + get { __data["description"] } + set { __data["description"] = newValue } + } + + public var enabled: Bool { + get { __data["enabled"] } + set { __data["enabled"] = newValue } + } + + public var name: String { + get { __data["name"] } + set { __data["name"] = newValue } + } + + public var profileId: ID { + get { __data["profileId"] } + set { __data["profileId"] = newValue } + } + + public var workspaceId: String { + get { __data["workspaceId"] } + set { __data["workspaceId"] = newValue } + } +} diff --git a/packages/frontend/apps/ios/App/Packages/AffineGraphQL/Sources/Schema/InputObjects/ReorderWorkspaceByokConfigsInput.graphql.swift b/packages/frontend/apps/ios/App/Packages/AffineGraphQL/Sources/Schema/InputObjects/RotateWorkspaceByokCredentialInput.graphql.swift similarity index 53% rename from packages/frontend/apps/ios/App/Packages/AffineGraphQL/Sources/Schema/InputObjects/ReorderWorkspaceByokConfigsInput.graphql.swift rename to packages/frontend/apps/ios/App/Packages/AffineGraphQL/Sources/Schema/InputObjects/RotateWorkspaceByokCredentialInput.graphql.swift index db64493001..2c9da08113 100644 --- a/packages/frontend/apps/ios/App/Packages/AffineGraphQL/Sources/Schema/InputObjects/ReorderWorkspaceByokConfigsInput.graphql.swift +++ b/packages/frontend/apps/ios/App/Packages/AffineGraphQL/Sources/Schema/InputObjects/RotateWorkspaceByokCredentialInput.graphql.swift @@ -3,7 +3,7 @@ import ApolloAPI -public struct ReorderWorkspaceByokConfigsInput: InputObject { +public struct RotateWorkspaceByokCredentialInput: InputObject { public private(set) var __data: InputDict public init(_ data: InputDict) { @@ -11,25 +11,25 @@ public struct ReorderWorkspaceByokConfigsInput: InputObject { } public init( - ids: [ID], - storage: GraphQLEnum, + credential: String, + profileId: ID, workspaceId: String ) { __data = InputDict([ - "ids": ids, - "storage": storage, + "credential": credential, + "profileId": profileId, "workspaceId": workspaceId ]) } - public var ids: [ID] { - get { __data["ids"] } - set { __data["ids"] = newValue } + public var credential: String { + get { __data["credential"] } + set { __data["credential"] = newValue } } - public var storage: GraphQLEnum { - get { __data["storage"] } - set { __data["storage"] = newValue } + public var profileId: ID { + get { __data["profileId"] } + set { __data["profileId"] = newValue } } public var workspaceId: String { diff --git a/packages/frontend/apps/ios/App/Packages/AffineGraphQL/Sources/Schema/InputObjects/SubmitAudioTranscriptionInput.graphql.swift b/packages/frontend/apps/ios/App/Packages/AffineGraphQL/Sources/Schema/InputObjects/SubmitAudioTranscriptionInput.graphql.swift index e05bed4c94..7f3fa49ded 100644 --- a/packages/frontend/apps/ios/App/Packages/AffineGraphQL/Sources/Schema/InputObjects/SubmitAudioTranscriptionInput.graphql.swift +++ b/packages/frontend/apps/ios/App/Packages/AffineGraphQL/Sources/Schema/InputObjects/SubmitAudioTranscriptionInput.graphql.swift @@ -13,14 +13,12 @@ public struct SubmitAudioTranscriptionInput: InputObject { public init( quality: GraphQLNullable = nil, sliceManifest: GraphQLNullable<[AudioSliceManifestItemInput]> = nil, - sourceAudio: GraphQLNullable = nil, - strategy: GraphQLNullable = nil + sourceAudio: GraphQLNullable = nil ) { __data = InputDict([ "quality": quality, "sliceManifest": sliceManifest, - "sourceAudio": sourceAudio, - "strategy": strategy + "sourceAudio": sourceAudio ]) } @@ -38,9 +36,4 @@ public struct SubmitAudioTranscriptionInput: InputObject { get { __data["sourceAudio"] } set { __data["sourceAudio"] = newValue } } - - public var strategy: GraphQLNullable { - get { __data["strategy"] } - set { __data["strategy"] = newValue } - } } diff --git a/packages/frontend/apps/ios/App/Packages/AffineGraphQL/Sources/Schema/InputObjects/TestWorkspaceByokConfigInput.graphql.swift b/packages/frontend/apps/ios/App/Packages/AffineGraphQL/Sources/Schema/InputObjects/TestWorkspaceByokConfigInput.graphql.swift deleted file mode 100644 index b2fd78d7e9..0000000000 --- a/packages/frontend/apps/ios/App/Packages/AffineGraphQL/Sources/Schema/InputObjects/TestWorkspaceByokConfigInput.graphql.swift +++ /dev/null @@ -1,60 +0,0 @@ -// @generated -// This file was automatically generated and should not be edited. - -import ApolloAPI - -public struct TestWorkspaceByokConfigInput: InputObject { - public private(set) var __data: InputDict - - public init(_ data: InputDict) { - __data = data - } - - public init( - apiKey: GraphQLNullable = nil, - configId: GraphQLNullable = nil, - endpoint: GraphQLNullable = nil, - provider: GraphQLEnum, - storage: GraphQLEnum, - workspaceId: String - ) { - __data = InputDict([ - "apiKey": apiKey, - "configId": configId, - "endpoint": endpoint, - "provider": provider, - "storage": storage, - "workspaceId": workspaceId - ]) - } - - public var apiKey: GraphQLNullable { - get { __data["apiKey"] } - set { __data["apiKey"] = newValue } - } - - public var configId: GraphQLNullable { - get { __data["configId"] } - set { __data["configId"] = newValue } - } - - public var endpoint: GraphQLNullable { - get { __data["endpoint"] } - set { __data["endpoint"] = newValue } - } - - public var provider: GraphQLEnum { - get { __data["provider"] } - set { __data["provider"] = newValue } - } - - public var storage: GraphQLEnum { - get { __data["storage"] } - set { __data["storage"] = newValue } - } - - public var workspaceId: String { - get { __data["workspaceId"] } - set { __data["workspaceId"] = newValue } - } -} diff --git a/packages/frontend/apps/ios/App/Packages/AffineGraphQL/Sources/Schema/InputObjects/WorkspaceByokCapabilityInput.graphql.swift b/packages/frontend/apps/ios/App/Packages/AffineGraphQL/Sources/Schema/InputObjects/WorkspaceByokCapabilityInput.graphql.swift new file mode 100644 index 0000000000..3ba9a38407 --- /dev/null +++ b/packages/frontend/apps/ios/App/Packages/AffineGraphQL/Sources/Schema/InputObjects/WorkspaceByokCapabilityInput.graphql.swift @@ -0,0 +1,53 @@ +// @generated +// This file was automatically generated and should not be edited. + +import ApolloAPI + +public struct WorkspaceByokCapabilityInput: InputObject { + public private(set) var __data: InputDict + + public init(_ data: InputDict) { + __data = data + } + + public init( + attachmentKinds: [String], + attachmentSources: [String], + features: [String], + input: [String], + output: [String] + ) { + __data = InputDict([ + "attachmentKinds": attachmentKinds, + "attachmentSources": attachmentSources, + "features": features, + "input": input, + "output": output + ]) + } + + public var attachmentKinds: [String] { + get { __data["attachmentKinds"] } + set { __data["attachmentKinds"] = newValue } + } + + public var attachmentSources: [String] { + get { __data["attachmentSources"] } + set { __data["attachmentSources"] = newValue } + } + + public var features: [String] { + get { __data["features"] } + set { __data["features"] = newValue } + } + + public var input: [String] { + get { __data["input"] } + set { __data["input"] = newValue } + } + + public var output: [String] { + get { __data["output"] } + set { __data["output"] = newValue } + } +} diff --git a/packages/frontend/apps/ios/App/Packages/AffineGraphQL/Sources/Schema/InputObjects/WorkspaceByokEndpointInput.graphql.swift b/packages/frontend/apps/ios/App/Packages/AffineGraphQL/Sources/Schema/InputObjects/WorkspaceByokEndpointInput.graphql.swift new file mode 100644 index 0000000000..9360207534 --- /dev/null +++ b/packages/frontend/apps/ios/App/Packages/AffineGraphQL/Sources/Schema/InputObjects/WorkspaceByokEndpointInput.graphql.swift @@ -0,0 +1,32 @@ +// @generated +// This file was automatically generated and should not be edited. + +import ApolloAPI + +public struct WorkspaceByokEndpointInput: InputObject { + public private(set) var __data: InputDict + + public init(_ data: InputDict) { + __data = data + } + + public init( + kind: String, + url: GraphQLNullable = nil + ) { + __data = InputDict([ + "kind": kind, + "url": url + ]) + } + + public var kind: String { + get { __data["kind"] } + set { __data["kind"] = newValue } + } + + public var url: GraphQLNullable { + get { __data["url"] } + set { __data["url"] = newValue } + } +} diff --git a/packages/frontend/apps/ios/App/Packages/AffineGraphQL/Sources/Schema/InputObjects/WorkspaceByokModelDeclarationInput.graphql.swift b/packages/frontend/apps/ios/App/Packages/AffineGraphQL/Sources/Schema/InputObjects/WorkspaceByokModelDeclarationInput.graphql.swift new file mode 100644 index 0000000000..61901e24d5 --- /dev/null +++ b/packages/frontend/apps/ios/App/Packages/AffineGraphQL/Sources/Schema/InputObjects/WorkspaceByokModelDeclarationInput.graphql.swift @@ -0,0 +1,32 @@ +// @generated +// This file was automatically generated and should not be edited. + +import ApolloAPI + +public struct WorkspaceByokModelDeclarationInput: InputObject { + public private(set) var __data: InputDict + + public init(_ data: InputDict) { + __data = data + } + + public init( + capabilities: [WorkspaceByokCapabilityInput], + modelId: String + ) { + __data = InputDict([ + "capabilities": capabilities, + "modelId": modelId + ]) + } + + public var capabilities: [WorkspaceByokCapabilityInput] { + get { __data["capabilities"] } + set { __data["capabilities"] = newValue } + } + + public var modelId: String { + get { __data["modelId"] } + set { __data["modelId"] = newValue } + } +} diff --git a/packages/frontend/apps/ios/App/Packages/AffineGraphQL/Sources/Schema/InputObjects/WorkspaceByokProfileDefinitionInput.graphql.swift b/packages/frontend/apps/ios/App/Packages/AffineGraphQL/Sources/Schema/InputObjects/WorkspaceByokProfileDefinitionInput.graphql.swift new file mode 100644 index 0000000000..1a515f93a4 --- /dev/null +++ b/packages/frontend/apps/ios/App/Packages/AffineGraphQL/Sources/Schema/InputObjects/WorkspaceByokProfileDefinitionInput.graphql.swift @@ -0,0 +1,39 @@ +// @generated +// This file was automatically generated and should not be edited. + +import ApolloAPI + +public struct WorkspaceByokProfileDefinitionInput: InputObject { + public private(set) var __data: InputDict + + public init(_ data: InputDict) { + __data = data + } + + public init( + endpoint: WorkspaceByokEndpointInput, + models: [WorkspaceByokModelDeclarationInput], + version: SafeInt + ) { + __data = InputDict([ + "endpoint": endpoint, + "models": models, + "version": version + ]) + } + + public var endpoint: WorkspaceByokEndpointInput { + get { __data["endpoint"] } + set { __data["endpoint"] = newValue } + } + + public var models: [WorkspaceByokModelDeclarationInput] { + get { __data["models"] } + set { __data["models"] = newValue } + } + + public var version: SafeInt { + get { __data["version"] } + set { __data["version"] = newValue } + } +} diff --git a/packages/frontend/apps/ios/App/Packages/AffineGraphQL/Sources/Schema/Objects/AdminMailDeliveryAnalytics.graphql.swift b/packages/frontend/apps/ios/App/Packages/AffineGraphQL/Sources/Schema/Objects/AdminMailDeliveryAnalytics.graphql.swift new file mode 100644 index 0000000000..c51b63b705 --- /dev/null +++ b/packages/frontend/apps/ios/App/Packages/AffineGraphQL/Sources/Schema/Objects/AdminMailDeliveryAnalytics.graphql.swift @@ -0,0 +1,12 @@ +// @generated +// This file was automatically generated and should not be edited. + +import ApolloAPI + +public extension Objects { + static let AdminMailDeliveryAnalytics = ApolloAPI.Object( + typename: "AdminMailDeliveryAnalytics", + implementedInterfaces: [], + keyFields: nil + ) +} \ No newline at end of file diff --git a/packages/frontend/apps/ios/App/Packages/AffineGraphQL/Sources/Schema/Objects/AdminMailDeliveryPoint.graphql.swift b/packages/frontend/apps/ios/App/Packages/AffineGraphQL/Sources/Schema/Objects/AdminMailDeliveryPoint.graphql.swift new file mode 100644 index 0000000000..fb141b6678 --- /dev/null +++ b/packages/frontend/apps/ios/App/Packages/AffineGraphQL/Sources/Schema/Objects/AdminMailDeliveryPoint.graphql.swift @@ -0,0 +1,12 @@ +// @generated +// This file was automatically generated and should not be edited. + +import ApolloAPI + +public extension Objects { + static let AdminMailDeliveryPoint = ApolloAPI.Object( + typename: "AdminMailDeliveryPoint", + implementedInterfaces: [], + keyFields: nil + ) +} \ No newline at end of file diff --git a/packages/frontend/apps/ios/App/Packages/AffineGraphQL/Sources/Schema/Objects/AdminMailDeliverySeries.graphql.swift b/packages/frontend/apps/ios/App/Packages/AffineGraphQL/Sources/Schema/Objects/AdminMailDeliverySeries.graphql.swift new file mode 100644 index 0000000000..bb62e6d929 --- /dev/null +++ b/packages/frontend/apps/ios/App/Packages/AffineGraphQL/Sources/Schema/Objects/AdminMailDeliverySeries.graphql.swift @@ -0,0 +1,12 @@ +// @generated +// This file was automatically generated and should not be edited. + +import ApolloAPI + +public extension Objects { + static let AdminMailDeliverySeries = ApolloAPI.Object( + typename: "AdminMailDeliverySeries", + implementedInterfaces: [], + keyFields: nil + ) +} \ No newline at end of file diff --git a/packages/frontend/apps/ios/App/Packages/AffineGraphQL/Sources/Schema/Objects/AdminMailDeliverySummary.graphql.swift b/packages/frontend/apps/ios/App/Packages/AffineGraphQL/Sources/Schema/Objects/AdminMailDeliverySummary.graphql.swift new file mode 100644 index 0000000000..b9fc7210fe --- /dev/null +++ b/packages/frontend/apps/ios/App/Packages/AffineGraphQL/Sources/Schema/Objects/AdminMailDeliverySummary.graphql.swift @@ -0,0 +1,12 @@ +// @generated +// This file was automatically generated and should not be edited. + +import ApolloAPI + +public extension Objects { + static let AdminMailDeliverySummary = ApolloAPI.Object( + typename: "AdminMailDeliverySummary", + implementedInterfaces: [], + keyFields: nil + ) +} \ No newline at end of file diff --git a/packages/frontend/apps/ios/App/Packages/AffineGraphQL/Sources/Schema/Objects/CopilotModelType.graphql.swift b/packages/frontend/apps/ios/App/Packages/AffineGraphQL/Sources/Schema/Objects/AuthSigningKeyType.graphql.swift similarity index 67% rename from packages/frontend/apps/ios/App/Packages/AffineGraphQL/Sources/Schema/Objects/CopilotModelType.graphql.swift rename to packages/frontend/apps/ios/App/Packages/AffineGraphQL/Sources/Schema/Objects/AuthSigningKeyType.graphql.swift index 1aa617b3b4..011caef8d9 100644 --- a/packages/frontend/apps/ios/App/Packages/AffineGraphQL/Sources/Schema/Objects/CopilotModelType.graphql.swift +++ b/packages/frontend/apps/ios/App/Packages/AffineGraphQL/Sources/Schema/Objects/AuthSigningKeyType.graphql.swift @@ -4,8 +4,8 @@ import ApolloAPI public extension Objects { - static let CopilotModelType = ApolloAPI.Object( - typename: "CopilotModelType", + static let AuthSigningKeyType = ApolloAPI.Object( + typename: "AuthSigningKeyType", implementedInterfaces: [], keyFields: nil ) diff --git a/packages/frontend/apps/ios/App/Packages/AffineGraphQL/Sources/Schema/Objects/RevealedAccessToken.graphql.swift b/packages/frontend/apps/ios/App/Packages/AffineGraphQL/Sources/Schema/Objects/CopilotRouteOptions.graphql.swift similarity index 66% rename from packages/frontend/apps/ios/App/Packages/AffineGraphQL/Sources/Schema/Objects/RevealedAccessToken.graphql.swift rename to packages/frontend/apps/ios/App/Packages/AffineGraphQL/Sources/Schema/Objects/CopilotRouteOptions.graphql.swift index f20d72f93b..726f08e566 100644 --- a/packages/frontend/apps/ios/App/Packages/AffineGraphQL/Sources/Schema/Objects/RevealedAccessToken.graphql.swift +++ b/packages/frontend/apps/ios/App/Packages/AffineGraphQL/Sources/Schema/Objects/CopilotRouteOptions.graphql.swift @@ -4,8 +4,8 @@ import ApolloAPI public extension Objects { - static let RevealedAccessToken = ApolloAPI.Object( - typename: "RevealedAccessToken", + static let CopilotRouteOptions = ApolloAPI.Object( + typename: "CopilotRouteOptions", implementedInterfaces: [], keyFields: nil ) diff --git a/packages/frontend/apps/ios/App/Packages/AffineGraphQL/Sources/Schema/Objects/TokenType.graphql.swift b/packages/frontend/apps/ios/App/Packages/AffineGraphQL/Sources/Schema/Objects/CopilotRouteTarget.graphql.swift similarity index 67% rename from packages/frontend/apps/ios/App/Packages/AffineGraphQL/Sources/Schema/Objects/TokenType.graphql.swift rename to packages/frontend/apps/ios/App/Packages/AffineGraphQL/Sources/Schema/Objects/CopilotRouteTarget.graphql.swift index 5f56cc4757..a483838518 100644 --- a/packages/frontend/apps/ios/App/Packages/AffineGraphQL/Sources/Schema/Objects/TokenType.graphql.swift +++ b/packages/frontend/apps/ios/App/Packages/AffineGraphQL/Sources/Schema/Objects/CopilotRouteTarget.graphql.swift @@ -4,8 +4,8 @@ import ApolloAPI public extension Objects { - static let TokenType = ApolloAPI.Object( - typename: "tokenType", + static let CopilotRouteTarget = ApolloAPI.Object( + typename: "CopilotRouteTarget", implementedInterfaces: [], keyFields: nil ) diff --git a/packages/frontend/apps/ios/App/Packages/AffineGraphQL/Sources/Schema/Objects/CopilotModelsType.graphql.swift b/packages/frontend/apps/ios/App/Packages/AffineGraphQL/Sources/Schema/Objects/McpCredentialType.graphql.swift similarity index 67% rename from packages/frontend/apps/ios/App/Packages/AffineGraphQL/Sources/Schema/Objects/CopilotModelsType.graphql.swift rename to packages/frontend/apps/ios/App/Packages/AffineGraphQL/Sources/Schema/Objects/McpCredentialType.graphql.swift index a74b943e57..89e30c0bee 100644 --- a/packages/frontend/apps/ios/App/Packages/AffineGraphQL/Sources/Schema/Objects/CopilotModelsType.graphql.swift +++ b/packages/frontend/apps/ios/App/Packages/AffineGraphQL/Sources/Schema/Objects/McpCredentialType.graphql.swift @@ -4,8 +4,8 @@ import ApolloAPI public extension Objects { - static let CopilotModelsType = ApolloAPI.Object( - typename: "CopilotModelsType", + static let McpCredentialType = ApolloAPI.Object( + typename: "McpCredentialType", implementedInterfaces: [], keyFields: nil ) diff --git a/packages/frontend/apps/ios/App/Packages/AffineGraphQL/Sources/Schema/Objects/RevealedMcpCredentialType.graphql.swift b/packages/frontend/apps/ios/App/Packages/AffineGraphQL/Sources/Schema/Objects/RevealedMcpCredentialType.graphql.swift new file mode 100644 index 0000000000..843129e5fb --- /dev/null +++ b/packages/frontend/apps/ios/App/Packages/AffineGraphQL/Sources/Schema/Objects/RevealedMcpCredentialType.graphql.swift @@ -0,0 +1,12 @@ +// @generated +// This file was automatically generated and should not be edited. + +import ApolloAPI + +public extension Objects { + static let RevealedMcpCredentialType = ApolloAPI.Object( + typename: "RevealedMcpCredentialType", + implementedInterfaces: [], + keyFields: nil + ) +} \ No newline at end of file diff --git a/packages/frontend/apps/ios/App/Packages/AffineGraphQL/Sources/Schema/Objects/WorkspaceByokCapabilityType.graphql.swift b/packages/frontend/apps/ios/App/Packages/AffineGraphQL/Sources/Schema/Objects/WorkspaceByokCapabilityType.graphql.swift new file mode 100644 index 0000000000..f02436402f --- /dev/null +++ b/packages/frontend/apps/ios/App/Packages/AffineGraphQL/Sources/Schema/Objects/WorkspaceByokCapabilityType.graphql.swift @@ -0,0 +1,12 @@ +// @generated +// This file was automatically generated and should not be edited. + +import ApolloAPI + +public extension Objects { + static let WorkspaceByokCapabilityType = ApolloAPI.Object( + typename: "WorkspaceByokCapabilityType", + implementedInterfaces: [], + keyFields: nil + ) +} \ No newline at end of file diff --git a/packages/frontend/apps/ios/App/Packages/AffineGraphQL/Sources/Schema/Objects/WorkspaceByokEndpointType.graphql.swift b/packages/frontend/apps/ios/App/Packages/AffineGraphQL/Sources/Schema/Objects/WorkspaceByokEndpointType.graphql.swift new file mode 100644 index 0000000000..d5e168ddcb --- /dev/null +++ b/packages/frontend/apps/ios/App/Packages/AffineGraphQL/Sources/Schema/Objects/WorkspaceByokEndpointType.graphql.swift @@ -0,0 +1,12 @@ +// @generated +// This file was automatically generated and should not be edited. + +import ApolloAPI + +public extension Objects { + static let WorkspaceByokEndpointType = ApolloAPI.Object( + typename: "WorkspaceByokEndpointType", + implementedInterfaces: [], + keyFields: nil + ) +} \ No newline at end of file diff --git a/packages/frontend/apps/ios/App/Packages/AffineGraphQL/Sources/Schema/Objects/WorkspaceByokKeyConfigType.graphql.swift b/packages/frontend/apps/ios/App/Packages/AffineGraphQL/Sources/Schema/Objects/WorkspaceByokKeyConfigType.graphql.swift deleted file mode 100644 index 2086215fea..0000000000 --- a/packages/frontend/apps/ios/App/Packages/AffineGraphQL/Sources/Schema/Objects/WorkspaceByokKeyConfigType.graphql.swift +++ /dev/null @@ -1,12 +0,0 @@ -// @generated -// This file was automatically generated and should not be edited. - -import ApolloAPI - -public extension Objects { - static let WorkspaceByokKeyConfigType = ApolloAPI.Object( - typename: "WorkspaceByokKeyConfigType", - implementedInterfaces: [], - keyFields: nil - ) -} \ No newline at end of file diff --git a/packages/frontend/apps/ios/App/Packages/AffineGraphQL/Sources/Schema/Objects/TestWorkspaceByokConfigResultType.graphql.swift b/packages/frontend/apps/ios/App/Packages/AffineGraphQL/Sources/Schema/Objects/WorkspaceByokModelDeclarationType.graphql.swift similarity index 61% rename from packages/frontend/apps/ios/App/Packages/AffineGraphQL/Sources/Schema/Objects/TestWorkspaceByokConfigResultType.graphql.swift rename to packages/frontend/apps/ios/App/Packages/AffineGraphQL/Sources/Schema/Objects/WorkspaceByokModelDeclarationType.graphql.swift index 283d8a5389..bc94cba3e6 100644 --- a/packages/frontend/apps/ios/App/Packages/AffineGraphQL/Sources/Schema/Objects/TestWorkspaceByokConfigResultType.graphql.swift +++ b/packages/frontend/apps/ios/App/Packages/AffineGraphQL/Sources/Schema/Objects/WorkspaceByokModelDeclarationType.graphql.swift @@ -4,8 +4,8 @@ import ApolloAPI public extension Objects { - static let TestWorkspaceByokConfigResultType = ApolloAPI.Object( - typename: "TestWorkspaceByokConfigResultType", + static let WorkspaceByokModelDeclarationType = ApolloAPI.Object( + typename: "WorkspaceByokModelDeclarationType", implementedInterfaces: [], keyFields: nil ) diff --git a/packages/frontend/apps/ios/App/Packages/AffineGraphQL/Sources/Schema/Objects/WorkspaceByokProbeType.graphql.swift b/packages/frontend/apps/ios/App/Packages/AffineGraphQL/Sources/Schema/Objects/WorkspaceByokProbeType.graphql.swift new file mode 100644 index 0000000000..f5c87e406f --- /dev/null +++ b/packages/frontend/apps/ios/App/Packages/AffineGraphQL/Sources/Schema/Objects/WorkspaceByokProbeType.graphql.swift @@ -0,0 +1,12 @@ +// @generated +// This file was automatically generated and should not be edited. + +import ApolloAPI + +public extension Objects { + static let WorkspaceByokProbeType = ApolloAPI.Object( + typename: "WorkspaceByokProbeType", + implementedInterfaces: [], + keyFields: nil + ) +} \ No newline at end of file diff --git a/packages/frontend/apps/ios/App/Packages/AffineGraphQL/Sources/Schema/Objects/WorkspaceByokCapabilityWarningType.graphql.swift b/packages/frontend/apps/ios/App/Packages/AffineGraphQL/Sources/Schema/Objects/WorkspaceByokProfileDefinitionType.graphql.swift similarity index 61% rename from packages/frontend/apps/ios/App/Packages/AffineGraphQL/Sources/Schema/Objects/WorkspaceByokCapabilityWarningType.graphql.swift rename to packages/frontend/apps/ios/App/Packages/AffineGraphQL/Sources/Schema/Objects/WorkspaceByokProfileDefinitionType.graphql.swift index 7000188c32..fd6c1f5570 100644 --- a/packages/frontend/apps/ios/App/Packages/AffineGraphQL/Sources/Schema/Objects/WorkspaceByokCapabilityWarningType.graphql.swift +++ b/packages/frontend/apps/ios/App/Packages/AffineGraphQL/Sources/Schema/Objects/WorkspaceByokProfileDefinitionType.graphql.swift @@ -4,8 +4,8 @@ import ApolloAPI public extension Objects { - static let WorkspaceByokCapabilityWarningType = ApolloAPI.Object( - typename: "WorkspaceByokCapabilityWarningType", + static let WorkspaceByokProfileDefinitionType = ApolloAPI.Object( + typename: "WorkspaceByokProfileDefinitionType", implementedInterfaces: [], keyFields: nil ) diff --git a/packages/frontend/apps/ios/App/Packages/AffineGraphQL/Sources/Schema/Objects/WorkspaceByokProfileType.graphql.swift b/packages/frontend/apps/ios/App/Packages/AffineGraphQL/Sources/Schema/Objects/WorkspaceByokProfileType.graphql.swift new file mode 100644 index 0000000000..5667d4d562 --- /dev/null +++ b/packages/frontend/apps/ios/App/Packages/AffineGraphQL/Sources/Schema/Objects/WorkspaceByokProfileType.graphql.swift @@ -0,0 +1,12 @@ +// @generated +// This file was automatically generated and should not be edited. + +import ApolloAPI + +public extension Objects { + static let WorkspaceByokProfileType = ApolloAPI.Object( + typename: "WorkspaceByokProfileType", + implementedInterfaces: [], + keyFields: nil + ) +} \ No newline at end of file diff --git a/packages/frontend/apps/ios/App/Packages/AffineGraphQL/Sources/Schema/SchemaMetadata.graphql.swift b/packages/frontend/apps/ios/App/Packages/AffineGraphQL/Sources/Schema/SchemaMetadata.graphql.swift index 2e71aca130..a034c7c7bc 100644 --- a/packages/frontend/apps/ios/App/Packages/AffineGraphQL/Sources/Schema/SchemaMetadata.graphql.swift +++ b/packages/frontend/apps/ios/App/Packages/AffineGraphQL/Sources/Schema/SchemaMetadata.graphql.swift @@ -25,6 +25,10 @@ public enum SchemaMetadata: ApolloAPI.SchemaMetadata { "AdminDashboardMinutePoint": AffineGraphQL.Objects.AdminDashboardMinutePoint, "AdminDashboardValueDayPoint": AffineGraphQL.Objects.AdminDashboardValueDayPoint, "AdminLicensePreview": AffineGraphQL.Objects.AdminLicensePreview, + "AdminMailDeliveryAnalytics": AffineGraphQL.Objects.AdminMailDeliveryAnalytics, + "AdminMailDeliveryPoint": AffineGraphQL.Objects.AdminMailDeliveryPoint, + "AdminMailDeliverySeries": AffineGraphQL.Objects.AdminMailDeliverySeries, + "AdminMailDeliverySummary": AffineGraphQL.Objects.AdminMailDeliverySummary, "AdminSharedLinkTopItem": AffineGraphQL.Objects.AdminSharedLinkTopItem, "AdminWorkspace": AffineGraphQL.Objects.AdminWorkspace, "AdminWorkspaceMember": AffineGraphQL.Objects.AdminWorkspaceMember, @@ -34,6 +38,7 @@ public enum SchemaMetadata: ApolloAPI.SchemaMetadata { "AggregateResultObjectType": AffineGraphQL.Objects.AggregateResultObjectType, "AppConfigValidateResult": AffineGraphQL.Objects.AppConfigValidateResult, "AudioSliceManifestItemType": AffineGraphQL.Objects.AudioSliceManifestItemType, + "AuthSigningKeyType": AffineGraphQL.Objects.AuthSigningKeyType, "BlobUploadInit": AffineGraphQL.Objects.BlobUploadInit, "BlobUploadPart": AffineGraphQL.Objects.BlobUploadPart, "BlobUploadedPart": AffineGraphQL.Objects.BlobUploadedPart, @@ -56,9 +61,9 @@ public enum SchemaMetadata: ApolloAPI.SchemaMetadata { "CopilotContextFile": AffineGraphQL.Objects.CopilotContextFile, "CopilotHistories": AffineGraphQL.Objects.CopilotHistories, "CopilotHistoriesTypeEdge": AffineGraphQL.Objects.CopilotHistoriesTypeEdge, - "CopilotModelType": AffineGraphQL.Objects.CopilotModelType, - "CopilotModelsType": AffineGraphQL.Objects.CopilotModelsType, "CopilotQuota": AffineGraphQL.Objects.CopilotQuota, + "CopilotRouteOptions": AffineGraphQL.Objects.CopilotRouteOptions, + "CopilotRouteTarget": AffineGraphQL.Objects.CopilotRouteTarget, "CopilotWorkspaceConfig": AffineGraphQL.Objects.CopilotWorkspaceConfig, "CopilotWorkspaceFile": AffineGraphQL.Objects.CopilotWorkspaceFile, "CopilotWorkspaceFileTypeEdge": AffineGraphQL.Objects.CopilotWorkspaceFileTypeEdge, @@ -85,6 +90,7 @@ public enum SchemaMetadata: ApolloAPI.SchemaMetadata { "License": AffineGraphQL.Objects.License, "LimitedUserType": AffineGraphQL.Objects.LimitedUserType, "ListedBlob": AffineGraphQL.Objects.ListedBlob, + "McpCredentialType": AffineGraphQL.Objects.McpCredentialType, "MeetingActionItemType": AffineGraphQL.Objects.MeetingActionItemType, "MeetingSummaryV2Type": AffineGraphQL.Objects.MeetingSummaryV2Type, "Mutation": AffineGraphQL.Objects.Mutation, @@ -107,7 +113,7 @@ public enum SchemaMetadata: ApolloAPI.SchemaMetadata { "ReleaseVersionType": AffineGraphQL.Objects.ReleaseVersionType, "RemoveAvatar": AffineGraphQL.Objects.RemoveAvatar, "ReplyObjectType": AffineGraphQL.Objects.ReplyObjectType, - "RevealedAccessToken": AffineGraphQL.Objects.RevealedAccessToken, + "RevealedMcpCredentialType": AffineGraphQL.Objects.RevealedMcpCredentialType, "SearchDocObjectType": AffineGraphQL.Objects.SearchDocObjectType, "SearchNodeObjectType": AffineGraphQL.Objects.SearchNodeObjectType, "SearchResultObjectType": AffineGraphQL.Objects.SearchResultObjectType, @@ -116,7 +122,6 @@ public enum SchemaMetadata: ApolloAPI.SchemaMetadata { "StreamObject": AffineGraphQL.Objects.StreamObject, "SubscriptionPrice": AffineGraphQL.Objects.SubscriptionPrice, "SubscriptionType": AffineGraphQL.Objects.SubscriptionType, - "TestWorkspaceByokConfigResultType": AffineGraphQL.Objects.TestWorkspaceByokConfigResultType, "TimeWindow": AffineGraphQL.Objects.TimeWindow, "TranscriptionItemType": AffineGraphQL.Objects.TranscriptionItemType, "TranscriptionQualityType": AffineGraphQL.Objects.TranscriptionQualityType, @@ -128,8 +133,12 @@ public enum SchemaMetadata: ApolloAPI.SchemaMetadata { "UserQuotaUsageType": AffineGraphQL.Objects.UserQuotaUsageType, "UserSettingsType": AffineGraphQL.Objects.UserSettingsType, "UserType": AffineGraphQL.Objects.UserType, - "WorkspaceByokCapabilityWarningType": AffineGraphQL.Objects.WorkspaceByokCapabilityWarningType, - "WorkspaceByokKeyConfigType": AffineGraphQL.Objects.WorkspaceByokKeyConfigType, + "WorkspaceByokCapabilityType": AffineGraphQL.Objects.WorkspaceByokCapabilityType, + "WorkspaceByokEndpointType": AffineGraphQL.Objects.WorkspaceByokEndpointType, + "WorkspaceByokModelDeclarationType": AffineGraphQL.Objects.WorkspaceByokModelDeclarationType, + "WorkspaceByokProbeType": AffineGraphQL.Objects.WorkspaceByokProbeType, + "WorkspaceByokProfileDefinitionType": AffineGraphQL.Objects.WorkspaceByokProfileDefinitionType, + "WorkspaceByokProfileType": AffineGraphQL.Objects.WorkspaceByokProfileType, "WorkspaceByokSettingsType": AffineGraphQL.Objects.WorkspaceByokSettingsType, "WorkspaceByokUsagePointType": AffineGraphQL.Objects.WorkspaceByokUsagePointType, "WorkspaceCalendarItemObjectType": AffineGraphQL.Objects.WorkspaceCalendarItemObjectType, @@ -140,8 +149,7 @@ public enum SchemaMetadata: ApolloAPI.SchemaMetadata { "WorkspaceQuotaType": AffineGraphQL.Objects.WorkspaceQuotaType, "WorkspaceRolePermissions": AffineGraphQL.Objects.WorkspaceRolePermissions, "WorkspaceType": AffineGraphQL.Objects.WorkspaceType, - "WorkspaceUserType": AffineGraphQL.Objects.WorkspaceUserType, - "tokenType": AffineGraphQL.Objects.TokenType + "WorkspaceUserType": AffineGraphQL.Objects.WorkspaceUserType ] public static func objectType(forTypename typename: String) -> ApolloAPI.Object? { diff --git a/packages/frontend/apps/ios/App/Packages/Intelligents/Package.swift b/packages/frontend/apps/ios/App/Packages/Intelligents/Package.swift index 72a171829b..9cbf51d8bc 100644 --- a/packages/frontend/apps/ios/App/Packages/Intelligents/Package.swift +++ b/packages/frontend/apps/ios/App/Packages/Intelligents/Package.swift @@ -16,7 +16,7 @@ let package = Package( dependencies: [ .package(path: "../AffineGraphQL"), .package(path: "../AffineResources"), - .package(url: "https://github.com/apollographql/apollo-ios.git", from: "1.25.7"), + .package(url: "https://github.com/apollographql/apollo-ios.git", from: "1.25.4"), .package(url: "https://github.com/apple/swift-collections.git", from: "1.6.0"), .package(url: "https://github.com/SnapKit/SnapKit.git", from: "5.7.1"), .package(url: "https://github.com/SwifterSwift/SwifterSwift.git", from: "6.2.0"), diff --git a/tests/affine-cloud-copilot/e2e/utils/test-utils.ts b/tests/affine-cloud-copilot/e2e/utils/test-utils.ts index d371713607..b5b0c9fb6c 100644 --- a/tests/affine-cloud-copilot/e2e/utils/test-utils.ts +++ b/tests/affine-cloud-copilot/e2e/utils/test-utils.ts @@ -1,8 +1,5 @@ import { skipOnboarding } from '@affine-test/kit/playwright'; -import { - createRandomAIUser, - switchDefaultChatModel, -} from '@affine-test/kit/utils/cloud'; +import { createRandomAIUser } from '@affine-test/kit/utils/cloud'; import { openHomePage, setCoreUrl } from '@affine-test/kit/utils/load-page'; import { clickNewPageButton, @@ -63,7 +60,6 @@ export class TestUtils { public async setupTestEnvironment(page: Page, defaultModel?: string) { const selectedModel = defaultModel ?? 'gpt-5.6-luna'; - await switchDefaultChatModel(selectedModel); await skipOnboarding(page.context()); await page.context().addInitScript(model => { diff --git a/tests/kit/src/utils/cloud.ts b/tests/kit/src/utils/cloud.ts index a74b696b0d..951866b6ee 100644 --- a/tests/kit/src/utils/cloud.ts +++ b/tests/kit/src/utils/cloud.ts @@ -155,21 +155,6 @@ export async function cleanupWorkspace(workspaceId: string): Promise { }); } -export async function switchDefaultChatModel(model: string) { - await runPrisma(async client => { - const prompt = await client.aiPrompt.findFirst({ - where: { name: 'Chat With AFFiNE AI' }, - select: { id: true }, - }); - if (!prompt) return; - - await client.aiPrompt.update({ - where: { id: prompt.id }, - data: { model }, - }); - }); -} - export async function createRandomAIUser(): Promise<{ name: string; email: string;