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

feat(server): improve context management (#15448)

Browse Source 
#### PR Dependency Tree


* **PR #15448** 👈

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

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

* **New Features**
* Added workspace artifact upload, browsing, removal, deduplication, and
library ownership support.
* Copilot now supports scoped document and artifact search, canvas
reading, live editor context, and frontend tools.
* Added scope and focus selectors with source-resolution receipts in
chat.
* Added embedding health, progress, synchronization, and retrieval
capabilities.
* Added BYOK policy visibility, provider restrictions, endpoint dialect
selection, and validation.
* Added delegated editor interactions and userdata document
authorization.

* **Bug Fixes**
* Improved attachment handling, cancellation, access control, retrieval
fallbacks, workspace synchronization, and configuration validation.
<!-- end of auto-generated comment: release notes by coderabbit.ai -->
This commit is contained in:
DarkSky
2026-08-10 09:27:58 +08:00
committed by GitHub
parent 42322d13fe
commit ee899a267b
311 changed files with 20468 additions and 14806 deletions
Download Patch File Download Diff File Expand all files Collapse all files
packages/backend/server/migrations/20260803095500_converge_copilot_runtime/migration.sql
+50 -36
View File
@@ -10,49 +10,63 @@ WHERE "id" IN (
'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 "definition" JSONB NOT NULL DEFAULT '{}',
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';
DROP CONSTRAINT "ai_sessions_metadata_prompt_name_fkey";
ALTER TABLE "ai_transcript_tasks"
DROP COLUMN "strategy";
ALTER COLUMN "strategy" SET DEFAULT '';
DROP TABLE "ai_prompts_messages";
DROP TABLE "ai_prompts_metadata";
ALTER TABLE "ai_sessions_messages" ADD COLUMN "scope_snapshot" JSONB;
ALTER TABLE "ai_sessions_metadata" ADD COLUMN "focus" JSONB;
ALTER TYPE "AiPromptRole" RENAME TO "AiSessionMessageRole";
CREATE TABLE "workspace_artifacts" (
"id" UUID NOT NULL,
"workspace_id" VARCHAR NOT NULL,
"content_hash" VARCHAR NOT NULL,
"display_name" VARCHAR,
"file_name" VARCHAR,
"canonical_media_type" VARCHAR NOT NULL,
"size_bytes" BIGINT NOT NULL,
"storage_scope" VARCHAR NOT NULL,
"storage_key" TEXT NOT NULL,
"status" VARCHAR NOT NULL,
"library_owned" BOOLEAN NOT NULL DEFAULT false,
"reservation_expires_at" TIMESTAMPTZ(3),
"ready_at" TIMESTAMPTZ(3),
"created_at" TIMESTAMPTZ(3) NOT NULL DEFAULT CURRENT_TIMESTAMP,
"updated_at" TIMESTAMPTZ(3) NOT NULL DEFAULT CURRENT_TIMESTAMP,
CONSTRAINT "workspace_artifacts_pkey" PRIMARY KEY ("id"),
CONSTRAINT "workspace_artifacts_library_display_name_check"
CHECK (NOT "library_owned" OR NULLIF(BTRIM("display_name"), '') IS NOT NULL),
CONSTRAINT "workspace_artifacts_workspace_id_fkey" FOREIGN KEY ("workspace_id") REFERENCES "workspaces"("id") ON DELETE CASCADE
);
CREATE UNIQUE INDEX "workspace_artifacts_workspace_id_content_hash_key" ON "workspace_artifacts"("workspace_id", "content_hash");
CREATE UNIQUE INDEX "workspace_artifacts_workspace_id_id_key" ON "workspace_artifacts"("workspace_id", "id");
CREATE INDEX "workspace_artifacts_workspace_id_status_idx" ON "workspace_artifacts"("workspace_id", "status");
CREATE INDEX "workspace_artifacts_status_reservation_expires_at_idx" ON "workspace_artifacts"("status", "reservation_expires_at");
CREATE TABLE "ai_message_artifacts" (
"message_id" VARCHAR NOT NULL,
"workspace_id" VARCHAR NOT NULL,
"artifact_id" UUID NOT NULL,
"role" VARCHAR NOT NULL,
"display_name" VARCHAR,
"metadata" JSONB,
"created_at" TIMESTAMPTZ(3) NOT NULL DEFAULT CURRENT_TIMESTAMP,
CONSTRAINT "ai_message_artifacts_pkey" PRIMARY KEY ("message_id", "artifact_id", "role"),
CONSTRAINT "ai_message_artifacts_message_id_fkey" FOREIGN KEY ("message_id") REFERENCES "ai_sessions_messages"("id") ON DELETE CASCADE,
CONSTRAINT "ai_message_artifacts_workspace_id_artifact_id_fkey" FOREIGN KEY ("workspace_id", "artifact_id") REFERENCES "workspace_artifacts"("workspace_id", "id") ON DELETE CASCADE
);
CREATE INDEX "ai_message_artifacts_workspace_id_artifact_id_idx" ON "ai_message_artifacts"("workspace_id", "artifact_id");
-- After stable and beta no longer run binaries built with the 115-migration
-- schema, remove the old provider keys, obsolete local-lease rows, ai_contexts,
-- ai_context_embeddings, and ai_workspace_embeddings in one cleanup migration.
packages/backend/server/schema.prisma
+74 -131
View File
@@ -191,7 +191,6 @@ model Workspace {
docs WorkspaceDoc[]
blobs Blob[]
ignoredDocs AiWorkspaceIgnoredDocs[]
embedFiles AiWorkspaceFiles[]
byokConfigs AiWorkspaceByokConfig[]
aiUsageEvents AiUsageEvent[]
comments Comment[]
@@ -209,6 +208,7 @@ model Workspace {
docAccessPolicies DocAccessPolicy[]
docGrants DocGrant[]
mcpCredentials McpCredential[]
artifacts WorkspaceArtifact[]
@@index([lastCheckEmbeddings])
@@index([createdAt])
@@ -652,8 +652,6 @@ model Snapshot {
// we need to clear all hanging updates and snapshots before enable the foreign key on workspaceId
// workspace Workspace @relation(fields: [workspaceId], references: [id], onDelete: Cascade)
embedding AiWorkspaceEmbedding[]
@@id([workspaceId, id])
@@index([workspaceId, updatedAt])
@@map("snapshots")
@@ -710,6 +708,11 @@ enum AiSessionMessageRole {
system
assistant
user
// the database type keeps the legacy name so the previous release, whose
// Prisma client casts enum values as "AiPromptRole", can keep writing
// ai_sessions_messages while it runs against the same database
@@map("AiPromptRole")
}
model AiSessionMessage {
@@ -721,10 +724,12 @@ model AiSessionMessage {
streamObjects Json? @db.Json
attachments Json? @db.Json
params Json? @db.Json
scopeSnapshot Json? @map("scope_snapshot") @db.JsonB
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)
session AiSession @relation(fields: [sessionId], references: [id], onDelete: Cascade)
artifacts AiMessageArtifact[]
@@index([sessionId])
@@index([sessionId, compatSubmissionId])
@@ -747,10 +752,10 @@ model AiSession {
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)
focus Json? @db.JsonB
user User @relation(fields: [userId], references: [id], onDelete: Cascade)
messages AiSessionMessage[]
context AiContext[]
actionRuns AiActionRun[]
//NOTE:
@@ -764,6 +769,50 @@ model AiSession {
@@map("ai_sessions_metadata")
}
model WorkspaceArtifact {
id String @id @default(uuid()) @db.Uuid
workspaceId String @map("workspace_id") @db.VarChar
contentHash String @map("content_hash") @db.VarChar
displayName String? @map("display_name") @db.VarChar
fileName String? @map("file_name") @db.VarChar
canonicalMediaType String @map("canonical_media_type") @db.VarChar
sizeBytes BigInt @map("size_bytes")
storageScope String @map("storage_scope") @db.VarChar
storageKey String @map("storage_key") @db.Text
status String @db.VarChar
libraryOwned Boolean @default(false) @map("library_owned")
reservationExpiresAt DateTime? @map("reservation_expires_at") @db.Timestamptz(3)
readyAt DateTime? @map("ready_at") @db.Timestamptz(3)
createdAt DateTime @default(now()) @map("created_at") @db.Timestamptz(3)
updatedAt DateTime @updatedAt @map("updated_at") @db.Timestamptz(3)
workspace Workspace @relation(fields: [workspaceId], references: [id], onDelete: Cascade)
messages AiMessageArtifact[]
@@unique([workspaceId, contentHash])
@@unique([workspaceId, id])
@@index([workspaceId, status])
@@index([status, reservationExpiresAt])
@@map("workspace_artifacts")
}
model AiMessageArtifact {
messageId String @map("message_id") @db.VarChar
workspaceId String @map("workspace_id") @db.VarChar
artifactId String @map("artifact_id") @db.Uuid
role String @db.VarChar
displayName String? @map("display_name") @db.VarChar
metadata Json? @db.JsonB
createdAt DateTime @default(now()) @map("created_at") @db.Timestamptz(3)
message AiSessionMessage @relation(fields: [messageId], references: [id], onDelete: Cascade)
artifact WorkspaceArtifact @relation(fields: [workspaceId, artifactId], references: [workspaceId, id], onDelete: Cascade)
@@id([messageId, artifactId, role])
@@index([workspaceId, artifactId])
@@map("ai_message_artifacts")
}
model AiActionRun {
id String @id @default(uuid()) @db.VarChar
userId String @map("user_id") @db.VarChar
@@ -821,59 +870,6 @@ model AiTranscriptTask {
@@map("ai_transcript_tasks")
}
model AiContext {
id String @id @default(uuid()) @db.VarChar
sessionId String @map("session_id") @db.VarChar
config Json @db.Json
createdAt DateTime @default(now()) @map("created_at") @db.Timestamptz(3)
updatedAt DateTime @updatedAt @map("updated_at") @db.Timestamptz(3)
embeddings AiContextEmbedding[]
session AiSession @relation(fields: [sessionId], references: [id], onDelete: Cascade)
@@map("ai_contexts")
}
model AiContextEmbedding {
id String @id @default(uuid()) @db.VarChar
contextId String @map("context_id") @db.VarChar
fileId String @map("file_id") @db.VarChar
// a file can be divided into multiple chunks and embedded separately.
chunk Int @db.Integer
content String @db.VarChar
embedding Unsupported("vector(1024)")
createdAt DateTime @default(now()) @map("created_at") @db.Timestamptz(3)
updatedAt DateTime @updatedAt @map("updated_at") @db.Timestamptz(3)
context AiContext @relation(fields: [contextId], references: [id], onDelete: Cascade)
@@unique([contextId, fileId, chunk])
@@index([embedding], map: "ai_context_embeddings_idx")
@@map("ai_context_embeddings")
}
model AiWorkspaceEmbedding {
workspaceId String @map("workspace_id") @db.VarChar
docId String @map("doc_id") @db.VarChar
// a doc can be divided into multiple chunks and embedded separately.
chunk Int @db.Integer
content String @db.VarChar
embedding Unsupported("vector(1024)")
createdAt DateTime @default(now()) @map("created_at") @db.Timestamptz(3)
updatedAt DateTime @updatedAt @map("updated_at") @db.Timestamptz(3)
// workspace level search not available for non-cloud workspaces
// so we can match this record with the snapshot one by one
snapshot Snapshot @relation(fields: [workspaceId, docId], references: [workspaceId, id], onDelete: Cascade)
@@id([workspaceId, docId, chunk])
@@index([embedding], map: "ai_workspace_embeddings_idx")
@@map("ai_workspace_embeddings")
}
model AiWorkspaceIgnoredDocs {
workspaceId String @map("workspace_id") @db.VarChar
docId String @map("doc_id") @db.VarChar
@@ -886,78 +882,26 @@ model AiWorkspaceIgnoredDocs {
@@map("ai_workspace_ignored_docs")
}
model AiWorkspaceFiles {
workspaceId String @map("workspace_id") @db.VarChar
fileId String @map("file_id") @db.VarChar
blobId String @default("") @map("blob_id") @db.VarChar
fileName String @map("file_name") @db.VarChar
mimeType String @map("mime_type") @db.VarChar
size Int @db.Integer
createdAt DateTime @default(now()) @map("created_at") @db.Timestamptz(3)
workspace Workspace @relation(fields: [workspaceId], references: [id], onDelete: Cascade)
embeddings AiWorkspaceFileEmbedding[]
@@id([workspaceId, fileId])
@@map("ai_workspace_files")
}
model AiWorkspaceFileEmbedding {
workspaceId String @map("workspace_id") @db.VarChar
fileId String @map("file_id") @db.VarChar
// a file can be divided into multiple chunks and embedded separately.
chunk Int @db.Integer
content String @db.VarChar
embedding Unsupported("vector(1024)")
createdAt DateTime @default(now()) @map("created_at") @db.Timestamptz(3)
file AiWorkspaceFiles @relation(fields: [workspaceId, fileId], references: [workspaceId, fileId], onDelete: Cascade)
@@id([workspaceId, fileId, chunk])
@@index([embedding], map: "ai_workspace_file_embeddings_idx")
@@map("ai_workspace_file_embeddings")
}
model AiWorkspaceBlobEmbedding {
workspaceId String @map("workspace_id") @db.VarChar
blobId String @map("blob_id") @db.VarChar
// a file can be divided into multiple chunks and embedded separately.
chunk Int @db.Integer
content String @db.VarChar
embedding Unsupported("vector(1024)")
createdAt DateTime @default(now()) @map("created_at") @db.Timestamptz(3)
blob Blob @relation(fields: [workspaceId, blobId], references: [workspaceId, key], onDelete: Cascade)
@@id([workspaceId, blobId, chunk])
@@index([embedding], map: "ai_workspace_blob_embeddings_idx")
@@map("ai_workspace_blob_embeddings")
}
model AiWorkspaceByokConfig {
id String @id @default(uuid()) @db.VarChar
workspaceId String @map("workspace_id") @db.VarChar
provider String @db.VarChar
name String @db.VarChar
description String? @db.VarChar
encryptedApiKey String @map("encrypted_api_key") @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)
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
createdBy String? @map("created_by") @db.VarChar
updatedBy String? @map("updated_by") @db.VarChar
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
workspaceId String @map("workspace_id") @db.VarChar
provider String @db.VarChar
name String @db.VarChar
description String? @db.VarChar
encryptedApiKey String @map("encrypted_api_key") @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)
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
createdBy String? @map("created_by") @db.VarChar
updatedBy String? @map("updated_by") @db.VarChar
createdAt DateTime @default(now()) @map("created_at") @db.Timestamptz(3)
updatedAt DateTime @updatedAt @map("updated_at") @db.Timestamptz(3)
workspace Workspace @relation(fields: [workspaceId], references: [id], onDelete: Cascade)
@@ -1209,8 +1153,7 @@ model Blob {
createdAt DateTime @default(now()) @map("created_at") @db.Timestamptz(3)
deletedAt DateTime? @map("deleted_at") @db.Timestamptz(3)
workspace Workspace @relation(fields: [workspaceId], references: [id], onDelete: Cascade)
AiWorkspaceBlobEmbedding AiWorkspaceBlobEmbedding[]
workspace Workspace @relation(fields: [workspaceId], references: [id], onDelete: Cascade)
@@id([workspaceId, key])
@@index([workspaceId, status, deletedAt])
packages/backend/server/scripts/genconfig.ts
+6 -3
View File
@@ -13,9 +13,12 @@ import {
const IGNORED_MODULES = new Set(['db', 'redis', 'graphql']);
function getDescriptors() {
return getAllDescriptors().filter(
({ module }) => !IGNORED_MODULES.has(module)
);
return getAllDescriptors()
.filter(({ module }) => !IGNORED_MODULES.has(module))
.map(({ module, descriptors }) => ({
module,
descriptors: descriptors.filter(({ descriptor }) => !descriptor.internal),
}));
}
interface PropertySchema {
packages/backend/server/scripts/repair-pgvector-embedding-tables.sql
-143
View File
@@ -1,143 +0,0 @@
DO $$
DECLARE
has_hnsw BOOLEAN;
BEGIN
IF NOT EXISTS (SELECT 1 FROM pg_extension WHERE extname = 'vector') THEN
BEGIN
CREATE EXTENSION IF NOT EXISTS "vector";
EXCEPTION
WHEN OTHERS THEN
RAISE NOTICE 'pgvector extension is not available. Skip repairing copilot embedding tables.';
RETURN;
END;
END IF;
SELECT EXISTS (SELECT 1 FROM pg_am WHERE amname = 'hnsw') INTO has_hnsw;
IF NOT has_hnsw THEN
RAISE NOTICE 'pgvector HNSW index access method is not available. Skip repairing copilot embedding indexes.';
END IF;
IF to_regclass('public.ai_contexts') IS NOT NULL THEN
CREATE TABLE IF NOT EXISTS "ai_context_embeddings" (
"id" VARCHAR NOT NULL,
"context_id" VARCHAR NOT NULL,
"file_id" VARCHAR NOT NULL,
"chunk" INTEGER NOT NULL,
"content" VARCHAR NOT NULL,
"embedding" vector(1024) NOT NULL,
"created_at" TIMESTAMPTZ(3) NOT NULL DEFAULT CURRENT_TIMESTAMP,
"updated_at" TIMESTAMPTZ(3) NOT NULL,
CONSTRAINT "ai_context_embeddings_pkey" PRIMARY KEY ("id")
);
IF has_hnsw THEN
CREATE INDEX IF NOT EXISTS "ai_context_embeddings_idx"
ON "ai_context_embeddings" USING hnsw ("embedding" vector_cosine_ops);
END IF;
CREATE UNIQUE INDEX IF NOT EXISTS "ai_context_embeddings_context_id_file_id_chunk_key"
ON "ai_context_embeddings"("context_id", "file_id", "chunk");
IF NOT EXISTS (
SELECT 1 FROM pg_constraint
WHERE conname = 'ai_context_embeddings_context_id_fkey'
AND conrelid = 'public.ai_context_embeddings'::regclass
) THEN
ALTER TABLE "ai_context_embeddings"
ADD CONSTRAINT "ai_context_embeddings_context_id_fkey"
FOREIGN KEY ("context_id") REFERENCES "ai_contexts"("id")
ON DELETE CASCADE ON UPDATE CASCADE;
END IF;
END IF;
IF to_regclass('public.snapshots') IS NOT NULL THEN
CREATE TABLE IF NOT EXISTS "ai_workspace_embeddings" (
"workspace_id" VARCHAR NOT NULL,
"doc_id" VARCHAR NOT NULL,
"chunk" INTEGER NOT NULL,
"content" VARCHAR NOT NULL,
"embedding" vector(1024) NOT NULL,
"created_at" TIMESTAMPTZ(3) NOT NULL DEFAULT CURRENT_TIMESTAMP,
"updated_at" TIMESTAMPTZ(3) NOT NULL,
CONSTRAINT "ai_workspace_embeddings_pkey"
PRIMARY KEY ("workspace_id", "doc_id", "chunk")
);
IF has_hnsw THEN
CREATE INDEX IF NOT EXISTS "ai_workspace_embeddings_idx"
ON "ai_workspace_embeddings" USING hnsw ("embedding" vector_cosine_ops);
END IF;
IF NOT EXISTS (
SELECT 1 FROM pg_constraint
WHERE conname = 'ai_workspace_embeddings_workspace_id_doc_id_fkey'
AND conrelid = 'public.ai_workspace_embeddings'::regclass
) THEN
ALTER TABLE "ai_workspace_embeddings"
ADD CONSTRAINT "ai_workspace_embeddings_workspace_id_doc_id_fkey"
FOREIGN KEY ("workspace_id", "doc_id")
REFERENCES "snapshots"("workspace_id", "guid")
ON DELETE CASCADE ON UPDATE CASCADE;
END IF;
END IF;
IF to_regclass('public.ai_workspace_files') IS NOT NULL THEN
CREATE TABLE IF NOT EXISTS "ai_workspace_file_embeddings" (
"workspace_id" VARCHAR NOT NULL,
"file_id" VARCHAR NOT NULL,
"chunk" INTEGER NOT NULL,
"content" VARCHAR NOT NULL,
"embedding" vector(1024) NOT NULL,
"created_at" TIMESTAMPTZ(3) NOT NULL DEFAULT CURRENT_TIMESTAMP,
CONSTRAINT "ai_workspace_file_embeddings_pkey"
PRIMARY KEY ("workspace_id", "file_id", "chunk")
);
IF has_hnsw THEN
CREATE INDEX IF NOT EXISTS "ai_workspace_file_embeddings_idx"
ON "ai_workspace_file_embeddings" USING hnsw ("embedding" vector_cosine_ops);
END IF;
IF NOT EXISTS (
SELECT 1 FROM pg_constraint
WHERE conname = 'ai_workspace_file_embeddings_workspace_id_file_id_fkey'
AND conrelid = 'public.ai_workspace_file_embeddings'::regclass
) THEN
ALTER TABLE "ai_workspace_file_embeddings"
ADD CONSTRAINT "ai_workspace_file_embeddings_workspace_id_file_id_fkey"
FOREIGN KEY ("workspace_id", "file_id")
REFERENCES "ai_workspace_files"("workspace_id", "file_id")
ON DELETE CASCADE ON UPDATE CASCADE;
END IF;
END IF;
IF to_regclass('public.blobs') IS NOT NULL THEN
CREATE TABLE IF NOT EXISTS "ai_workspace_blob_embeddings" (
"workspace_id" VARCHAR NOT NULL,
"blob_id" VARCHAR NOT NULL,
"chunk" INTEGER NOT NULL,
"content" VARCHAR NOT NULL,
"embedding" vector(1024) NOT NULL,
"created_at" TIMESTAMPTZ(3) NOT NULL DEFAULT CURRENT_TIMESTAMP,
CONSTRAINT "ai_workspace_blob_embeddings_pkey"
PRIMARY KEY ("workspace_id", "blob_id", "chunk")
);
IF has_hnsw THEN
CREATE INDEX IF NOT EXISTS "ai_workspace_blob_embeddings_idx"
ON "ai_workspace_blob_embeddings" USING hnsw ("embedding" vector_cosine_ops);
END IF;
IF NOT EXISTS (
SELECT 1 FROM pg_constraint
WHERE conname = 'ai_workspace_blob_embeddings_workspace_id_blob_id_fkey'
AND conrelid = 'public.ai_workspace_blob_embeddings'::regclass
) THEN
ALTER TABLE "ai_workspace_blob_embeddings"
ADD CONSTRAINT "ai_workspace_blob_embeddings_workspace_id_blob_id_fkey"
FOREIGN KEY ("workspace_id", "blob_id")
REFERENCES "blobs"("workspace_id", "key")
ON DELETE CASCADE ON UPDATE CASCADE;
END IF;
END IF;
END $$;
packages/backend/server/scripts/self-host-predeploy.js
-15
View File
@@ -47,20 +47,6 @@ function runPrismaMigrations() {
});
}
function repairPgvectorEmbeddingTables() {
console.log('repairing copilot pgvector embedding tables.');
const sql = fs.readFileSync(
path.join(import.meta.dirname, 'repair-pgvector-embedding-tables.sql'),
'utf-8'
);
execSync('yarn prisma db execute --stdin --schema schema.prisma', {
encoding: 'utf-8',
env: process.env,
input: sql,
stdio: ['pipe', 'inherit', 'inherit'],
});
}
function runDataMigrations() {
console.log('running data migrations.');
execSync('yarn cli run', {
@@ -109,5 +95,4 @@ function fixFailedMigrations() {
prepare();
fixFailedMigrations();
runPrismaMigrations();
repairPgvectorEmbeddingTables();
runDataMigrations();
packages/backend/server/src/__tests__/copilot/byok.spec.ts
+14 -25
View File
@@ -4,7 +4,6 @@ import { PrismaClient } from '@prisma/client';
import type { TestFn } from 'ava';
import ava from 'ava';
import { Config } from '../../base';
import type { CurrentUser } from '../../core/auth';
import { BackendRuntimeProvider } from '../../core/backend-runtime';
import type { WorkspaceType } from '../../core/workspaces';
@@ -17,7 +16,6 @@ type Context = {
db: PrismaClient;
models: Models;
runtime: BackendRuntimeProvider;
config: Config;
resolver: WorkspaceByokResolver;
};
@@ -34,7 +32,6 @@ const testPrivateKey = privateKey
.toString();
const definition = {
version: 1,
endpoint: { kind: 'provider_default' },
models: [
{
@@ -59,7 +56,6 @@ test.before(async t => {
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);
});
@@ -73,31 +69,24 @@ test.after.always(async t => {
else process.env.AFFINE_PRIVATE_KEY = previousKey;
});
test('BYOK settings expose the configured custom endpoint policy', async t => {
test('BYOK settings expose the native effective 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);
const previous = t.context.config.copilot.byok.allowCustomEndpoint;
t.context.config.copilot.byok.allowCustomEndpoint = true;
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;
}
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.deepEqual(settings.policy, await t.context.runtime.getByokPolicy());
});
test('native BYOK runtime owns multi-model profile CAS, ordering, and credential rotation', async t => {
packages/backend/server/src/__tests__/copilot/capability-runtime.spec.ts
+60 -9
View File
@@ -196,28 +196,37 @@ test('text streaming consumes native generic events', async t => {
);
});
test('product event consumer attributes BYOK usage from structured identity', async t => {
test('product event consumer records route activity and real usage', async t => {
const records: unknown[] = [];
const activity: string[] = [];
const failures: string[] = [];
const models = {
copilotUsage: { create: async (value: unknown) => records.push(value) },
copilotWorkspaceByokConfig: {
touchUsed: async () => {},
markFailure: async () => {},
touchUsed: async (_workspaceId: string, profileId: string) =>
activity.push(profileId),
markFailure: async (
_workspaceId: string,
_profileId: string,
errorKind: string
) => failures.push(errorKind),
},
} as unknown as Models;
const consumer = new CopilotRuntimeEventConsumer(models);
const route = {
profileId: 'profile-1',
source: 'server' as const,
provider: 'openai',
model: 'opaque/model:B',
};
await consumer.consume(
[
{
type: 'usage',
route: {
profileId: 'profile-1',
source: 'server',
provider: 'openai',
model: 'opaque/model:B',
},
route,
usage: { input_tokens: 3, output_tokens: 2, total_tokens: 5 },
},
{ type: 'route_selected', route },
],
{ workspaceId: 'workspace-1', featureKind: 'chat' }
);
@@ -230,6 +239,48 @@ test('product event consumer attributes BYOK usage from structured identity', as
completionTokens: 2,
totalTokens: 5,
});
t.deepEqual(activity, ['profile-1']);
await consumer.consume(
[{ type: 'route_selected', route: { ...route, profileId: 'profile-2' } }],
{ workspaceId: 'workspace-1', featureKind: 'chat' }
);
t.is(records.length, 1);
t.deepEqual(activity, ['profile-1', 'profile-2']);
await consumer.consume(
[
{
type: 'usage',
route: { ...route, source: 'local', profileId: 'local-1' },
usage: {
prompt_tokens: 0,
completion_tokens: 0,
total_tokens: 0,
cached_tokens: 0,
},
},
{
type: 'route_selected',
route: { ...route, source: 'local', profileId: 'local-1' },
},
{
type: 'route_selected',
route: { ...route, source: 'affine_cloud', profileId: 'managed-1' },
},
{ type: 'route_failed', route, errorKind: 'upstream_error' },
],
{ workspaceId: 'workspace-1', featureKind: 'chat' }
);
t.like(records[1], {
providerSource: 'byok_local',
promptTokens: 0,
completionTokens: 0,
totalTokens: 0,
cachedTokens: 0,
});
t.deepEqual(activity, ['profile-1', 'profile-2']);
t.deepEqual(failures, ['upstream_error']);
});
test('tool callback validates arguments and preserves call identity', async t => {
packages/backend/server/src/__tests__/copilot/conversation-host.spec.ts
+42 -3
View File
@@ -29,7 +29,10 @@ function fixture(
sessionId,
content: 'hello',
attachments: [],
params: { tone: 'brief' },
params: {
tone: 'brief',
scopeSelectors: [{ kind: 'document', id: 'doc-2' }],
},
createdAt: new Date('2026-01-01T00:00:00.000Z'),
},
],
@@ -43,6 +46,7 @@ function fixture(
userId: 'user-1',
workspaceId: 'workspace-1',
docId: 'doc-1',
focus: { selectors: [] },
prompt: {
name: 'Chat With AFFiNE AI',
config: {},
@@ -96,9 +100,41 @@ function fixture(
const policy = {
hasQuota: async () => quota,
} as unknown as ConversationPolicy;
const runtime = {
putWorkspaceArtifact: async () => {
throw new Error('unexpected attachment');
},
compileTurnScope: async (input: {
selectors: unknown[];
preferredSourceIds?: string[];
}) => ({
version: 1,
resolvedAt: '2026-01-01T00:00:00.000Z',
selectors: input.selectors,
requiredDocIds: [],
requiredArtifactIds: [],
preferredSourceIds: input.preferredSourceIds ?? [],
retrieval: {
mode: input.selectors.length ? 'required' : 'workspace',
requiredDocIds: [],
requiredArtifactIds: [],
preferredSourceIds: input.preferredSourceIds ?? [],
},
}),
};
const attachmentAdmission = {
admitPromptAttachments: async () => [],
};
return {
host: new ConversationHost(sessions, submissionStore, mutex, policy),
host: new ConversationHost(
sessions,
submissionStore,
mutex,
policy,
runtime as never,
attachmentAdmission as never
),
sessionId,
token,
durable,
@@ -119,6 +155,9 @@ test('compat submission becomes one durable user turn and replays idempotently',
});
t.is(first.latestTurn?.content, 'hello');
t.deepEqual(first.latestTurn?.metadata, { tone: 'brief' });
t.deepEqual(first.latestTurn?.scopeSnapshot?.selectors, [
{ kind: 'document', id: 'doc-2', source: 'draft' },
]);
t.is(state.appendCount(), 1);
t.false(state.submissions.has(state.token));
t.truthy(state.accepted.get(state.token));
@@ -179,7 +218,7 @@ test('compat submission cannot be consumed by another session', async t => {
sessionId: 'session-other',
content: 'secret',
attachments: [],
params: { tone: 'brief' },
params: { tone: 'brief', scopeSelectors: [] },
createdAt: new Date(),
});
packages/backend/server/src/__tests__/copilot/copilot.e2e.ts
+48 -80
View File
@@ -2,6 +2,7 @@ import '../../plugins/copilot';
import { randomUUID } from 'node:crypto';
import { createCopilotMessageMutation } from '@affine/graphql';
import { McpAccessMode, PrismaClient } from '@prisma/client';
import type { TestFn } from 'ava';
import ava from 'ava';
@@ -9,24 +10,15 @@ import ava from 'ava';
import { Config } from '../../base';
import { ServerFeature, ServerService } from '../../core';
import { AuthService } from '../../core/auth';
import {
ContextCategories,
DocRole,
Models,
WorkspaceMemberStatus,
WorkspaceRole,
} from '../../models';
import { Models } 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 { installMockCopilotRuntime } from '../mocks';
import { createTestingApp, createWorkspace, type TestingApp } from '../utils';
import {
addContextCategory,
addContextFile,
chatWithImages,
chatWithText,
createCopilotContext,
createCopilotMessage,
createCopilotSession,
getCopilotSession,
@@ -89,7 +81,7 @@ test('disabled copilot hides its server feature and rejects every API transport'
}
});
test('session, compat message, text SSE and durable history share one public contract', async t => {
test('session, message, local context restriction and durable history share one public contract', async t => {
const { app } = t.context;
await app.signupV1();
const workspace = await createWorkspace(app);
@@ -139,6 +131,49 @@ test('session, compat message, text SSE and durable history share one public con
);
t.is(history.messages.filter(message => message.role === 'user').length, 1);
t.not(history.messages[0].id, token);
const localSessionId = await createCopilotSession(
app,
randomUUID(),
null,
'Chat With AFFiNE AI'
);
t.truthy(await createCopilotMessage(app, localSessionId, 'local hello'));
const localContextResponse = await app
.POST('/graphql')
.set('x-operation-name', createCopilotMessageMutation.op)
.send({
query: createCopilotMessageMutation.query,
variables: {
options: {
sessionId: localSessionId,
content: 'local context',
params: {
scopeSelectors: [{ kind: 'document', id: randomUUID() }],
},
},
},
})
.expect(200);
t.is(
localContextResponse.body.errors?.[0]?.message,
"Local workspaces don't support attachments or references."
);
await t.throwsAsync(
app.gql({
query: createCopilotMessageMutation,
variables: {
options: {
sessionId: localSessionId,
content: 'local attachment',
blobs: [new File(['attachment'], 'attachment.txt')],
},
},
}),
{
message: "Local workspaces don't support attachments or references.",
}
);
});
test('chat and history endpoints reject a different user', async t => {
@@ -181,73 +216,6 @@ test('image SSE emits persisted attachment events for action sessions', async t
t.truthy(attachment?.data);
});
test('context API rechecks write access and filters unreadable category docs', async t => {
const { app } = t.context;
const models = app.get(Models);
const owner = await app.signupV1();
const workspace = await createWorkspace(app);
const member = await app.signupV1();
await models.workspaceUser.set(
workspace.id,
member.id,
WorkspaceRole.Collaborator,
{ status: WorkspaceMemberStatus.Accepted }
);
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'))
);
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 hidden = await app.create(Mockers.DocSnapshot, {
workspaceId: workspace.id,
user: owner,
});
await app.create(Mockers.DocMeta, {
workspaceId: workspace.id,
docId: readable.id,
title: 'readable',
});
await app.create(Mockers.DocMeta, {
workspaceId: workspace.id,
docId: hidden.id,
title: 'hidden',
defaultRole: DocRole.None,
});
const category = await addContextCategory(
app,
contextId,
ContextCategories.Collection,
'favorites',
[readable.id, hidden.id]
);
t.deepEqual(
category.docs.map(doc => doc.id),
[readable.id]
);
});
test('MCP credentials remain endpoint-bound through rotate, revoke and expiry', async t => {
const { app } = t.context;
const auth = app.get(AuthService);
@@ -282,7 +250,7 @@ test('MCP credentials remain endpoint-bound through rotate, revoke and expiry',
(await provider.for(user.id, target.id, McpAccessMode.READ_ONLY)).tools.map(
tool => tool.name
),
['read_document', 'semantic_search', 'keyword_search']
['read_document', 'doc_search']
);
const rotated = await credentials.rotate(
packages/backend/server/src/__tests__/copilot/runtime-boundaries.spec.ts
+861 -15
View File
@@ -1,20 +1,35 @@
import { EventEmitter } from 'node:events';
import type { DelegatedToolRequest } from '@affine/realtime';
import type { PrismaClient } from '@prisma/client';
import ava from 'ava';
import { firstValueFrom } from 'rxjs';
import type { Config, JobQueue } from '../../base';
import {
AccessDenied,
type Config,
type EventBus,
type JobQueue,
} from '../../base';
import { ServerFeature, type ServerService } from '../../core';
import type { DocReader } from '../../core/doc';
import type { PermissionAccess } from '../../core/permission';
import { type RealtimePublisher, RealtimeRegistry } from '../../core/realtime';
import type { CanvasProjectionV1 } from '../../core/utils/blocksuite';
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,
promptMessageFromTurn,
type Turn,
turnFromChatMessage,
} from '../../plugins/copilot/core';
import { CopilotCronJobs } from '../../plugins/copilot/cron';
import { DelegatedEditorRealtimeProvider } from '../../plugins/copilot/delegated/realtime';
import { DelegatedEditorService } from '../../plugins/copilot/delegated/service';
import type { NativeEmbeddingService } from '../../plugins/copilot/embedding/native';
import {
CopilotFeatureGuard,
CopilotFeatureService,
@@ -22,17 +37,671 @@ import {
import type { PromptService } from '../../plugins/copilot/prompt';
import type { ResolvedPrompt } from '../../plugins/copilot/prompt/spec';
import { TextStreamParser } from '../../plugins/copilot/providers/utils';
import { ArtifactRetrievalService } from '../../plugins/copilot/retrieval/artifact';
import { DocumentRetrievalService } from '../../plugins/copilot/retrieval/document';
import {
projectActionEventToChatEvent,
projectActionResultToAssistantTurn,
} from '../../plugins/copilot/runtime/action-output-projector';
import type { ActionStreamHost } from '../../plugins/copilot/runtime/hosts/action-stream-host';
import {
collectAttachmentFootnotes,
collectDocumentFootnotes,
formatAttachmentFootnotes,
formatDocumentFootnotes,
} from '../../plugins/copilot/runtime/tool/footnotes';
import { NativeProviderAdapter } from '../../plugins/copilot/runtime/tool/native-adapter';
import type { TurnOrchestrator } from '../../plugins/copilot/runtime/turn-orchestrator';
import { ChatSession } from '../../plugins/copilot/session';
import {
ChatSession,
type ChatSessionService,
} from '../../plugins/copilot/session';
import type { CopilotStorage } from '../../plugins/copilot/storage';
import {
createArtifactReadTool,
createArtifactSearchTool,
} from '../../plugins/copilot/tools/artifact';
import { buildDocCanvasGetter } from '../../plugins/copilot/tools/doc-canvas-read';
import { buildDocumentSearch } from '../../plugins/copilot/tools/doc-search';
import type { IndexerService } from '../../plugins/indexer/service';
const test = ava;
test('delegated editor requests require exact identity and cancel on interruption', async t => {
const published: Array<{ event: Record<string, unknown> }> = [];
const publisher = {
publish: (
_topic: string,
_input: unknown,
event: Record<string, unknown>
) => published.push({ event }),
} as unknown as RealtimePublisher;
const delegated = new DelegatedEditorService(publisher);
delegated.upsert('user-1', 'connection-1', {
clientId: 'client-1',
sessionId: 'session-1',
workspaceId: 'workspace-1',
docId: 'doc-1',
editorStateId: 'state-1',
mode: 'page',
readonly: false,
focused: true,
capabilities: ['frontend_get_editor_state', 'frontend_read_selection'],
});
const result = delegated.execute(
{
user: 'user-1',
session: 'session-1',
workspace: 'workspace-1',
},
'frontend_get_editor_state',
{},
undefined,
{
runId: '3e476e0f-5841-4ab5-afca-610eca612ef1',
toolCallId: 'call_provider_1',
}
);
const request = published[0].event as unknown as DelegatedToolRequest;
t.is(request.toolCallId, 'call_provider_1');
const registry = new RealtimeRegistry();
new DelegatedEditorRealtimeProvider(
registry,
{ broadcast: () => {} } as unknown as EventBus,
{} as ChatSessionService,
delegated
).onModuleInit();
t.notThrows(() =>
registry.getRequest('copilot.delegated.tool.respond').input.parse({
requestId: request.requestId,
runId: request.runId,
toolCallId: request.toolCallId,
sessionId: request.sessionId,
workspaceId: request.workspaceId,
docId: request.docId,
clientId: request.clientId,
editorStateId: request.editorStateId,
result: { mode: 'page' },
})
);
t.false(
delegated.receive('user-1', {
...request,
editorStateId: 'stale-state',
result: { mode: 'page' },
})
);
t.false(
delegated.receive('user-1', {
...request,
workspaceId: 'workspace-2',
result: { editor_state_id: 'state-1', mode: 'page' },
})
);
t.true(
delegated.receive('user-1', {
...request,
result: { editor_state_id: 'state-1', mode: 'page' },
})
);
t.deepEqual(await result, {
editor_state_id: 'state-1',
mode: 'page',
});
const selection = delegated.execute(
{
user: 'user-1',
session: 'session-1',
workspace: 'workspace-1',
},
'frontend_read_selection',
{}
);
const selectionRequest = published.at(-1)
?.event as unknown as DelegatedToolRequest;
t.true(
delegated.receive('user-1', {
...selectionRequest,
result: { editor_state_id: 'state-1', text: 'live content' },
})
);
t.deepEqual(await selection, {
editor_state_id: 'state-1',
text: 'live content',
source: {
type: 'document',
workspace_id: 'workspace-1',
doc_id: 'doc-1',
revision: 'state-1',
},
});
const controller = new AbortController();
const aborted = delegated.execute(
{
user: 'user-1',
session: 'session-1',
workspace: 'workspace-1',
},
'frontend_get_editor_state',
{},
controller.signal
);
controller.abort();
t.like(await aborted, { error: { code: 'ABORTED', retryable: false } });
t.is(published.at(-1)?.event.type, 'cancel');
const preAbortedController = new AbortController();
preAbortedController.abort();
const preAborted = await delegated.execute(
{
user: 'user-1',
session: 'session-1',
workspace: 'workspace-1',
},
'frontend_get_editor_state',
{},
preAbortedController.signal
);
t.like(preAborted, { error: { code: 'ABORTED', retryable: false } });
const disconnected = delegated.execute(
{
user: 'user-1',
session: 'session-1',
workspace: 'workspace-1',
},
'frontend_get_editor_state',
{}
);
delegated.onDisconnect({ connectionId: 'connection-1' });
t.like(await disconnected, {
error: { code: 'FRONTEND_DISCONNECTED', retryable: true },
});
t.like(published.at(-1)?.event, { type: 'cancel', reason: 'disconnect' });
});
test('canvas reads expose top-level and frame-owned canvas blocks', async t => {
const projection: CanvasProjectionV1 = {
version: 1,
docId: 'doc-1',
revision: 'revision-1',
title: 'Canvas',
counts: {},
warnings: [],
blocks: [
{
id: 'page-1',
type: 'paragraph',
visibility: 'page',
text: 'Page only',
childIds: [],
},
{
id: 'frame-1',
type: 'frame',
visibility: 'edgeless',
childIds: ['edgeless-1', 'shape-1'],
},
{
id: 'edgeless-1',
type: 'edgeless-text',
visibility: 'edgeless',
text: 'Frame text',
childIds: [],
},
{
id: 'edgeless-2',
type: 'edgeless-text',
visibility: 'edgeless',
text: 'Top-level text',
childIds: [],
},
],
elements: [
{ id: 'shape-1', type: 'shape', frameId: 'frame-1', childIds: [] },
{ id: 'shape-2', type: 'shape', childIds: [] },
],
};
const getter = buildDocCanvasGetter(
{
user: () => ({
workspace: () => ({ doc: () => ({ can: async () => true }) }),
}),
} as unknown as PermissionAccess,
{ getDocCanvas: async () => projection } as unknown as DocReader,
{
workspace: { get: async () => ({ id: 'workspace-1' }) },
} as unknown as Models
);
const options = { user: 'user-1', workspace: 'workspace-1' };
const overview = await getter(
options,
'doc-1',
{ kind: 'overview' },
undefined,
50
);
t.deepEqual(
'blocks' in overview ? overview.blocks.map(block => block.id) : [],
['edgeless-2', 'frame-1']
);
t.deepEqual(
'elements' in overview ? overview.elements.map(element => element.id) : [],
['shape-2']
);
const frame = await getter(
options,
'doc-1',
{ kind: 'frame', frame_id: 'frame-1' },
undefined,
50
);
t.deepEqual('blocks' in frame ? frame.blocks.map(block => block.id) : [], [
'edgeless-1',
'frame-1',
]);
t.deepEqual(
'elements' in frame ? frame.elements.map(element => element.id) : [],
['shape-1']
);
const scopedGetter = buildDocCanvasGetter(
{} as PermissionAccess,
{} as DocReader,
{} as Models,
{ mode: 'selected', allowedDocIds: ['doc-2'] }
);
const outsideScope = await scopedGetter(
options,
'doc-1',
{ kind: 'overview' },
undefined,
50
);
t.like(outsideScope, { code: 'DOC_SCOPE_DENIED' });
});
test('document tools enforce the user-selected hard scope', async t => {
const hit = {
docId: 'doc-1',
title: 'Doc',
excerpt: 'excerpt',
visibility: 'page' as const,
score: 1,
unitId: 'block:1',
};
const searchCalls: Array<string[] | undefined> = [];
const retrieval = {
search: async (
_options: unknown,
_query: string,
docIds: string[] | undefined,
_limit: number
) => {
searchCalls.push(docIds);
return {
retrievalMode: 'hybrid',
degradedReason: undefined,
hits: [hit],
};
},
} as unknown as DocumentRetrievalService;
const options = { user: 'user-1', workspace: 'workspace-1' };
const readableAc = {
user: () => ({
workspace: () => ({
docs: async <T extends { docId: string }>(candidates: T[]) =>
candidates.filter(candidate => candidate.docId !== 'hidden-doc'),
}),
}),
} as unknown as PermissionAccess;
const documentModels = {
doc: {
findMetas: async (ids: Array<{ docId: string }>) =>
ids.map(({ docId }) => ({
docId,
title: `title-${docId}`,
updatedAt: new Date(1),
})),
},
} as unknown as Models;
const lexicalIndexer = {
searchDocsByKeyword: async () => [
{
docId: 'shared-doc',
title: 'Lexical title',
highlight: 'lexical passage',
unitId: 'block:shared',
visibility: 'page',
projectionVersion: '1',
sourceHash: 'hash',
},
],
} as unknown as IndexerService;
const vectorSearch = {
canEmbedding: true,
matchWorkspaceDocCandidates: async () => [
{
docId: 'shared-doc',
chunk: 0,
content: 'vector passage',
distance: 0.1,
unitId: 'block:shared',
visibility: 'page' as const,
},
{
docId: 'hidden-doc',
chunk: 0,
content: 'hidden passage',
distance: 0.2,
unitId: 'block:hidden',
visibility: 'page' as const,
},
],
rerankWorkspaceDocs: async (
_workspaceId: string,
_query: string,
candidates: Array<{
docId: string;
chunk: number;
content: string;
distance: number;
unitId: string;
visibility: 'page';
}>
) => candidates,
};
const hybrid = new DocumentRetrievalService(
{ indexer: { enabled: true } } as Config,
readableAc,
lexicalIndexer,
vectorSearch,
documentModels
);
const hybridResult = await hybrid.search(options, 'query', undefined, 10);
t.is(hybridResult.retrievalMode, 'hybrid');
t.deepEqual(
hybridResult.hits.map(result => result.docId),
['shared-doc']
);
t.true(hybridResult.hits[0].score > 1 / 61);
const lexicalOnly = new DocumentRetrievalService(
{ indexer: { enabled: true } } as Config,
readableAc,
lexicalIndexer,
{ ...vectorSearch, canEmbedding: false },
documentModels
);
const lexicalResult = await lexicalOnly.search(
options,
'query',
undefined,
10
);
t.is(lexicalResult.retrievalMode, 'lexical');
t.is(lexicalResult.degradedReason, 'VECTOR_UNAVAILABLE');
const vectorOnly = new DocumentRetrievalService(
{ indexer: { enabled: false } } as Config,
readableAc,
lexicalIndexer,
vectorSearch,
documentModels
);
const vectorResult = await vectorOnly.search(options, 'query', undefined, 10);
t.is(vectorResult.retrievalMode, 'vector');
t.is(vectorResult.degradedReason, 'LEXICAL_UNAVAILABLE');
t.deepEqual(
vectorResult.hits.map(result => result.docId),
['shared-doc']
);
// model omits doc_ids: pinned scope applies
let search = buildDocumentSearch(retrieval, options, {
mode: 'selected',
allowedDocIds: ['pinned-1'],
});
let result: any = await search('query', undefined, 10);
t.deepEqual(searchCalls.pop(), ['pinned-1']);
t.is(result.hits[0].doc_id, 'doc-1');
t.is(result.hits[0].source.doc_id, 'doc-1');
// model-provided ids cannot replace the complete user-selected scope
search = buildDocumentSearch(retrieval, options, {
mode: 'selected',
allowedDocIds: ['pinned-1'],
});
result = await search('query', ['other-1'], 10);
t.deepEqual(searchCalls.pop(), ['pinned-1']);
t.is(result.hits[0].doc_id, 'doc-1');
// an empty array keeps the pinned scope
search = buildDocumentSearch(retrieval, options, {
mode: 'selected',
allowedDocIds: ['pinned-1'],
});
await search('query', [], 10);
t.deepEqual(searchCalls.pop(), ['pinned-1']);
// an explicitly selected empty category remains an empty hard scope
search = buildDocumentSearch(retrieval, options, {
mode: 'selected',
allowedDocIds: [],
});
result = await search('query', undefined, 10);
t.is(searchCalls.length, 0);
t.is(result.scope_mode, 'selected');
t.is(result.scope_doc_count, 0);
t.deepEqual(result.hits, []);
// no pinned scope: omission searches the whole workspace
search = buildDocumentSearch(retrieval, options);
await search('query', undefined, 10);
t.is(searchCalls.pop(), undefined);
// missing identity is a non-retryable tool error
const unauthenticated: any = await buildDocumentSearch(retrieval, undefined, {
mode: 'selected',
allowedDocIds: ['pinned-1'],
})('query', undefined, 10);
t.is(unauthenticated.code, 'INVALID_CONTEXT');
t.is(searchCalls.length, 0);
const artifactCalls: Array<{
kind: string;
sourceKey?: string;
requiredArtifactIds: string[];
}> = [];
const artifactScope = {
mode: 'required' as const,
requiredDocIds: [],
requiredArtifactIds: ['6ba7b810-9dad-11d1-80b4-00c04fd430c8'],
preferredSourceIds: [],
};
const artifactEmbedding = {
match: async (
_workspaceId: string,
_query: string,
kind: string,
retrievalScope: typeof artifactScope,
_limit: number,
signal?: AbortSignal
) => {
signal?.throwIfAborted();
artifactCalls.push({
kind,
requiredArtifactIds: retrievalScope.requiredArtifactIds,
});
return [];
},
readSourceContent: async (
_workspaceId: string,
kind: string,
sourceKey: string,
retrievalScope: typeof artifactScope
) => {
artifactCalls.push({
kind,
sourceKey,
requiredArtifactIds: retrievalScope.requiredArtifactIds,
});
if (!retrievalScope.requiredArtifactIds.includes(sourceKey)) {
throw new Error('embedding_source_out_of_scope');
}
return {
content: 'artifact body',
revision: 'revision-1',
mimeType: 'text/plain',
name: 'note.txt',
truncated: false,
};
},
} as unknown as NativeEmbeddingService;
const artifactRetrieval = new ArtifactRetrievalService(
{
user: () => ({
workspace: () => ({
allowLocal: () => ({ can: async () => true }),
}),
}),
} as unknown as PermissionAccess,
artifactEmbedding,
{
workspaceArtifact: {
findMany: async () => [
{
id: artifactScope.requiredArtifactIds[0],
displayName: null,
canonicalMediaType: 'text/plain',
},
],
},
aiMessageArtifact: {
findMany: async () => [
{
artifactId: artifactScope.requiredArtifactIds[0],
displayName: 'original-note.txt',
},
],
},
} as unknown as PrismaClient
);
const artifactOptions = {
user: 'user-1',
workspace: 'workspace-1',
billingUnitId: 'message-1',
retrievalScope: artifactScope,
};
const artifactSearch = createArtifactSearchTool(
artifactRetrieval,
artifactOptions
);
const artifactSearchResult = await artifactSearch.execute?.(
{ query: 'query' },
{}
);
t.deepEqual(artifactCalls.shift(), {
kind: 'artifact',
requiredArtifactIds: artifactScope.requiredArtifactIds,
});
t.deepEqual(artifactCalls.shift(), {
kind: 'artifact',
sourceKey: artifactScope.requiredArtifactIds[0],
requiredArtifactIds: artifactScope.requiredArtifactIds,
});
t.like(artifactSearchResult, {
hits: [
{
excerpt: 'artifact body',
source: { type: 'artifact', name: 'original-note.txt' },
},
],
});
const artifactRead = createArtifactReadTool(
artifactRetrieval,
artifactOptions
);
const artifactReadResult = await artifactRead.execute?.(
{ artifact_id: artifactScope.requiredArtifactIds[0] },
{}
);
t.like(artifactReadResult, {
source: {
artifact_id: artifactScope.requiredArtifactIds[0],
name: 'original-note.txt',
},
});
const fallbackArtifactRetrieval = new ArtifactRetrievalService(
{
user: () => ({
workspace: () => ({
allowLocal: () => ({ can: async () => true }),
}),
}),
} as unknown as PermissionAccess,
artifactEmbedding,
{
workspaceArtifact: { findMany: async () => [] },
aiMessageArtifact: { findMany: async () => [] },
} as unknown as PrismaClient
);
t.like(
await fallbackArtifactRetrieval.read({
userId: 'user-1',
workspaceId: 'workspace-1',
artifactId: artifactScope.requiredArtifactIds[0],
retrieval: artifactScope,
}),
{ name: 'note.txt', mimeType: 'text/plain' }
);
const deniedArtifactRetrieval = new ArtifactRetrievalService(
{
user: () => ({
workspace: () => ({
allowLocal: () => ({ can: async () => false }),
}),
}),
} as unknown as PermissionAccess,
artifactEmbedding,
{} as PrismaClient
);
await t.throwsAsync(
deniedArtifactRetrieval.read({
userId: 'user-1',
workspaceId: 'workspace-1',
artifactId: artifactScope.requiredArtifactIds[0],
retrieval: artifactScope,
}),
{ instanceOf: AccessDenied }
);
const deniedArtifactRead = await artifactRead.execute?.(
{ artifact_id: '6ba7b811-9dad-11d1-80b4-00c04fd430c8' },
{}
);
t.like(deniedArtifactRead, { code: 'ARTIFACT_UNAVAILABLE' });
const abortedSearch = new AbortController();
abortedSearch.abort();
await t.throwsAsync(
artifactRetrieval.search({
userId: 'user-1',
workspaceId: 'workspace-1',
query: 'query',
retrieval: artifactScope,
limit: 5,
signal: abortedSearch.signal,
}),
{ name: 'AbortError' }
);
});
test('copilot config controls the server feature and request admission', t => {
const config = { copilot: { enabled: false } } as Config;
const features = new Set<ServerFeature>();
@@ -91,6 +760,7 @@ test('chat session preserves prompt params, attachments, stash and revert semant
userId: 'user-1',
workspaceId: 'workspace-1',
docId: 'doc-1',
focus: { selectors: [] },
prompt,
turns: [turn('session-1', 'user', 'persisted')],
},
@@ -127,13 +797,7 @@ test('chat session preserves prompt params, attachments, stash and revert semant
{
role: 'assistant',
content: 'answer',
attachments: [
{
kind: 'file_handle',
fileHandle: 'file-1',
mimeType: 'application/pdf',
},
],
attachments: undefined,
params: { word: 'world' },
},
]);
@@ -144,6 +808,13 @@ test('chat session preserves prompt params, attachments, stash and revert semant
saved[0].map(item => item.content),
['answer']
);
t.deepEqual(saved[0][0].attachments, [
{
kind: 'file_handle',
fileHandle: 'file-1',
mimeType: 'application/pdf',
},
]);
session.pushTurn(turn('session-1', 'user', 'retry'));
session.pushTurn(turn('session-1', 'assistant', 'retry answer'));
@@ -205,8 +876,23 @@ test('chat message adapters preserve and canonicalize assistant render trace', t
t.deepEqual(chatMessageFromTurn(converted), {
...message,
attachments: undefined,
scopeSnapshot: undefined,
streamObjects: converted.renderTrace,
});
t.deepEqual(
promptMessageFromTurn({
...converted,
attachments: [
{
attachment: 'data:text/plain;base64,dGV4dA==',
mimeType: 'text/plain',
},
{ attachment: 'data:image/png;base64,aW1hZ2U=', mimeType: 'image/png' },
],
}).attachments,
[{ attachment: 'data:image/png;base64,aW1hZ2U=', mimeType: 'image/png' }]
);
});
test('action output projection preserves public SSE and assistant-turn contracts', t => {
@@ -216,6 +902,7 @@ test('action output projection preserves public SSE and assistant-turn contracts
userId: 'user-1',
workspaceId: 'workspace-1',
docId: 'doc-1',
focus: { selectors: [] },
prompt,
turns: [],
},
@@ -256,9 +943,96 @@ test('action output projection preserves public SSE and assistant-turn contracts
}),
null
);
t.is(
formatDocumentFootnotes([
{
type: 'document',
workspace_id: 'workspace-1',
doc_id: 'doc-1',
title: 'Getting Started',
revision: 'revision-1',
visibility: 'edgeless',
},
{
type: 'document',
workspace_id: 'workspace-1',
doc_id: 'doc-1',
title: 'Getting Started',
revision: 'revision-1',
visibility: 'edgeless',
element_id: 'element-1',
},
]),
'\n\n[^doc-1]\n\n[^doc-1]: {"type":"doc","docId":"doc-1","title":"Getting Started"}'
);
t.is(
formatAttachmentFootnotes([
{
artifactId: 'artifact-1',
fileName: 'notes.txt',
fileType: 'text/plain',
},
]),
'\n\n[^attachment-1]\n\n[^attachment-1]: {"type":"attachment","artifactId":"artifact-1","fileName":"notes.txt","fileType":"text/plain"}'
);
t.deepEqual(
collectDocumentFootnotes({
type: 'tool_result',
call_id: 'call-1',
name: 'frontend_read_selection',
arguments: {},
output: {
source: {
type: 'document',
workspace_id: 'workspace-1',
doc_id: 'doc-1',
revision: 'state-1',
},
},
}),
[
{
type: 'document',
workspace_id: 'workspace-1',
doc_id: 'doc-1',
title: '',
revision: 'state-1',
visibility: undefined,
block_id: undefined,
element_id: undefined,
frame_id: undefined,
},
]
);
t.deepEqual(
collectAttachmentFootnotes({
type: 'tool_result',
call_id: 'call-2',
name: 'artifact_search',
arguments: {},
output: {
hits: [
{
source: {
type: 'artifact',
workspace_id: 'workspace-1',
artifact_id: 'artifact-1',
},
},
],
},
}),
[
{
artifactId: 'artifact-1',
fileName: 'Attachment',
fileType: 'application/octet-stream',
},
]
);
});
test('text stream parser keeps reasoning and tool output distinct from answer text', t => {
test('text stream parser keeps reasoning and tool output distinct from answer text', async t => {
const parser = new TextStreamParser();
const output = [
parser.parse({ type: 'reasoning-delta', text: 'Think' }),
@@ -286,6 +1060,83 @@ test('text stream parser keeps reasoning and tool output distinct from answer te
() => parser.parse({ type: 'error', error: { message: 'failed' } }),
{ message: 'failed' }
);
const adapter = new NativeProviderAdapter(async function* () {
yield {
type: 'citation',
index: 1,
url: 'https://affine.pro',
};
yield {
type: 'tool_result',
call_id: 'call-1',
name: 'artifact_read',
arguments: {},
output: {
artifactId: 'artifact-1',
fileName: 'notes.txt',
fileType: 'text/plain',
},
};
yield {
type: 'tool_result',
call_id: 'call-2',
name: 'frontend_read_selection',
arguments: {},
output: {
text: 'live content',
source: {
type: 'document',
workspace_id: 'workspace-1',
doc_id: 'doc-1',
revision: 'state-1',
},
},
};
yield { type: 'done' };
});
const streamObjects = [];
for await (const item of adapter.streamObject({
model: 'test',
messages: [],
})) {
streamObjects.push(item);
}
t.deepEqual(streamObjects.at(-1), {
type: 'text-delta',
textDelta: '\n\n[^doc-1]\n\n[^doc-1]: {"type":"doc","docId":"doc-1"}',
});
const streamOutput = streamObjects
.filter(item => item.type === 'text-delta')
.map(item => item.textDelta)
.join('');
t.true(streamOutput.includes('"url":"https%3A%2F%2Faffine.pro"'));
t.true(streamOutput.includes('[^attachment-1]'));
t.true(streamOutput.includes('"artifactId":"artifact-1"'));
const textAdapter = new NativeProviderAdapter(async function* () {
yield {
type: 'tool_result',
call_id: 'call-1',
name: 'artifact_read',
arguments: {},
output: {
artifactId: 'artifact-1',
fileName: 'notes.txt',
fileType: 'text/plain',
},
};
yield { type: 'done' };
});
let textOutput = '';
for await (const chunk of textAdapter.streamText({
model: 'test',
messages: [],
})) {
textOutput += chunk;
}
t.true(textOutput.includes('[^attachment-1]'));
t.true(textOutput.includes('"artifactId":"artifact-1"'));
});
test('history prompt preload excludes system messages and precedes durable history', t => {
@@ -364,11 +1215,6 @@ test('title policy and cron scheduling retain background-job invariants', async
{},
{ jobId: 'daily-copilot-generate-missing-titles' },
],
[
'copilot.workspace.cleanupTrashedDocEmbeddings',
{},
{ jobId: 'daily-copilot-cleanup-trashed-doc-embeddings' },
],
[
'copilot.session.generateTitle',
{ sessionId: 'session-1' },
packages/backend/server/src/__tests__/e2e/create-app.ts
+40 -6
View File
@@ -3,7 +3,11 @@ import assert from 'node:assert';
import { gqlFetcherFactory } from '@affine/graphql';
import { INestApplication, ModuleMetadata } from '@nestjs/common';
import { NestApplication } from '@nestjs/core';
import { Test, TestingModuleBuilder } from '@nestjs/testing';
import {
Test,
type TestingModule,
TestingModuleBuilder,
} from '@nestjs/testing';
import { PrismaClient } from '@prisma/client';
import cookieParser from 'cookie-parser';
import graphqlUploadExpress from 'graphql-upload/graphqlUploadExpress.mjs';
@@ -22,6 +26,7 @@ import {
import { ThrottlerStorage } from '../../base/throttler';
import { SocketIoAdapter } from '../../base/websocket';
import { AuthGuard, AuthService } from '../../core/auth';
import { BACKEND_RUNTIME_CONFIG_PATHS } from '../../core/backend-runtime';
import { Mailer } from '../../core/mail';
import { Models } from '../../models';
import {
@@ -33,6 +38,7 @@ import {
MockUserInput,
} from '../mocks';
import { parseCookies, TEST_LOG_LEVEL } from '../utils';
import { createTestRuntimeConfig } from '../utils/runtime-config';
interface TestingAppMetadata {
tapModule?(m: TestingModuleBuilder): void;
@@ -235,6 +241,9 @@ export class TestingApp extends NestApplication {
export async function createApp(
metadata: TestingAppMetadata = {}
): Promise<TestingApp> {
const runtimeConfig = await createTestRuntimeConfig(
new ConfigFactory().config.db.datasourceUrl
);
const { buildAppModule } = await import('../../app.module');
const { tapModule, tapApp } = metadata;
@@ -244,27 +253,36 @@ export async function createApp(
builder.overrideProvider(Mailer).useValue(new MockMailer());
builder.overrideProvider(JobQueue).useValue(new MockJobQueue());
builder
.overrideProvider(BACKEND_RUNTIME_CONFIG_PATHS)
.useValue([runtimeConfig.configPath]);
// when custom override happens
if (tapModule) {
tapModule(builder);
}
const module = await builder.compile();
let module: TestingModule;
try {
module = await builder.compile();
} catch (error) {
await runtimeConfig.cleanup();
throw error;
}
module.get(ConfigFactory).override({
storages: {
avatar: {
storage: {
provider: 'assetpack',
bucket: 'avatars',
config: { path: '/tmp/affine-test-storage' },
config: { path: runtimeConfig.storagePath },
},
},
blob: {
storage: {
provider: 'assetpack',
bucket: 'blobs',
config: { path: '/tmp/affine-test-storage' },
config: { path: runtimeConfig.storagePath },
},
},
},
@@ -272,7 +290,7 @@ export async function createApp(
storage: {
provider: 'assetpack',
bucket: 'copilot',
config: { path: '/tmp/affine-test-storage' },
config: { path: runtimeConfig.storagePath },
},
},
});
@@ -284,6 +302,17 @@ export async function createApp(
bodyParser: true,
rawBody: true,
});
const close = app.close.bind(app);
let closePromise: Promise<void> | undefined;
app.close = () => {
return (closePromise ??= (async () => {
try {
await close();
} finally {
await runtimeConfig.cleanup();
}
})());
};
const logger = new AFFiNELogger();
logger.setLogLevels([TEST_LOG_LEVEL]);
@@ -309,7 +338,12 @@ export async function createApp(
tapApp(app);
}
await app.init();
try {
await app.init();
} catch (error) {
await app.close();
throw error;
}
return app;
}
packages/backend/server/src/__tests__/e2e/doc-service/controller.spec.ts
+3 -1
View File
@@ -26,7 +26,9 @@ e2e('should get doc markdown success', async t => {
.expect(200)
.expect('Content-Type', 'application/json; charset=utf-8');
t.snapshot(res.body);
const { revision, ...body } = res.body;
t.regex(revision, /^\d+$/);
t.snapshot(body);
});
e2e('should get doc markdown return null when doc not exists', async t => {
packages/backend/server/src/__tests__/e2e/storage/r2-proxy.spec.ts
+2 -2
View File
@@ -369,7 +369,7 @@ e2e.serial('should proxy single upload with valid signature', async t => {
e2e.serial('should proxy multipart upload and return etag', async t => {
const { workspace } = await setupWorkspace();
const key = 'multipart-object';
const key = sha256Base64urlWithPadding(Buffer.from('multipart-object'));
const totalSize = MULTIPART_THRESHOLD + 1024;
const init = await createBlobUpload(workspace.id, key, totalSize, 'bin');
@@ -404,7 +404,7 @@ e2e.serial(
'should resume multipart upload and return uploaded parts',
async t => {
const { workspace } = await setupWorkspace();
const key = 'multipart-resume';
const key = sha256Base64urlWithPadding(Buffer.from('multipart-resume'));
const totalSize = MULTIPART_THRESHOLD + 1024;
const init1 = await createBlobUpload(workspace.id, key, totalSize, 'bin');
packages/backend/server/src/__tests__/models/__snapshots__/copilot-context.spec.ts.md
-247
View File
@@ -1,247 +0,0 @@
# Snapshot report for `src/__tests__/models/copilot-context.spec.ts`
The actual snapshot is saved in `copilot-context.spec.ts.snap`.
Generated by [AVA](https://avajs.dev).
## should get null for non-exist job
> should return null for non-exist job
null
## should insert embedding by doc id
> should match file embedding
[
{
fileId: 'file-id',
},
]
> should return empty array when embedding is deleted
[]
> should match workspace embedding
[
{
docId: 'doc1',
},
]
> should return empty array when doc is ignored
[]
> should return workspace embedding
[
{
docId: 'doc1',
},
]
> should return empty array when embedding deleted
[]
## should check embedding table
> should return true when embedding table is available
true
## should merge doc status correctly
> basic doc status merge
[
{
id: 'doc1',
status: 'processing',
},
{
id: 'doc2',
status: 'processing',
},
{
id: 'doc3',
status: 'failed',
},
{
id: 'doc4',
status: 'processing',
},
]
> mixed doc status merge
[
{
id: 'doc5',
status: 'finished',
},
{
id: 'doc5',
status: 'finished',
},
{
id: 'doc6',
status: 'processing',
},
{
id: 'doc6',
status: 'failed',
},
{
id: 'doc7',
status: 'processing',
},
]
> edge cases results
[
{
case: 0,
length: 1,
statuses: [
'processing',
],
},
{
case: 1,
length: 1,
statuses: [
'processing',
],
},
{
case: 2,
length: 100,
statuses: [
'processing',
'processing',
'processing',
'processing',
'processing',
'processing',
'processing',
'processing',
'processing',
'processing',
'processing',
'processing',
'processing',
'processing',
'processing',
'processing',
'processing',
'processing',
'processing',
'processing',
'processing',
'processing',
'processing',
'processing',
'processing',
'processing',
'processing',
'processing',
'processing',
'processing',
'processing',
'processing',
'processing',
'processing',
'processing',
'processing',
'processing',
'processing',
'processing',
'processing',
'processing',
'processing',
'processing',
'processing',
'processing',
'processing',
'processing',
'processing',
'processing',
'processing',
'processing',
'processing',
'processing',
'processing',
'processing',
'processing',
'processing',
'processing',
'processing',
'processing',
'processing',
'processing',
'processing',
'processing',
'processing',
'processing',
'processing',
'processing',
'processing',
'processing',
'processing',
'processing',
'processing',
'processing',
'processing',
'processing',
'processing',
'processing',
'processing',
'processing',
'processing',
'processing',
'processing',
'processing',
'processing',
'processing',
'processing',
'processing',
'processing',
'processing',
'processing',
'processing',
'processing',
'processing',
'processing',
'processing',
'processing',
'processing',
'processing',
'processing',
],
},
]
## should handle concurrent mergeDocStatus calls
> concurrent calls results
[
{
call: 1,
status: 'finished',
},
{
call: 2,
status: 'finished',
},
{
call: 3,
status: 'processing',
},
]
packages/backend/server/src/__tests__/models/__snapshots__/copilot-context.spec.ts.snap
BIN
View File
Binary file not shown.
packages/backend/server/src/__tests__/models/__snapshots__/copilot-session.spec.ts.md
+2 -2
View File
@@ -559,11 +559,11 @@ Generated by [AVA](https://avajs.dev).
> attach and detach operation results
{
attachPhase: {
afterAttach: {
bothSessionsPresent: true,
docSessionCount: 2,
},
detachPhase: {
afterDetach: {
originalDocSessionRemains: true,
workspaceSessionExists: true,
},
packages/backend/server/src/__tests__/models/__snapshots__/copilot-session.spec.ts.snap
BIN
View File
Binary file not shown.
packages/backend/server/src/__tests__/models/__snapshots__/copilot-workspace.spec.ts.md
-140
View File
@@ -1,140 +0,0 @@
# Snapshot report for `src/__tests__/models/copilot-workspace.spec.ts`
The actual snapshot is saved in `copilot-workspace.spec.ts.snap`.
Generated by [AVA](https://avajs.dev).
## should manage copilot workspace ignored docs
> should add ignored doc
1
> should return added doc
[
{
docId: 'doc1',
},
]
> should return ignored docs in workspace
[
'doc1',
]
> should not change if ignored doc exists
0
> should not add ignored doc again
[
{
docId: 'doc1',
},
]
> should add new ignored doc
1
> should add ignored doc
[
{
docId: 'new_doc',
},
{
docId: 'doc1',
},
]
> should remove ignored doc
[
{
docId: 'new_doc',
},
]
## should insert and search embedding
> should match workspace file embedding
[
{
blobId: 'blob1',
chunk: 0,
content: 'content',
distance: 0,
mimeType: 'text/plain',
name: 'file1',
},
]
> should match workspace blob embedding
[
{
blobId: 'blob-test',
chunk: 0,
content: 'blob content',
distance: 0,
},
]
> should find docs to embed
1
> should not find docs to embed
0
> should find docs to embed
1
> should not find docs to embed
0
## should check need to be embedded
> document with no embedding should need embedding
true
> document with recent embedding should not need embedding
false
> document updated after embedding and older-than-10m should need embedding
true
> should not need embedding when only 10-minute window passed without updates
false
> should need embedding when doc updated and last embedding older than 10 minutes
true
## should filter outdated doc id style in embedding status
> should include modern doc format
{
embedded: 0,
total: 1,
}
> should count docs after filtering outdated
{
embedded: 1,
total: 1,
}
packages/backend/server/src/__tests__/models/__snapshots__/copilot-workspace.spec.ts.snap
BIN
View File
Binary file not shown.
packages/backend/server/src/__tests__/models/copilot-context.spec.ts
-417
View File
@@ -1,417 +0,0 @@
import { randomUUID } from 'node:crypto';
import { PrismaClient, User, Workspace } from '@prisma/client';
import ava, { TestFn } from 'ava';
import Sinon from 'sinon';
import { Config } from '../../base';
import {
ContextEmbedStatus,
CopilotContextModel,
CopilotSessionModel,
CopilotWorkspaceConfigModel,
UserModel,
WorkspaceModel,
} from '../../models';
import { createTestingModule, type TestingModule } from '../utils';
import { cleanObject } from '../utils/copilot';
interface Context {
config: Config;
module: TestingModule;
db: PrismaClient;
user: UserModel;
workspace: WorkspaceModel;
copilotSession: CopilotSessionModel;
copilotContext: CopilotContextModel;
copilotWorkspace: CopilotWorkspaceConfigModel;
}
const test = ava as TestFn<Context>;
test.before(async t => {
const module = await createTestingModule();
t.context.user = module.get(UserModel);
t.context.workspace = module.get(WorkspaceModel);
t.context.copilotSession = module.get(CopilotSessionModel);
t.context.copilotContext = module.get(CopilotContextModel);
t.context.copilotWorkspace = module.get(CopilotWorkspaceConfigModel);
t.context.db = module.get(PrismaClient);
t.context.config = module.get(Config);
t.context.module = module;
});
let user: User;
let workspace: Workspace;
let sessionId: string;
let docId = 'doc1';
test.beforeEach(async t => {
await t.context.module.initTestingDB();
user = await t.context.user.create({
email: 'test@affine.pro',
});
workspace = await t.context.workspace.create(user.id);
sessionId = await t.context.copilotSession.create({
sessionId: randomUUID(),
workspaceId: workspace.id,
docId,
userId: user.id,
title: null,
promptName: 'prompt-name',
promptAction: null,
});
});
test.after(async t => {
await t.context.module.close();
});
test('should create a copilot context', async t => {
const { id: contextId } = await t.context.copilotContext.create(sessionId);
t.truthy(contextId);
const context = await t.context.copilotContext.get(contextId);
t.is(context?.id, contextId, 'should get context by id');
const config = await t.context.copilotContext.getConfig(contextId);
t.is(config?.workspaceId, workspace.id, 'should get context config');
const context1 = await t.context.copilotContext.getBySessionId(sessionId);
t.is(context1?.id, contextId, 'should get context by session id');
});
test('should get null for non-exist job', async t => {
const job = await t.context.copilotContext.get('non-exist');
t.snapshot(job, 'should return null for non-exist job');
});
test('should update context', async t => {
const { id: contextId } = await t.context.copilotContext.create(sessionId);
const config = (await t.context.copilotContext.getConfig(contextId))!;
t.assert(config, 'should get context config');
const doc = {
id: docId,
createdAt: Date.now(),
};
config.docs.push(doc);
await t.context.copilotContext.update(contextId, { config });
const config1 = await t.context.copilotContext.getConfig(contextId);
t.deepEqual(config1, config);
});
test('should insert embedding by doc id', async t => {
const { id: contextId } = await t.context.copilotContext.create(sessionId);
{
await t.context.copilotContext.insertFileEmbedding(contextId, 'file-id', [
{
index: 0,
content: 'content',
embedding: Array.from({ length: 1024 }, () => 1),
},
]);
{
const ret = await t.context.copilotContext.matchFileEmbedding(
Array.from({ length: 1024 }, () => 0.9),
contextId,
1,
1
);
t.snapshot(
cleanObject(ret, ['chunk', 'content', 'distance']),
'should match file embedding'
);
}
{
await t.context.copilotContext.deleteFileEmbedding(contextId, 'file-id');
const ret = await t.context.copilotContext.matchFileEmbedding(
Array.from({ length: 1024 }, () => 0.9),
contextId,
1,
1
);
t.snapshot(ret, 'should return empty array when embedding is deleted');
}
}
{
await t.context.db.snapshot.create({
data: {
workspaceId: workspace.id,
id: docId,
blob: Buffer.from([1, 1]),
state: Buffer.from([1, 1]),
updatedAt: new Date(),
createdAt: new Date(),
},
});
await t.context.copilotContext.insertWorkspaceEmbedding(
workspace.id,
docId,
[
{
index: 0,
content: 'content',
embedding: Array.from({ length: 1024 }, () => 1),
},
]
);
{
const ret = await t.context.copilotContext.listWorkspaceDocEmbedding(
workspace.id,
[docId]
);
t.true(
ret.includes(docId),
'should return doc id when embedding is inserted'
);
}
{
const ret = await t.context.copilotContext.matchWorkspaceEmbedding(
Array.from({ length: 1024 }, () => 0.9),
workspace.id,
1,
1
);
t.snapshot(
cleanObject(ret, ['chunk', 'content', 'distance']),
'should match workspace embedding'
);
}
{
await t.context.copilotWorkspace.updateIgnoredDocs(workspace.id, [docId]);
const ret = await t.context.copilotContext.matchWorkspaceEmbedding(
Array.from({ length: 1024 }, () => 0.9),
workspace.id,
1,
1
);
t.snapshot(ret, 'should return empty array when doc is ignored');
}
{
await t.context.copilotWorkspace.updateIgnoredDocs(
workspace.id,
undefined,
[docId]
);
const ret = await t.context.copilotContext.matchWorkspaceEmbedding(
Array.from({ length: 1024 }, () => 0.9),
workspace.id,
1,
1
);
t.snapshot(
cleanObject(ret, ['chunk', 'content', 'distance']),
'should return workspace embedding'
);
}
{
await t.context.copilotContext.deleteWorkspaceEmbedding(
workspace.id,
docId
);
const ret = await t.context.copilotContext.matchWorkspaceEmbedding(
Array.from({ length: 1024 }, () => 0.9),
workspace.id,
1,
1
);
t.snapshot(ret, 'should return empty array when embedding deleted');
}
}
});
test('should check embedding table', async t => {
{
const ret = await t.context.copilotContext.checkEmbeddingAvailable();
t.snapshot(ret, 'should return true when embedding table is available');
}
// {
// await t.context.db
// .$executeRaw`DROP TABLE IF EXISTS "ai_context_embeddings"`;
// const ret = await t.context.copilotContext.checkEmbeddingAvailable();
// t.false(ret, 'should return false when embedding table is not available');
// }
});
test('should merge doc status correctly', async t => {
const createDoc = (id: string, status?: string) => ({
id,
createdAt: Date.now(),
...(status && { status: status as any }),
});
const createDocWithEmbedding = async (docId: string) => {
await t.context.db.snapshot.create({
data: {
workspaceId: workspace.id,
id: docId,
blob: Buffer.from([1, 1]),
state: Buffer.from([1, 1]),
updatedAt: new Date(),
createdAt: new Date(),
},
});
await t.context.copilotContext.insertWorkspaceEmbedding(
workspace.id,
docId,
[
{
index: 0,
content: 'content',
embedding: Array.from({ length: 1024 }, () => 1),
},
]
);
};
const emptyResult = await t.context.copilotContext.mergeDocStatus(
workspace.id,
[]
);
t.deepEqual(emptyResult, []);
const basicDocs = [
createDoc('doc1'),
createDoc('doc2'),
createDoc('doc3', 'failed'),
createDoc('doc4', 'processing'),
];
const basicResult = await t.context.copilotContext.mergeDocStatus(
workspace.id,
basicDocs
);
t.snapshot(
basicResult.map(d => ({ id: d.id, status: d.status })),
'basic doc status merge'
);
{
await createDocWithEmbedding('doc5');
const mixedDocs = [
createDoc('doc5'),
createDoc('doc5', 'processing'),
createDoc('doc6'),
createDoc('doc6', 'failed'),
createDoc('doc7'),
];
const mixedResult = await t.context.copilotContext.mergeDocStatus(
workspace.id,
mixedDocs
);
t.snapshot(
mixedResult.map(d => ({ id: d.id, status: d.status })),
'mixed doc status merge'
);
const hasEmbeddingStub = Sinon.stub(
t.context.copilotContext,
'listWorkspaceDocEmbedding'
).resolves([]);
const stubResult = await t.context.copilotContext.mergeDocStatus(
workspace.id,
[createDoc('doc5')]
);
t.is(stubResult[0].status, ContextEmbedStatus.processing);
hasEmbeddingStub.restore();
}
{
const testCases = [
{
workspaceId: 'invalid-workspace',
docs: [{ id: 'doc1', createdAt: Date.now() }],
},
{
workspaceId: workspace.id,
docs: [{ id: 'doc1', createdAt: Date.now(), status: undefined as any }],
},
{
workspaceId: workspace.id,
docs: Array.from({ length: 100 }, (_, i) => ({
id: `doc-${i}`,
createdAt: Date.now() + i,
})),
},
];
const results = await Promise.all(
testCases.map(testCase =>
t.context.copilotContext.mergeDocStatus(
testCase.workspaceId,
testCase.docs
)
)
);
t.snapshot(
results.map((result, index) => ({
case: index,
length: result.length,
statuses: result.map(d => d.status),
})),
'edge cases results'
);
}
});
test('should handle concurrent mergeDocStatus calls', async t => {
await t.context.db.snapshot.create({
data: {
workspaceId: workspace.id,
id: 'concurrent-doc',
blob: Buffer.from([1, 1]),
state: Buffer.from([1, 1]),
updatedAt: new Date(),
createdAt: new Date(),
},
});
await t.context.copilotContext.insertWorkspaceEmbedding(
workspace.id,
'concurrent-doc',
[
{
index: 0,
content: 'content',
embedding: Array.from({ length: 1024 }, () => 1),
},
]
);
const concurrentDocs = [
[{ id: 'concurrent-doc', createdAt: Date.now() }],
[{ id: 'concurrent-doc', createdAt: Date.now() + 1000 }],
[{ id: 'non-existent-doc', createdAt: Date.now() }],
];
const results = await Promise.all(
concurrentDocs.map(docs =>
t.context.copilotContext.mergeDocStatus(workspace.id, docs)
)
);
t.snapshot(
results.map((result, index) => ({
call: index + 1,
status: result[0].status,
})),
'concurrent calls results'
);
});
packages/backend/server/src/__tests__/models/copilot-session.spec.ts
+97 -4
View File
@@ -895,13 +895,13 @@ test('should handle fork and session attachment operations', async t => {
t.snapshot(
{
attachPhase: {
afterAttach: {
docSessionCount: docSessionsAfterAttach.length,
bothSessionsPresent:
docSessionsAfterAttach.some(s => s.id === workspaceSessionId) &&
docSessionsAfterAttach.some(s => s.id === existingDocSessionId),
},
detachPhase: {
afterDetach: {
workspaceSessionExists: workspaceSessionsAfterDetach.some(
s => s.id === workspaceSessionId && !s.pinned
),
@@ -1000,27 +1000,120 @@ test('should cleanup empty sessions correctly', async t => {
test('should append durable message and account message cost', async t => {
const { copilotSession, db } = t.context;
const workspaceId = workspace.id;
if (!workspaceId) {
t.fail('Test workspace ID is missing');
return;
}
const { sessionId } = await createTestSession(t);
const artifact = await db.workspaceArtifact.create({
data: {
workspaceId,
contentHash: `test-${sessionId}`,
canonicalMediaType: 'text/plain',
sizeBytes: 5,
storageScope: 'copilot',
storageKey: `artifacts/${sessionId}`,
status: 'ready',
readyAt: new Date(),
},
});
const scopeSnapshot = {
version: 1,
resolvedAt: new Date().toISOString(),
selectors: [
{
kind: 'artifact' as const,
id: artifact.id,
source: 'message' as const,
},
],
requiredDocIds: [],
requiredArtifactIds: [artifact.id],
preferredSourceIds: [],
retrieval: {
mode: 'required' as const,
requiredDocIds: [],
requiredArtifactIds: [artifact.id],
preferredSourceIds: [],
},
};
const appended = await copilotSession.appendMessage({
sessionId,
userId: user.id,
message: {
role: 'user',
content: 'hello durable world',
attachments: [
{
kind: 'file_handle',
fileHandle: artifact.id,
mimeType: 'text/plain',
fileName: 'note.txt',
},
{
kind: 'file_handle',
fileHandle: artifact.id,
mimeType: 'text/plain',
fileName: 'duplicate-name.txt',
},
],
params: { foo: 'bar' },
scopeSnapshot,
createdAt: new Date(),
},
focus: {
selectors: [{ kind: 'document', id: 'doc-1', source: 'focus' }],
},
artifacts: [
{
artifactId: artifact.id,
role: 'attachment',
displayName: 'note.txt',
},
{
artifactId: artifact.id,
role: 'attachment',
displayName: 'duplicate-name.txt',
},
],
});
const afterAppend = await db.aiSession.findUniqueOrThrow({
where: { id: sessionId },
select: { messageCost: true },
select: { messageCost: true, focus: true },
});
t.truthy(appended.id);
const messageId = appended.id;
if (!messageId) {
t.fail('Appended message ID is missing');
return;
}
t.is(afterAppend.messageCost, 1);
t.is(appended.attachments?.length, 2);
t.deepEqual(appended.params, { foo: 'bar' });
t.deepEqual(appended.scopeSnapshot, scopeSnapshot);
t.deepEqual(afterAppend.focus, {
selectors: [{ kind: 'document', id: 'doc-1', source: 'focus' }],
});
const artifactReference = await db.aiMessageArtifact.findUniqueOrThrow({
where: {
messageId_artifactId_role: {
messageId,
artifactId: artifact.id,
role: 'attachment',
},
},
});
t.is(artifactReference.workspaceId, workspaceId);
t.is(artifactReference.displayName, 'note.txt');
t.is(
await db.aiMessageArtifact.count({
where: { messageId, artifactId: artifact.id, role: 'attachment' },
}),
1
);
const appendedBare = await copilotSession.appendMessage({
sessionId,
packages/backend/server/src/__tests__/models/copilot-workspace.spec.ts
+230 -412
View File
@@ -1,26 +1,23 @@
import { randomUUID } from 'node:crypto';
import { createHash } from 'node:crypto';
import { PrismaClient, User, Workspace } from '@prisma/client';
import ava, { TestFn } from 'ava';
import { Config } from '../../base';
import { CopilotContextModel } from '../../models/copilot-context';
import { BackendRuntimeProvider } from '../../core/backend-runtime';
import { WorkspaceBlobStorage } from '../../core/storage';
import { CopilotWorkspaceConfigModel } from '../../models/copilot-workspace';
import { DocModel } from '../../models/doc';
import { UserModel } from '../../models/user';
import { WorkspaceModel } from '../../models/workspace';
import { createTestingModule, type TestingModule } from '../utils';
import { cleanObject } from '../utils/copilot';
interface Context {
config: Config;
module: TestingModule;
db: PrismaClient;
doc: DocModel;
user: UserModel;
workspace: WorkspaceModel;
copilotContext: CopilotContextModel;
copilotWorkspace: CopilotWorkspaceConfigModel;
runtime: BackendRuntimeProvider;
db: PrismaClient;
storage: WorkspaceBlobStorage;
}
const test = ava as TestFn<Context>;
@@ -29,24 +26,19 @@ test.before(async t => {
const module = await createTestingModule();
t.context.user = module.get(UserModel);
t.context.workspace = module.get(WorkspaceModel);
t.context.copilotContext = module.get(CopilotContextModel);
t.context.copilotWorkspace = module.get(CopilotWorkspaceConfigModel);
t.context.runtime = module.get(BackendRuntimeProvider);
t.context.db = module.get(PrismaClient);
t.context.doc = module.get(DocModel);
t.context.config = module.get(Config);
t.context.storage = module.get(WorkspaceBlobStorage);
t.context.module = module;
});
let user: User;
let workspace: Workspace;
let docId = 'doc1';
test.beforeEach(async t => {
await t.context.module.initTestingDB();
user = await t.context.user.create({
email: 'test@affine.pro',
});
user = await t.context.user.create({ email: 'test@affine.pro' });
workspace = await t.context.workspace.create(user.id);
});
@@ -54,419 +46,245 @@ test.after(async t => {
await t.context.module.close();
});
test('should manage copilot workspace ignored docs', async t => {
const ignoredDocs = await t.context.copilotWorkspace.listIgnoredDocs(
workspace.id
test('should manage workspace ignored documents', async t => {
t.is(await t.context.copilotWorkspace.countIgnoredDocs(workspace.id), 0);
t.is(
await t.context.copilotWorkspace.updateIgnoredDocs(workspace.id, ['doc1']),
1
);
t.deepEqual(ignoredDocs, []);
{
const count = await t.context.copilotWorkspace.updateIgnoredDocs(
workspace.id,
[docId]
);
t.snapshot(count, 'should add ignored doc');
const ret = await t.context.copilotWorkspace.listIgnoredDocs(workspace.id);
t.snapshot(cleanObject(ret), 'should return added doc');
const check = await t.context.copilotWorkspace.checkIgnoredDocs(
workspace.id,
[docId]
);
t.snapshot(check, 'should return ignored docs in workspace');
}
{
const count = await t.context.copilotWorkspace.updateIgnoredDocs(
workspace.id,
[docId]
);
t.snapshot(count, 'should not change if ignored doc exists');
const ret = await t.context.copilotWorkspace.listIgnoredDocs(workspace.id);
t.snapshot(cleanObject(ret), 'should not add ignored doc again');
}
{
const count = await t.context.copilotWorkspace.updateIgnoredDocs(
workspace.id,
['new_doc']
);
t.snapshot(count, 'should add new ignored doc');
const ret = await t.context.copilotWorkspace.listIgnoredDocs(workspace.id);
t.snapshot(cleanObject(ret), 'should add ignored doc');
}
{
t.is(
await t.context.copilotWorkspace.updateIgnoredDocs(workspace.id, ['doc1']),
0
);
t.is(
await t.context.copilotWorkspace.updateIgnoredDocs(workspace.id, ['doc2']),
1
);
t.is(await t.context.copilotWorkspace.countIgnoredDocs(workspace.id), 2);
const firstPage = await t.context.copilotWorkspace.listIgnoredDocs(
workspace.id,
{ offset: 0, first: 1 }
);
t.is(firstPage.length, 1);
t.true(['doc1', 'doc2'].includes(firstPage[0].docId));
t.deepEqual(
await t.context.copilotWorkspace.checkIgnoredDocs(workspace.id, [
'doc1',
'doc2',
]),
['doc1', 'doc2']
);
t.is(
await t.context.copilotWorkspace.updateIgnoredDocs(
workspace.id,
undefined,
[docId]
);
const ret = await t.context.copilotWorkspace.listIgnoredDocs(workspace.id);
t.snapshot(cleanObject(ret), 'should remove ignored doc');
}
[],
['doc1', 'doc2']
),
2
);
t.is(await t.context.copilotWorkspace.countIgnoredDocs(workspace.id), 0);
});
test('should insert and search embedding', async t => {
{
const { fileId } = await t.context.copilotWorkspace.addFile(workspace.id, {
fileName: 'file1',
blobId: 'blob1',
test('workspace artifacts deduplicate bytes and remain workspace isolated', async t => {
const body = Buffer.from('shared artifact');
const first = await t.context.runtime.putWorkspaceArtifact(
{
workspaceId: workspace.id,
mimeType: 'text/plain',
size: 1,
});
await t.context.copilotWorkspace.insertFileEmbeddings(
workspace.id,
fileId,
[
{
index: 0,
content: 'content',
embedding: Array.from({ length: 1024 }, () => 1),
},
]
);
displayName: 'first.txt',
fileName: 'first.txt',
libraryOwned: false,
},
body
);
const repeated = await t.context.runtime.putWorkspaceArtifact(
{
const ret = await t.context.copilotWorkspace.matchFileEmbedding(
workspace.id,
Array.from({ length: 1024 }, () => 0.9),
1,
1
);
t.snapshot(
cleanObject(ret, ['fileId']),
'should match workspace file embedding'
);
}
}
workspaceId: workspace.id,
mimeType: 'text/plain',
displayName: 'repeated.txt',
fileName: 'repeated.txt',
libraryOwned: true,
},
body
);
t.is(repeated.id, first.id);
t.is(repeated.displayName, 'repeated.txt');
t.is(repeated.fileName, 'first.txt');
t.true(repeated.libraryOwned);
{
await t.context.db.blob.create({
data: {
workspaceId: workspace.id,
key: 'blob-test',
mime: 'text/plain',
size: 1,
},
});
const blobId = 'blob-test';
await t.context.copilotWorkspace.insertBlobEmbeddings(
const unnamed = await t.context.runtime.putWorkspaceArtifact(
{
workspaceId: workspace.id,
mimeType: 'application/octet-stream',
libraryOwned: false,
},
Buffer.from('unnamed artifact')
);
await t.throwsAsync(
t.context.runtime.setArtifactLibraryOwned(workspace.id, unnamed.id, true),
{ message: 'artifact_library_display_name_required' }
);
await t.throwsAsync(
t.context.runtime.setArtifactLibraryOwned(
workspace.id,
'6ba7b811-9dad-11d1-80b4-00c04fd430c8',
false
),
{ message: 'artifact_not_found' }
);
const blobId = createHash('sha256').update(body).digest('base64url');
await t.context.storage.put(workspace.id, blobId, body);
await t.throwsAsync(
t.context.runtime.ensureWorkspaceBlobArtifact({
workspaceId: workspace.id,
blobId,
[
{
index: 0,
content: 'blob content',
embedding: Array.from({ length: 1024 }, () => 1),
},
]
);
mimeType: 'text/plain',
libraryOwned: true,
}),
{ message: 'artifact_library_display_name_required' }
);
await t.context.db.workspaceArtifact.update({
where: { id: first.id },
data: {
status: 'reserving',
reservationExpiresAt: new Date(Date.now() + 60_000),
},
});
const aliased = await t.context.runtime.ensureWorkspaceBlobArtifact({
workspaceId: workspace.id,
blobId,
mimeType: 'text/plain',
libraryOwned: false,
});
t.is(aliased.id, first.id);
t.is(aliased.status, 'ready');
t.is(aliased.storageScope, 'copilot');
const otherWorkspace = await t.context.workspace.create(user.id);
const isolated = await t.context.runtime.putWorkspaceArtifact(
{
const ret = await t.context.copilotWorkspace.matchBlobEmbedding(
workspace.id,
Array.from({ length: 1024 }, () => 0.9),
1,
1
);
t.snapshot(cleanObject(ret), 'should match workspace blob embedding');
}
workspaceId: otherWorkspace.id,
mimeType: 'text/plain',
fileName: 'isolated.txt',
libraryOwned: false,
},
body
);
t.not(isolated.id, first.id);
t.is(isolated.contentHash, first.contentHash);
t.is(
await t.context.db.workspaceArtifact.count({
where: { contentHash: first.contentHash },
}),
2
);
await t.context.copilotWorkspace.removeBlob(workspace.id, blobId);
const session = await t.context.db.aiSession.create({
data: {
userId: user.id,
workspaceId: otherWorkspace.id,
promptName: 'Chat With AFFiNE AI',
},
});
const message = await t.context.db.aiSessionMessage.create({
data: { sessionId: session.id, role: 'user', content: 'attachment' },
});
await t.context.db.aiMessageArtifact.create({
data: {
messageId: message.id,
workspaceId: otherWorkspace.id,
artifactId: isolated.id,
role: 'attachment',
},
});
await t.context.workspace.delete(otherWorkspace.id);
t.is(
await t.context.db.aiMessageArtifact.count({
where: { artifactId: isolated.id },
}),
0
);
await t.context.runtime.setArtifactLibraryOwned(
workspace.id,
first.id,
false
);
await t.context.db.$executeRaw`UPDATE workspace_artifacts
SET created_at='2026-01-01T00:00:00.000Z', updated_at='2026-01-01T00:00:00.000Z'
WHERE id=${first.id}::uuid`;
const reused = await t.context.runtime.putWorkspaceArtifact(
{
const ret = await t.context.copilotWorkspace.matchBlobEmbedding(
workspace.id,
Array.from({ length: 1024 }, () => 0.9),
1,
1
);
t.deepEqual(ret, [], 'should not match after removal');
}
}
{
const docId = randomUUID();
await t.context.doc.upsert({
spaceId: workspace.id,
docId,
blob: Uint8Array.from([1, 2, 3]),
timestamp: Date.now(),
editorId: user.id,
});
const toBeEmbedDocIds = await t.context.copilotWorkspace.findDocsToEmbed(
workspace.id
);
t.snapshot(toBeEmbedDocIds.length, 'should find docs to embed');
await t.context.copilotContext.insertWorkspaceEmbedding(
workspace.id,
docId,
[
{
index: 0,
content: 'content',
embedding: Array.from({ length: 1024 }, () => 1),
},
]
);
const afterInsertEmbedding =
await t.context.copilotWorkspace.findDocsToEmbed(workspace.id);
t.snapshot(afterInsertEmbedding.length, 'should not find docs to embed');
}
{
const docId = randomUUID();
await t.context.doc.upsert({
spaceId: workspace.id,
docId,
blob: Uint8Array.from([1, 2, 3]),
timestamp: Date.now(),
editorId: user.id,
});
const toBeEmbedDocIds = await t.context.copilotWorkspace.findDocsToEmbed(
workspace.id
);
t.snapshot(toBeEmbedDocIds.length, 'should find docs to embed');
await t.context.copilotWorkspace.updateIgnoredDocs(workspace.id, [docId]);
const afterAddIgnoreDocs = await t.context.copilotWorkspace.findDocsToEmbed(
workspace.id
);
t.snapshot(afterAddIgnoreDocs.length, 'should not find docs to embed');
}
{
const docId = `foo$bar`;
await t.context.doc.upsert({
spaceId: workspace.id,
docId: docId,
blob: Uint8Array.from([1, 2, 3]),
timestamp: Date.now(),
editorId: user.id,
});
const results = await t.context.copilotWorkspace.findDocsToEmbed(
workspace.id
);
t.false(results.includes(docId), 'docs containing `$` should be excluded');
}
{
const docId = 'empty_doc';
await t.context.doc.upsert({
spaceId: workspace.id,
docId: docId,
blob: Uint8Array.from([0, 0]),
timestamp: Date.now(),
editorId: user.id,
});
const results = await t.context.copilotWorkspace.findDocsToEmbed(
workspace.id
);
t.false(results.includes(docId), 'empty documents should be excluded');
}
});
test('should check need to be embedded', async t => {
const docId = randomUUID();
await t.context.doc.upsert({
spaceId: workspace.id,
docId,
blob: Uint8Array.from([1, 2, 3]),
timestamp: Date.now(),
editorId: user.id,
workspaceId: workspace.id,
mimeType: 'text/plain',
displayName: 'reused.txt',
fileName: 'reused.txt',
libraryOwned: false,
},
body
);
t.is(reused.id, first.id);
t.is(await t.context.runtime.cleanupUnreferencedArtifacts(1), 0);
await t.context.db.$executeRaw`UPDATE workspace_artifacts
SET updated_at='2026-01-01T00:00:00.000Z'
WHERE id=${first.id}::uuid`;
const retainedSession = await t.context.db.aiSession.create({
data: {
userId: user.id,
workspaceId: workspace.id,
promptName: 'Chat With AFFiNE AI',
},
});
{
let needsEmbedding = await t.context.copilotWorkspace.checkDocNeedEmbedded(
workspace.id,
docId
);
t.snapshot(
needsEmbedding,
'document with no embedding should need embedding'
);
}
{
await t.context.copilotContext.insertWorkspaceEmbedding(
workspace.id,
docId,
[
{
index: 0,
content: 'content',
embedding: Array.from({ length: 1024 }, () => 1),
const retainedMessage = await t.context.db.aiSessionMessage.create({
data: {
sessionId: retainedSession.id,
role: 'user',
content: 'retained attachment',
artifacts: {
create: {
workspaceId: workspace.id,
artifactId: first.id,
role: 'attachment',
},
]
);
let needsEmbedding = await t.context.copilotWorkspace.checkDocNeedEmbedded(
workspace.id,
docId
);
t.snapshot(
needsEmbedding,
'document with recent embedding should not need embedding'
);
}
{
await t.context.doc.upsert({
spaceId: workspace.id,
docId,
blob: Uint8Array.from([4, 5, 6]),
timestamp: Date.now() + 1000, // Ensure timestamp is later
editorId: user.id,
});
// simulate an old embedding
const oldEmbeddingTime = new Date(Date.now() - 25 * 60 * 1000);
await t.context.db.aiWorkspaceEmbedding.updateMany({
where: { workspaceId: workspace.id, docId },
data: { updatedAt: oldEmbeddingTime },
});
let needsEmbedding = await t.context.copilotWorkspace.checkDocNeedEmbedded(
workspace.id,
docId
);
t.snapshot(
needsEmbedding,
'document updated after embedding and older-than-10m should need embedding'
);
}
{
// only time passed (>10m since last embedding) but no doc updates => should NOT re-embed
const baseNow = Date.now();
const docId2 = randomUUID();
const t0 = baseNow - 30 * 60 * 1000; // snapshot updated 30 minutes ago
const t1 = baseNow - 25 * 60 * 1000; // embedding updated 25 minutes ago
await t.context.doc.upsert({
spaceId: workspace.id,
docId: docId2,
blob: Uint8Array.from([1, 2, 3]),
timestamp: t0,
editorId: user.id,
});
await t.context.copilotContext.insertWorkspaceEmbedding(
workspace.id,
docId2,
[
{
index: 0,
content: 'content2',
embedding: Array.from({ length: 1024 }, () => 1),
},
]
);
await t.context.db.aiWorkspaceEmbedding.updateMany({
where: { workspaceId: workspace.id, docId: docId2 },
data: { updatedAt: new Date(t1) },
});
let needsEmbedding = await t.context.copilotWorkspace.checkDocNeedEmbedded(
workspace.id,
docId2
);
t.snapshot(
needsEmbedding,
'should not need embedding when only 10-minute window passed without updates'
);
const t2 = baseNow - 5 * 60 * 1000; // doc updated 5 minutes ago
await t.context.doc.upsert({
spaceId: workspace.id,
docId: docId2,
blob: Uint8Array.from([7, 8, 9]),
timestamp: t2,
editorId: user.id,
});
needsEmbedding = await t.context.copilotWorkspace.checkDocNeedEmbedded(
workspace.id,
docId2
);
t.snapshot(
needsEmbedding,
'should need embedding when doc updated and last embedding older than 10 minutes'
);
}
// --- new cases end ---
});
test('should check embedding table', async t => {
{
const ret = await t.context.copilotWorkspace.checkEmbeddingAvailable();
t.true(ret, 'should return true when embedding table is available');
}
// {
// await t.context.db
// .$executeRaw`DROP TABLE IF EXISTS "ai_workspace_file_embeddings"`;
// const ret = await t.context.copilotWorkspace.checkEmbeddingAvailable();
// t.false(ret, 'should return false when embedding table is not available');
// }
});
test('should filter outdated doc id style in embedding status', async t => {
const docId = randomUUID();
const outdatedDocId = `${workspace.id}:space:${docId}`;
await t.context.doc.upsert({
spaceId: workspace.id,
docId,
blob: Uint8Array.from([1, 2, 3]),
timestamp: Date.now(),
editorId: user.id,
},
},
});
await t.context.doc.upsert({
spaceId: workspace.id,
docId: outdatedDocId,
blob: Uint8Array.from([1, 2, 3]),
timestamp: Date.now(),
editorId: user.id,
t.is(await t.context.runtime.cleanupUnreferencedArtifacts(1), 0);
await t.context.db.aiSessionMessage.delete({
where: { id: retainedMessage.id },
});
t.is(await t.context.runtime.cleanupUnreferencedArtifacts(1), 1);
const [source] = await t.context.db.$queryRaw<
{ deletedAt: Date | null }[]
>`SELECT deleted_at AS "deletedAt" FROM embedding_sources
WHERE workspace_id=${workspace.id} AND source_kind='artifact' AND source_key=${first.id}`;
t.truthy(source?.deletedAt);
t.is(
await t.context.db.workspaceArtifact.count({ where: { id: first.id } }),
0
);
{
const status = await t.context.copilotWorkspace.getEmbeddingStatus(
workspace.id
);
t.snapshot(status, 'should include modern doc format');
}
{
await t.context.copilotContext.insertWorkspaceEmbedding(
workspace.id,
docId,
[
{
index: 0,
content: 'content',
embedding: Array.from({ length: 1024 }, () => 1),
},
]
);
const status = await t.context.copilotWorkspace.getEmbeddingStatus(
workspace.id
);
t.snapshot(status, 'should count docs after filtering outdated');
}
const deletingBody = Buffer.from('cleanup retry');
const deletingBlobId = createHash('sha256')
.update(deletingBody)
.digest('base64url');
await t.context.storage.put(workspace.id, deletingBlobId, deletingBody);
const deleting = await t.context.runtime.ensureWorkspaceBlobArtifact({
workspaceId: workspace.id,
blobId: deletingBlobId,
mimeType: 'text/plain',
libraryOwned: false,
});
t.is(deleting.storageScope, 'blob');
await t.context.db.workspaceArtifact.update({
where: { id: deleting.id },
data: { status: 'deleting' },
});
await t.context.storage.delete(workspace.id, deletingBlobId, true);
t.is(await t.context.runtime.cleanupUnreferencedArtifacts(1), 1);
t.is(
await t.context.db.workspaceArtifact.count({ where: { id: deleting.id } }),
0
);
});
packages/backend/server/src/__tests__/sync/gateway.spec.ts
+31
View File
@@ -721,6 +721,16 @@ test('workspace sync delete-doc should enforce doc permissions', async t => {
);
t.true(error.message.includes('Doc.Delete'));
const userdataError = getErrorResponse(
t,
await emitWithAck(socket, 'space:delete-doc', {
spaceType: 'workspace',
spaceId: workspace.id,
docId: `userdata$${owner.id}$${workspace.id}$docIntegrationRef`,
})
);
t.is(userdataError.name, 'SPACE_ACCESS_DENIED');
const ownerJoin = unwrapResponse(
t,
await emitWithAck<{ clientId: string; success: boolean }>(
@@ -805,6 +815,16 @@ test('workspace sync load-doc should enforce doc read permissions', async t => {
})
);
t.true(error.message.includes('Doc.Read'));
const userdataError = getErrorResponse(
t,
await emitWithAck(socket, 'space:load-doc', {
spaceType: 'workspace',
spaceId: workspace.id,
docId: `userdata$${owner.id}$${workspace.id}$favorite`,
})
);
t.is(userdataError.name, 'SPACE_ACCESS_DENIED');
} finally {
socket.disconnect();
}
@@ -869,6 +889,17 @@ test('workspace sync push-doc-update should enforce doc update permissions', asy
);
t.true(error.message.includes('Doc.Update'));
const userdataError = getErrorResponse(
t,
await emitWithAck(socket, 'space:push-doc-update', {
spaceType: 'workspace',
spaceId: workspace.id,
docId: `userdata$${owner.id}$${workspace.id}$settings`,
update: createYjsUpdateBase64(),
})
);
t.is(userdataError.name, 'SPACE_ACCESS_DENIED');
const updates = await db.update.count({
where: {
workspaceId: workspace.id,
packages/backend/server/src/__tests__/utils/blobs.ts
+3 -1
View File
@@ -1,3 +1,5 @@
import { createHash } from 'node:crypto';
import { type Blob } from '@prisma/client';
import { TestingApp } from './testing-app';
@@ -104,7 +106,7 @@ export async function setBlob(
.attach(
'0',
buffer,
`blob-${Math.random().toString(16).substring(2, 10)}.data`
createHash('sha256').update(buffer).digest('base64url')
)
.expect(200);
packages/backend/server/src/__tests__/utils/copilot.ts
-238
View File
@@ -1,26 +1,14 @@
import {
addContextCategoryMutation,
addContextDocMutation,
addContextFileMutation,
ContextCategories as GraphQLContextCategories,
createCopilotContextMutation,
createCopilotMessageMutation,
createCopilotSessionMutation,
forkCopilotSessionMutation,
getCopilotSessionQuery,
getTranscriptTaskQuery,
listContextObjectQuery,
listContextQuery,
matchFilesQuery,
matchWorkspaceDocsQuery,
removeContextDocMutation,
removeContextFileMutation,
settleTranscriptTaskMutation,
submitTranscriptTaskMutation,
updateCopilotSessionMutation,
} from '@affine/graphql';
import { ContextCategories } from '../../models';
import { TestingApp } from './testing-app';
export const cleanObject = (
@@ -132,232 +120,6 @@ export async function forkCopilotSession(
return res.forkCopilotSession;
}
export async function createCopilotContext(
app: TestingApp,
workspaceId: string,
sessionId: string
): Promise<string> {
const res = await app.gql({
query: createCopilotContextMutation,
variables: { workspaceId, sessionId },
});
return res.createCopilotContext;
}
export async function matchFiles(
app: TestingApp,
contextId: string,
content: string,
limit: number
): Promise<
| {
fileId: string;
chunk: number;
content: string;
distance: number | null;
}[]
| undefined
> {
const res = await app.gql({
query: matchFilesQuery,
variables: { contextId, content, limit, threshold: 1 },
});
return res.currentUser?.copilot?.contexts?.[0]?.matchFiles;
}
export async function matchWorkspaceDocs(
app: TestingApp,
contextId: string,
content: string,
limit: number
): Promise<
| {
docId: string;
chunk: number;
content: string;
distance: number | null;
}[]
| undefined
> {
const res = await app.gql({
query: matchWorkspaceDocsQuery,
variables: { contextId, content, limit, threshold: 1 },
});
return res.currentUser?.copilot?.contexts?.[0]?.matchWorkspaceDocs;
}
export async function listContext(
app: TestingApp,
workspaceId: string,
sessionId: string
): Promise<
{
id: string;
workspaceId: string;
}[]
> {
const res = await app.gql({
query: listContextQuery,
variables: { workspaceId, sessionId },
});
return (res.currentUser?.copilot?.contexts || []).filter(
(context): context is { id: string; workspaceId: string } => !!context.id
);
}
export async function addContextFile(
app: TestingApp,
contextId: string,
fileName: string,
content: Buffer
): Promise<{ id: string }> {
const res = await app.gql({
query: addContextFileMutation,
variables: {
content: new File([content], fileName, {
type: 'application/octet-stream',
}),
options: { contextId },
},
});
return res.addContextFile;
}
export async function removeContextFile(
app: TestingApp,
contextId: string,
fileId: string
): Promise<boolean> {
const res = await app.gql({
query: removeContextFileMutation,
variables: { options: { contextId, fileId } },
});
return res.removeContextFile;
}
export async function addContextDoc(
app: TestingApp,
contextId: string,
docId: string
): Promise<{ id: string }[]> {
const res = await app.gql({
query: addContextDocMutation,
variables: { options: { contextId, docId } },
});
return [res.addContextDoc];
}
export async function addContextCategory(
app: TestingApp,
contextId: string,
type: ContextCategories,
categoryId: string,
docs: string[]
): Promise<{ type: string; id: string; docs: { id: string }[] }> {
const graphqlType =
type === ContextCategories.Collection
? GraphQLContextCategories.Collection
: GraphQLContextCategories.Tag;
const res = await app.gql({
query: addContextCategoryMutation,
variables: { options: { contextId, type: graphqlType, categoryId, docs } },
});
return res.addContextCategory;
}
export async function removeContextDoc(
app: TestingApp,
contextId: string,
docId: string
): Promise<boolean> {
const res = await app.gql({
query: removeContextDocMutation,
variables: { options: { contextId, docId } },
});
return res.removeContextDoc;
}
export async function listContextDocAndFiles(
app: TestingApp,
workspaceId: string,
sessionId: string,
contextId: string
): Promise<
| {
docs: {
id: string;
status: string | null;
createdAt: number;
}[];
files: {
id: string;
name: string;
blobId: string;
chunkSize: number;
status: string;
error: string | null;
createdAt: number;
}[];
}
| undefined
> {
const res = await app.gql({
query: listContextObjectQuery,
variables: { workspaceId, sessionId, contextId },
});
const context = res.currentUser?.copilot?.contexts?.[0];
if (!context) {
return undefined;
}
return {
docs: context.docs,
files: context.files.map(({ mimeType: _mimeType, ...file }) => file),
};
}
export async function listContextCategories(
app: TestingApp,
workspaceId: string,
sessionId: string,
contextId: string
): Promise<
| {
collections: {
type: string;
id: string;
docs: {
id: string;
status: string | null;
createdAt: number;
}[];
}[];
}
| undefined
> {
const res = await app.gql({
query: listContextObjectQuery,
variables: { workspaceId, sessionId, contextId },
});
const context = res.currentUser?.copilot?.contexts?.[0];
if (!context) {
return undefined;
}
return { collections: context.collections };
}
export async function submitTranscriptTask(
app: TestingApp,
workspaceId: string,
packages/backend/server/src/__tests__/utils/runtime-config.ts
+40
View File
@@ -0,0 +1,40 @@
import { generateKeyPairSync } from 'node:crypto';
import { mkdtemp, rm, writeFile } from 'node:fs/promises';
import { tmpdir } from 'node:os';
import { join } from 'node:path';
const { privateKey } = generateKeyPairSync('ec', { namedCurve: 'P-256' });
const testPrivateKey = privateKey
.export({ format: 'pem', type: 'pkcs8' })
.toString();
export async function createTestRuntimeConfig(databaseUrl: string) {
const directory = await mkdtemp(join(tmpdir(), 'affine-server-test-'));
const storagePath = join(directory, 'storage');
const storage = (bucket: string) => ({
provider: 'assetpack',
bucket,
config: { path: storagePath },
});
const configPath = join(directory, 'config.json');
await writeFile(
configPath,
JSON.stringify({
crypto: { privateKey: testPrivateKey },
db: { datasourceUrl: databaseUrl },
storages: {
'avatar.storage': storage('avatars'),
'blob.storage': storage('blobs'),
},
copilot: {
enabled: true,
storage: storage('copilot'),
},
})
);
return {
configPath,
storagePath,
cleanup: () => rm(directory, { recursive: true, force: true }),
};
}
packages/backend/server/src/__tests__/utils/testing-module.ts
+36 -8
View File
@@ -13,6 +13,7 @@ import { AFFiNELogger, ConfigFactory, JobModule, JobQueue } from '../../base';
import { GqlModule } from '../../base/graphql';
import { ServerConfigModule } from '../../core';
import { AuthGuard, AuthModule } from '../../core/auth';
import { BACKEND_RUNTIME_CONFIG_PATHS } from '../../core/backend-runtime';
import { Mailer, MailModule } from '../../core/mail';
import { ModelsModule } from '../../models';
// for jsdoc inference
@@ -20,6 +21,7 @@ import { ModelsModule } from '../../models';
import type { createModule } from '../create-module';
import { createFactory, MockJobModule, MockJobQueue } from '../mocks';
import { MockMailer } from '../mocks/mailer.mock';
import { createTestRuntimeConfig } from './runtime-config';
import { initTestingDB, TEST_LOG_LEVEL } from './utils';
interface TestingModuleMetadata extends ModuleMetadata {
@@ -73,6 +75,9 @@ export async function createTestingModule(
moduleDef: TestingModuleMetadata = {},
autoInitialize = true
): Promise<TestingModule> {
const runtimeConfig = await createTestRuntimeConfig(
new ConfigFactory().config.db.datasourceUrl
);
// setting up
let imports = moduleDef.imports ?? [buildAppModule(globalThis.env)];
imports =
@@ -104,25 +109,34 @@ export async function createTestingModule(
builder.overrideProvider(Mailer).useClass(MockMailer);
builder.overrideProvider(JobQueue).useClass(MockJobQueue);
builder
.overrideProvider(BACKEND_RUNTIME_CONFIG_PATHS)
.useValue([runtimeConfig.configPath]);
if (moduleDef.tapModule) {
moduleDef.tapModule(builder);
}
const module = await builder.compile();
let module: BaseTestingModule;
try {
module = await builder.compile();
} catch (error) {
await runtimeConfig.cleanup();
throw error;
}
module.get(ConfigFactory).override({
storages: {
avatar: {
storage: {
provider: 'assetpack',
bucket: 'avatars',
config: { path: '/tmp/affine-test-storage' },
config: { path: runtimeConfig.storagePath },
},
},
blob: {
storage: {
provider: 'assetpack',
bucket: 'blobs',
config: { path: '/tmp/affine-test-storage' },
config: { path: runtimeConfig.storagePath },
},
},
},
@@ -131,7 +145,7 @@ export async function createTestingModule(
storage: {
provider: 'assetpack',
bucket: 'copilot',
config: { path: '/tmp/affine-test-storage' },
config: { path: runtimeConfig.storagePath },
},
},
});
@@ -146,9 +160,18 @@ export async function createTestingModule(
module.get(PrismaClient, { strict: false })
);
testingModule[Symbol.asyncDispose] = async () => {
await module.close();
const close = testingModule.close.bind(testingModule);
let closePromise: Promise<void> | undefined;
testingModule.close = () => {
return (closePromise ??= (async () => {
try {
await close();
} finally {
await runtimeConfig.cleanup();
}
})());
};
testingModule[Symbol.asyncDispose] = () => testingModule.close();
testingModule.mails = module.get(Mailer, { strict: false }) as MockMailer;
testingModule.queue = module.get(JobQueue, { strict: false }) as MockJobQueue;
@@ -160,8 +183,13 @@ export async function createTestingModule(
module.useLogger(logger);
if (autoInitialize) {
await testingModule.initTestingDB();
await testingModule.init();
try {
await testingModule.initTestingDB();
await testingModule.init();
} catch (error) {
await testingModule.close();
throw error;
}
}
return testingModule;
}
packages/backend/server/src/__tests__/workspace/blobs.e2e.ts
+9 -2
View File
@@ -273,7 +273,7 @@ test('should create pending blob upload with graphql fallback', async t => {
await app.signupV1('u1@affine.pro');
const workspace = await createWorkspace(app);
const key = `upload-${Math.random().toString(16).slice(2, 8)}`;
const key = sha256Base64urlWithPadding(Buffer.from('pending-upload'));
const size = 4;
const mime = 'text/plain';
@@ -351,7 +351,14 @@ test('should reject multipart upload part url on fs provider', async t => {
const workspace = await createWorkspace(app);
await t.throwsAsync(
() => getBlobUploadPartUrl(app, workspace.id, 'blob-key', 'upload', 1),
() =>
getBlobUploadPartUrl(
app,
workspace.id,
sha256Base64urlWithPadding(Buffer.from('blob-key')),
'upload',
1
),
{
message: 'Multipart upload is not supported',
}
packages/backend/server/src/base/config/__tests__/config.spec.ts
+16
View File
@@ -16,6 +16,12 @@ test('should create config', t => {
t.is(typeof config.auth.passwordRequirements.max, 'number');
t.is(typeof config.job.queue, 'object');
t.deepEqual(config.copilot.byok.allowedProviders, [
'openai',
'anthropic',
'gemini',
'fal',
]);
});
test('should override config', async t => {
@@ -89,6 +95,16 @@ test('should validate config', t => {
error.message,
'Invalid app config for module `auth` with key `passwordRequirements`. Minimum length of password must be less than maximum length.'
);
const [nativeError] = config.validate([
{
module: 'copilot',
key: 'byok.allowedProviders',
value: ['openai', 'openai'],
},
])!;
t.true(nativeError instanceof InvalidAppConfig);
t.regex(nativeError.message, /supported and unique/);
});
test('should override correctly', t => {
packages/backend/server/src/base/config/index.ts
+6 -1
View File
@@ -26,4 +26,9 @@ export class ConfigModule {
}
export { Config, ConfigFactory };
export { defineModuleConfig, type JSONSchema } from './register';
export {
defineModuleConfig,
defineNativeModuleConfig,
type JSONSchema,
type NativeAppConfigDescriptor,
} from './register';
packages/backend/server/src/base/config/register.ts
+100 -33
View File
@@ -8,22 +8,35 @@ import { z } from 'zod';
import { type EnvConfigType, parseEnvValue } from './env';
import { AppConfigByPath } from './types';
export type JSONSchema = { description?: string } & (
| { type?: undefined; oneOf?: JSONSchema[] }
| {
type: 'string' | 'number' | 'boolean';
enum?: string[];
}
| {
type: 'array';
items?: JSONSchema;
}
| {
type: 'object';
properties?: Record<string, JSONSchema>;
required?: string[];
}
);
export type JSONSchema = {
$id?: string;
$ref?: string;
$schema?: string;
additionalProperties?: boolean | JSONSchema;
allOf?: JSONSchema[];
anyOf?: JSONSchema[];
definitions?: Record<string, JSONSchema>;
description?: string;
default?: unknown;
enum?: unknown[];
format?: string;
items?: JSONSchema;
minItems?: number;
minLength?: number;
oneOf?: JSONSchema[];
pattern?: string;
properties?: Record<string, JSONSchema>;
required?: string[];
title?: string;
type?:
| 'string'
| 'number'
| 'boolean'
| 'array'
| 'object'
| 'null'
| Array<'string' | 'number' | 'boolean' | 'array' | 'object' | 'null'>;
};
type ConfigType = EnvConfigType | 'array' | 'object' | 'any';
export type ConfigDescriptor<T> = {
@@ -34,6 +47,7 @@ export type ConfigDescriptor<T> = {
default: T;
env?: [string, EnvConfigType];
link?: string;
internal?: boolean;
};
type ConfigDefineDescriptor<T> = {
@@ -44,6 +58,7 @@ type ConfigDefineDescriptor<T> = {
env?: string | [string, EnvConfigType];
link?: string;
schema?: JSONSchema;
internal?: boolean;
};
function typeFromShape(shape: z.ZodType<any>): ConfigType {
@@ -87,19 +102,17 @@ function shapeFromType(type: ConfigType): z.ZodType<any> {
}
function typeFromSchema(schema: JSONSchema): ConfigType {
if ('type' in schema) {
switch (schema.type) {
case 'string':
return 'string';
case 'number':
return 'float';
case 'boolean':
return 'boolean';
case 'array':
return 'array';
case 'object':
return 'object';
}
switch (schema.type) {
case 'string':
return 'string';
case 'number':
return 'float';
case 'boolean':
return 'boolean';
case 'array':
return 'array';
case 'object':
return 'object';
}
return 'any';
@@ -168,6 +181,7 @@ function standardizeDescriptor<T>(
},
env,
link: desc.link,
internal: desc.internal,
schema: {
type: schemaFromType(type),
description: desc.desc,
@@ -200,12 +214,20 @@ export const getDescriptors = once(() => {
export function defineModuleConfig<T extends keyof AppConfigSchema>(
module: T,
defs: ModuleConfigDescriptors<AppConfigByPath<T>>
) {
registerModuleConfig(
module,
defs as Record<string, ConfigDefineDescriptor<unknown>>
);
}
function registerModuleConfig(
module: string,
defs: Record<string, ConfigDefineDescriptor<unknown>>
) {
const descriptors: Record<string, ConfigDescriptor<any>> = {};
Object.entries(defs).forEach(([key, desc]) => {
descriptors[key] = standardizeDescriptor(
desc as ConfigDefineDescriptor<any>
);
descriptors[key] = standardizeDescriptor(desc);
});
APP_CONFIG_DESCRIPTORS[module] = {
@@ -214,7 +236,52 @@ export function defineModuleConfig<T extends keyof AppConfigSchema>(
};
}
const CONFIG_JSON_PATHS = [
export type NativeAppConfigDescriptor = {
key: string;
description: string;
defaultValue: unknown;
schema: JSONSchema;
internal: boolean;
};
export function defineNativeModuleConfig<T extends keyof AppConfigSchema>(
module: T,
descriptors: NativeAppConfigDescriptor[],
validate: (module: string, key: string, value: unknown) => string[],
nodeDefinitions: Partial<ModuleConfigDescriptors<AppConfigByPath<T>>> = {}
) {
registerModuleConfig(module, {
...nodeDefinitions,
...Object.fromEntries(
descriptors.map(descriptor => [
descriptor.key,
{
desc: descriptor.description,
default: descriptor.defaultValue,
schema: descriptor.schema,
internal: descriptor.internal,
validate: (value: unknown) => {
const errors = validate(module, descriptor.key, value);
return errors.length
? {
success: false as const,
error: new z.ZodError(
errors.map(message => ({
code: z.ZodIssueCode.custom,
message,
path: [],
}))
),
}
: { success: true as const, data: value };
},
},
])
),
} as Record<string, ConfigDefineDescriptor<unknown>>);
}
export const CONFIG_JSON_PATHS = [
join(env.projectRoot, 'config.json'),
`${homedir()}/.affine/config/config.json`,
];
packages/backend/server/src/base/error/def.ts
+21 -37
View File
@@ -1,5 +1,4 @@
import { STATUS_CODES } from 'node:http';
import { escape } from 'node:querystring';
import { HttpStatus, Logger } from '@nestjs/common';
import { ClsServiceManager } from 'nestjs-cls';
@@ -791,35 +790,6 @@ export const USER_FRIENDLY_ERRORS = {
message: ({ provider, kind, message }) =>
`Provider ${provider} failed with ${kind} error: ${message || 'unknown'}`,
},
copilot_invalid_context: {
type: 'invalid_input',
args: { contextId: 'string' },
message: ({ contextId }) => `Invalid copilot context ${contextId}.`,
},
copilot_context_file_not_supported: {
type: 'bad_request',
args: { fileName: 'string', message: 'string' },
message: ({ fileName, message }) =>
`File ${fileName} is not supported to use as context: ${message}`,
},
copilot_failed_to_modify_context: {
type: 'internal_server_error',
args: { contextId: 'string', message: 'string' },
message: ({ contextId, message }) =>
`Failed to modify context ${contextId}: ${message}`,
},
copilot_failed_to_match_context: {
type: 'internal_server_error',
args: { contextId: 'string', content: 'string', message: 'string' },
message: ({ contextId, content, message }) =>
`Failed to match context ${contextId} with "${escape(content)}": ${message}`,
},
copilot_failed_to_match_global_context: {
type: 'internal_server_error',
args: { workspaceId: 'string', content: 'string', message: 'string' },
message: ({ workspaceId, content, message }) =>
`Failed to match context in workspace ${workspaceId} with "${escape(content)}": ${message}`,
},
copilot_embedding_disabled: {
type: 'action_forbidden',
message: `Embedding feature is disabled, please contact the administrator to enable it in the workspace settings.`,
@@ -828,6 +798,27 @@ export const USER_FRIENDLY_ERRORS = {
type: 'action_forbidden',
message: `Embedding feature not available, you may need to install pgvector extension to your database`,
},
copilot_selected_sources_processing: {
type: 'bad_request',
message: `Selected sources are still processing. Try again shortly.`,
},
copilot_selected_sources_failed: {
type: 'bad_request',
message: `Selected sources could not be processed. Remove the failed source or try again.`,
},
copilot_selected_sources_unavailable: {
type: 'action_forbidden',
message: `Selected sources are not available for AI retrieval.`,
},
copilot_selected_sources_limit_exceeded: {
type: 'invalid_input',
message: `Too many or too much content was selected. Select fewer sources and try again.`,
},
copilot_failed_to_add_workspace_artifact: {
type: 'internal_server_error',
args: { message: 'string' },
message: ({ message }) => `Failed to add workspace artifact: ${message}`,
},
copilot_transcription_job_exists: {
type: 'bad_request',
message: 'Transcription job already exists',
@@ -840,13 +831,6 @@ export const USER_FRIENDLY_ERRORS = {
type: 'bad_request',
message: `Audio not provided.`,
},
copilot_failed_to_add_workspace_file_embedding: {
type: 'internal_server_error',
args: { message: 'string' },
message: ({ message }) =>
`Failed to add workspace file embedding: ${message}`,
},
// Quota & Limit errors
blob_quota_exceeded: {
type: 'quota_exceeded',
packages/backend/server/src/base/error/errors.gen.ts
+40 -73
View File
@@ -868,62 +868,6 @@ export class CopilotProviderSideError extends UserFriendlyError {
super('internal_server_error', 'copilot_provider_side_error', message, args);
}
}
@ObjectType()
class CopilotInvalidContextDataType {
@Field() contextId!: string
}
export class CopilotInvalidContext extends UserFriendlyError {
constructor(args: CopilotInvalidContextDataType, message?: string | ((args: CopilotInvalidContextDataType) => string)) {
super('invalid_input', 'copilot_invalid_context', message, args);
}
}
@ObjectType()
class CopilotContextFileNotSupportedDataType {
@Field() fileName!: string
@Field() message!: string
}
export class CopilotContextFileNotSupported extends UserFriendlyError {
constructor(args: CopilotContextFileNotSupportedDataType, message?: string | ((args: CopilotContextFileNotSupportedDataType) => string)) {
super('bad_request', 'copilot_context_file_not_supported', message, args);
}
}
@ObjectType()
class CopilotFailedToModifyContextDataType {
@Field() contextId!: string
@Field() message!: string
}
export class CopilotFailedToModifyContext extends UserFriendlyError {
constructor(args: CopilotFailedToModifyContextDataType, message?: string | ((args: CopilotFailedToModifyContextDataType) => string)) {
super('internal_server_error', 'copilot_failed_to_modify_context', message, args);
}
}
@ObjectType()
class CopilotFailedToMatchContextDataType {
@Field() contextId!: string
@Field() content!: string
@Field() message!: string
}
export class CopilotFailedToMatchContext extends UserFriendlyError {
constructor(args: CopilotFailedToMatchContextDataType, message?: string | ((args: CopilotFailedToMatchContextDataType) => string)) {
super('internal_server_error', 'copilot_failed_to_match_context', message, args);
}
}
@ObjectType()
class CopilotFailedToMatchGlobalContextDataType {
@Field() workspaceId!: string
@Field() content!: string
@Field() message!: string
}
export class CopilotFailedToMatchGlobalContext extends UserFriendlyError {
constructor(args: CopilotFailedToMatchGlobalContextDataType, message?: string | ((args: CopilotFailedToMatchGlobalContextDataType) => string)) {
super('internal_server_error', 'copilot_failed_to_match_global_context', message, args);
}
}
export class CopilotEmbeddingDisabled extends UserFriendlyError {
constructor(message?: string) {
@@ -937,6 +881,40 @@ export class CopilotEmbeddingUnavailable extends UserFriendlyError {
}
}
export class CopilotSelectedSourcesProcessing extends UserFriendlyError {
constructor(message?: string) {
super('bad_request', 'copilot_selected_sources_processing', message);
}
}
export class CopilotSelectedSourcesFailed extends UserFriendlyError {
constructor(message?: string) {
super('bad_request', 'copilot_selected_sources_failed', message);
}
}
export class CopilotSelectedSourcesUnavailable extends UserFriendlyError {
constructor(message?: string) {
super('action_forbidden', 'copilot_selected_sources_unavailable', message);
}
}
export class CopilotSelectedSourcesLimitExceeded extends UserFriendlyError {
constructor(message?: string) {
super('invalid_input', 'copilot_selected_sources_limit_exceeded', message);
}
}
@ObjectType()
class CopilotFailedToAddWorkspaceArtifactDataType {
@Field() message!: string
}
export class CopilotFailedToAddWorkspaceArtifact extends UserFriendlyError {
constructor(args: CopilotFailedToAddWorkspaceArtifactDataType, message?: string | ((args: CopilotFailedToAddWorkspaceArtifactDataType) => string)) {
super('internal_server_error', 'copilot_failed_to_add_workspace_artifact', message, args);
}
}
export class CopilotTranscriptionJobExists extends UserFriendlyError {
constructor(message?: string) {
super('bad_request', 'copilot_transcription_job_exists', message);
@@ -954,16 +932,6 @@ export class CopilotTranscriptionAudioNotProvided extends UserFriendlyError {
super('bad_request', 'copilot_transcription_audio_not_provided', message);
}
}
@ObjectType()
class CopilotFailedToAddWorkspaceFileEmbeddingDataType {
@Field() message!: string
}
export class CopilotFailedToAddWorkspaceFileEmbedding extends UserFriendlyError {
constructor(args: CopilotFailedToAddWorkspaceFileEmbeddingDataType, message?: string | ((args: CopilotFailedToAddWorkspaceFileEmbeddingDataType) => string)) {
super('internal_server_error', 'copilot_failed_to_add_workspace_file_embedding', message, args);
}
}
export class BlobQuotaExceeded extends UserFriendlyError {
constructor(message?: string) {
@@ -1317,17 +1285,16 @@ export enum ErrorNames {
COPILOT_PROMPT_INVALID,
COPILOT_PROVIDER_NOT_SUPPORTED,
COPILOT_PROVIDER_SIDE_ERROR,
COPILOT_INVALID_CONTEXT,
COPILOT_CONTEXT_FILE_NOT_SUPPORTED,
COPILOT_FAILED_TO_MODIFY_CONTEXT,
COPILOT_FAILED_TO_MATCH_CONTEXT,
COPILOT_FAILED_TO_MATCH_GLOBAL_CONTEXT,
COPILOT_EMBEDDING_DISABLED,
COPILOT_EMBEDDING_UNAVAILABLE,
COPILOT_SELECTED_SOURCES_PROCESSING,
COPILOT_SELECTED_SOURCES_FAILED,
COPILOT_SELECTED_SOURCES_UNAVAILABLE,
COPILOT_SELECTED_SOURCES_LIMIT_EXCEEDED,
COPILOT_FAILED_TO_ADD_WORKSPACE_ARTIFACT,
COPILOT_TRANSCRIPTION_JOB_EXISTS,
COPILOT_TRANSCRIPTION_JOB_NOT_FOUND,
COPILOT_TRANSCRIPTION_AUDIO_NOT_PROVIDED,
COPILOT_FAILED_TO_ADD_WORKSPACE_FILE_EMBEDDING,
BLOB_QUOTA_EXCEEDED,
STORAGE_QUOTA_EXCEEDED,
MEMBER_QUOTA_EXCEEDED,
@@ -1368,5 +1335,5 @@ registerEnumType(ErrorNames, {
export const ErrorDataUnionType = createUnionType({
name: 'ErrorDataUnion',
types: () =>
[GraphqlBadRequestDataType, HttpRequestErrorDataType, SsrfBlockedErrorDataType, ResponseTooLargeErrorDataType, ImageFormatNotSupportedDataType, QueryTooLongDataType, ValidationErrorDataType, WrongSignInCredentialsDataType, UnknownOauthProviderDataType, InvalidOauthCallbackCodeDataType, MissingOauthQueryParameterDataType, InvalidOauthResponseDataType, InvalidEmailDataType, InvalidPasswordLengthDataType, WorkspacePermissionNotFoundDataType, SpaceNotFoundDataType, MemberNotFoundInSpaceDataType, NotInSpaceDataType, AlreadyInSpaceDataType, SpaceAccessDeniedDataType, SpaceOwnerNotFoundDataType, SpaceShouldHaveOnlyOneOwnerDataType, DocNotFoundDataType, DocActionDeniedDataType, DocUpdateBlockedDataType, VersionRejectedDataType, InvalidHistoryTimestampDataType, DocHistoryNotFoundDataType, BlobNotFoundDataType, ExpectToGrantDocUserRolesDataType, ExpectToRevokeDocUserRolesDataType, ExpectToUpdateDocUserRoleDataType, NoMoreSeatDataType, UnsupportedSubscriptionPlanDataType, SubscriptionAlreadyExistsDataType, SubscriptionNotExistsDataType, SameSubscriptionRecurringDataType, SubscriptionPlanNotFoundDataType, CalendarProviderRequestErrorDataType, NoCopilotProviderAvailableDataType, CopilotFailedToGenerateEmbeddingDataType, CopilotDocNotFoundDataType, CopilotMessageNotFoundDataType, CopilotPromptNotFoundDataType, CopilotProviderNotSupportedDataType, CopilotProviderSideErrorDataType, CopilotInvalidContextDataType, CopilotContextFileNotSupportedDataType, CopilotFailedToModifyContextDataType, CopilotFailedToMatchContextDataType, CopilotFailedToMatchGlobalContextDataType, CopilotFailedToAddWorkspaceFileEmbeddingDataType, RuntimeConfigNotFoundDataType, InvalidRuntimeConfigTypeDataType, InvalidLicenseToActivateDataType, InvalidLicenseUpdateParamsDataType, UnsupportedClientVersionDataType, UnsupportedServerVersionDataType, MentionUserDocAccessDeniedDataType, InvalidAppConfigDataType, InvalidAppConfigInputDataType, InvalidSearchProviderRequestDataType, InvalidIndexerInputDataType] as const,
[GraphqlBadRequestDataType, HttpRequestErrorDataType, SsrfBlockedErrorDataType, ResponseTooLargeErrorDataType, ImageFormatNotSupportedDataType, QueryTooLongDataType, ValidationErrorDataType, WrongSignInCredentialsDataType, UnknownOauthProviderDataType, InvalidOauthCallbackCodeDataType, MissingOauthQueryParameterDataType, InvalidOauthResponseDataType, InvalidEmailDataType, InvalidPasswordLengthDataType, WorkspacePermissionNotFoundDataType, SpaceNotFoundDataType, MemberNotFoundInSpaceDataType, NotInSpaceDataType, AlreadyInSpaceDataType, SpaceAccessDeniedDataType, SpaceOwnerNotFoundDataType, SpaceShouldHaveOnlyOneOwnerDataType, DocNotFoundDataType, DocActionDeniedDataType, DocUpdateBlockedDataType, VersionRejectedDataType, InvalidHistoryTimestampDataType, DocHistoryNotFoundDataType, BlobNotFoundDataType, ExpectToGrantDocUserRolesDataType, ExpectToRevokeDocUserRolesDataType, ExpectToUpdateDocUserRoleDataType, NoMoreSeatDataType, UnsupportedSubscriptionPlanDataType, SubscriptionAlreadyExistsDataType, SubscriptionNotExistsDataType, SameSubscriptionRecurringDataType, SubscriptionPlanNotFoundDataType, CalendarProviderRequestErrorDataType, NoCopilotProviderAvailableDataType, CopilotFailedToGenerateEmbeddingDataType, CopilotDocNotFoundDataType, CopilotMessageNotFoundDataType, CopilotPromptNotFoundDataType, CopilotProviderNotSupportedDataType, CopilotProviderSideErrorDataType, CopilotFailedToAddWorkspaceArtifactDataType, RuntimeConfigNotFoundDataType, InvalidRuntimeConfigTypeDataType, InvalidLicenseToActivateDataType, InvalidLicenseUpdateParamsDataType, UnsupportedClientVersionDataType, UnsupportedServerVersionDataType, MentionUserDocAccessDeniedDataType, InvalidAppConfigDataType, InvalidAppConfigInputDataType, InvalidSearchProviderRequestDataType, InvalidIndexerInputDataType] as const,
});
packages/backend/server/src/base/index.ts
+2
View File
@@ -10,7 +10,9 @@ export {
Config,
ConfigFactory,
defineModuleConfig,
defineNativeModuleConfig,
type JSONSchema,
type NativeAppConfigDescriptor,
} from './config';
export * from './cors';
export * from './error';
packages/backend/server/src/core/backend-runtime/__tests__/job.spec.ts
+175 -2
View File
@@ -1,3 +1,6 @@
import { readFileSync } from 'node:fs';
import { join } from 'node:path';
import { ScheduleModule } from '@nestjs/schedule';
import ava, { TestFn } from 'ava';
import Sinon from 'sinon';
@@ -6,26 +9,50 @@ import {
createTestingModule,
type TestingModule,
} from '../../../__tests__/utils';
import {
CopilotSelectedSourcesFailed,
CopilotSelectedSourcesLimitExceeded,
CopilotSelectedSourcesProcessing,
CopilotSelectedSourcesUnavailable,
} from '../../../base';
import { Models } from '../../../models';
import { BackendRuntimeModule, BackendRuntimeProvider } from '../index';
import { BackendRuntimeHousekeepingJob } from '../job';
import {
BackendRuntimeEmbeddingJob,
BackendRuntimeHousekeepingJob,
} from '../job';
interface Context {
module: TestingModule;
embeddingJob: BackendRuntimeEmbeddingJob;
job: BackendRuntimeHousekeepingJob;
getSnapshot: Sinon.SinonStub;
allowEmbedding: Sinon.SinonStub;
runtime: {
cleanupExpiredRuntimeStates: Sinon.SinonStub;
cleanupExpiredRuntimeGates: Sinon.SinonStub;
cleanupExpiredRollingQuota: Sinon.SinonStub;
cleanupUnreferencedArtifacts: Sinon.SinonStub;
reconcileEmbeddingWorkspaces: Sinon.SinonStub;
embeddingHealth: Sinon.SinonStub;
syncEmbeddingState: Sinon.SinonStub;
};
}
const test = ava as TestFn<Context>;
test.before(async t => {
const snapshot = readFileSync(
join(process.cwd(), 'src/__tests__/__fixtures__/test-doc.snapshot.bin')
);
t.context.runtime = {
cleanupExpiredRuntimeStates: Sinon.stub(),
cleanupExpiredRuntimeGates: Sinon.stub(),
cleanupExpiredRollingQuota: Sinon.stub(),
cleanupUnreferencedArtifacts: Sinon.stub(),
reconcileEmbeddingWorkspaces: Sinon.stub(),
embeddingHealth: Sinon.stub().resolves({ enabled: true }),
syncEmbeddingState: Sinon.stub(),
};
t.context.module = await createTestingModule({
imports: [ScheduleModule.forRoot(), BackendRuntimeModule],
@@ -35,6 +62,25 @@ test.before(async t => {
.useValue(t.context.runtime);
},
});
const models = t.context.module.get(Models);
t.context.getSnapshot = Sinon.stub(models.doc, 'getSnapshot').resolves({
workspaceId: 'workspace-1',
id: 'doc-1',
blob: snapshot,
size: BigInt(snapshot.length),
state: null,
createdAt: new Date('2026-01-01T00:00:00Z'),
updatedAt: new Date('2026-01-02T00:00:00Z'),
createdBy: null,
updatedBy: null,
createdByUser: null,
updatedByUser: null,
});
t.context.allowEmbedding = Sinon.stub(
models.workspace,
'allowEmbedding'
).resolves(true);
t.context.embeddingJob = t.context.module.get(BackendRuntimeEmbeddingJob);
t.context.job = t.context.module.get(BackendRuntimeHousekeepingJob);
});
@@ -42,21 +88,148 @@ test.beforeEach(t => {
t.context.runtime.cleanupExpiredRuntimeStates.reset();
t.context.runtime.cleanupExpiredRuntimeGates.reset();
t.context.runtime.cleanupExpiredRollingQuota.reset();
t.context.runtime.cleanupUnreferencedArtifacts.reset();
t.context.runtime.reconcileEmbeddingWorkspaces.reset();
t.context.runtime.embeddingHealth.resetHistory();
t.context.runtime.syncEmbeddingState.reset();
t.context.getSnapshot.resetHistory();
t.context.allowEmbedding.resetHistory();
});
test.after.always(async t => {
Sinon.restore();
await t.context.module.close();
});
test('backend-runtime housekeeping cleans runtime state and gate batches', async t => {
test('backend-runtime jobs ingest documents and clean runtime state', async t => {
await t.context.embeddingJob.onDocSnapshotUpdated({
workspaceId: 'workspace-1',
docId: 'doc-1',
blob: Buffer.alloc(0),
});
const { payload } = await t.context.module.queue.waitFor(
'backendRuntime.syncDocumentEmbedding'
);
await t.context.embeddingJob.syncDocument(payload);
t.is(t.context.getSnapshot.callCount, 1);
t.is(t.context.runtime.syncEmbeddingState.callCount, 1);
t.like(t.context.runtime.syncEmbeddingState.firstCall.args[0], {
workspaceId: 'workspace-1',
enabled: true,
reconcileDocuments: true,
});
t.is(
t.context.runtime.syncEmbeddingState.firstCall.args[0].documents[0].docId,
'doc-1'
);
t.true(
t.context.runtime.syncEmbeddingState.firstCall.args[0].documents[0].units
.length > 0
);
const documentJobCount = t.context.module.queue.count(
'backendRuntime.syncDocumentEmbedding'
);
await t.context.embeddingJob.onDocSnapshotUpdated({
workspaceId: 'workspace-1',
docId: 'db$docProperties',
blob: Buffer.alloc(0),
});
t.is(
t.context.module.queue.count('backendRuntime.syncDocumentEmbedding'),
documentJobCount
);
await t.context.embeddingJob.onDocSnapshotUpdated({
workspaceId: 'workspace-1',
docId: 'workspace-1',
blob: Buffer.alloc(0),
});
const reconcile = await t.context.module.queue.waitFor(
'backendRuntime.reconcileDocumentEmbeddings'
);
await t.context.embeddingJob.reconcileDocuments(reconcile.payload);
t.like(t.context.runtime.syncEmbeddingState.secondCall.args[0], {
workspaceId: 'workspace-1',
enabled: true,
reconcileDocuments: true,
});
await t.context.embeddingJob.prepareSelectedDocuments('workspace-1', [
'doc-1',
'doc-1',
]);
t.like(t.context.runtime.syncEmbeddingState.thirdCall.args[0], {
workspaceId: 'workspace-1',
enabled: true,
reconcileDocuments: false,
priority: 1000,
waitForReadyMs: 90_000,
});
t.is(
t.context.runtime.syncEmbeddingState.thirdCall.args[0].documents.length,
1
);
for (const [nativeError, expectedError] of [
['embedding_selected_sources_processing', CopilotSelectedSourcesProcessing],
['embedding_selected_sources_failed', CopilotSelectedSourcesFailed],
[
'embedding_selected_sources_unavailable',
CopilotSelectedSourcesUnavailable,
],
] as const) {
t.context.runtime.syncEmbeddingState.rejects(new Error(nativeError));
const error = await t.throwsAsync(() =>
t.context.embeddingJob.prepareSelectedDocuments('workspace-1', ['doc-1'])
);
t.true(error instanceof expectedError);
}
t.context.runtime.syncEmbeddingState.resolves(undefined);
await t.throwsAsync(
() =>
t.context.embeddingJob.prepareSelectedDocuments(
'workspace-1',
Array.from({ length: 65 }, (_, index) => `doc-${index}`)
),
{ instanceOf: CopilotSelectedSourcesLimitExceeded }
);
t.context.getSnapshot.resolves(null);
await t.throwsAsync(
() =>
t.context.embeddingJob.prepareSelectedDocuments('workspace-1', [
'missing-doc',
]),
{ instanceOf: CopilotSelectedSourcesUnavailable }
);
const callsBeforeMissingBackgroundDoc =
t.context.runtime.syncEmbeddingState.callCount;
await t.context.embeddingJob.syncDocument({
workspaceId: 'workspace-1',
docId: 'missing-doc',
});
t.is(
t.context.runtime.syncEmbeddingState.callCount,
callsBeforeMissingBackgroundDoc + 1
);
t.deepEqual(
t.context.runtime.syncEmbeddingState.lastCall.args[0].documents,
[]
);
t.context.runtime.cleanupExpiredRuntimeStates.onCall(0).resolves(1000);
t.context.runtime.cleanupExpiredRuntimeStates.onCall(1).resolves(2);
t.context.runtime.cleanupExpiredRuntimeGates.resolves(1);
t.context.runtime.cleanupExpiredRollingQuota.resolves(1);
t.context.runtime.cleanupUnreferencedArtifacts.resolves(1);
t.context.runtime.reconcileEmbeddingWorkspaces.resolves(2);
await t.context.job.cleanExpiredRuntimeHousekeeping();
t.is(t.context.runtime.cleanupExpiredRuntimeStates.callCount, 2);
t.is(t.context.runtime.cleanupExpiredRuntimeGates.callCount, 1);
t.is(t.context.runtime.cleanupExpiredRollingQuota.callCount, 1);
t.is(t.context.runtime.cleanupUnreferencedArtifacts.callCount, 1);
t.is(t.context.runtime.reconcileEmbeddingWorkspaces.callCount, 1);
});
packages/backend/server/src/core/backend-runtime/__tests__/provider.spec.ts
+3 -1
View File
@@ -29,12 +29,14 @@ test('backend-runtime provider starts once, runs migrations once, and reports he
await provider.start();
await provider.onConfigChanged({ updates: { mailer: {} } });
await provider.onConfigChanged({ updates: { copilot: {} } });
await provider.onConfigChanged({ updates: { storages: {} } });
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.is(runtime.reloadConfig.callCount, 2);
t.true(runtime.reloadConfig.alwaysCalledWithExactly(privateKey));
t.true(health.databaseConnected);
t.is(runtime.stop.callCount, 1);
});
packages/backend/server/src/core/backend-runtime/index.ts
+20 -4
View File
@@ -1,16 +1,32 @@
import { Global, Module } from '@nestjs/common';
import { BackendRuntimeHousekeepingJob } from './job';
import { BackendRuntimeProvider } from './provider';
import {
BackendRuntimeEmbeddingJob,
BackendRuntimeHousekeepingJob,
} from './job';
import {
BACKEND_RUNTIME_CONFIG_PATHS,
BackendRuntimeProvider,
} from './provider';
@Global()
@Module({
providers: [BackendRuntimeProvider, BackendRuntimeHousekeepingJob],
exports: [BackendRuntimeProvider],
providers: [
{
provide: BACKEND_RUNTIME_CONFIG_PATHS,
useValue: undefined,
},
BackendRuntimeProvider,
BackendRuntimeEmbeddingJob,
BackendRuntimeHousekeepingJob,
],
exports: [BackendRuntimeProvider, BackendRuntimeEmbeddingJob],
})
export class BackendRuntimeModule {}
export { BackendRuntimeEmbeddingJob } from './job';
export {
BACKEND_RUNTIME_CONFIG_PATHS,
BackendRuntimeProvider,
type RuntimeInviteAbuseAction,
type RuntimeInviteAbuseClaimedAction,
packages/backend/server/src/core/backend-runtime/job.ts
+183 -2
View File
@@ -1,12 +1,189 @@
import { Injectable, Logger } from '@nestjs/common';
import { Cron, CronExpression } from '@nestjs/schedule';
import { JobQueue, OnJob } from '../../base';
import {
CopilotSelectedSourcesFailed,
CopilotSelectedSourcesLimitExceeded,
CopilotSelectedSourcesProcessing,
CopilotSelectedSourcesUnavailable,
JobQueue,
OnEvent,
OnJob,
} from '../../base';
import { Models } from '../../models';
import { projectDocSearch } from '../utils/blocksuite';
import { BackendRuntimeProvider } from './provider';
const SELECTED_DOCUMENT_LIMIT = 64;
const SELECTED_DOCUMENT_UNIT_LIMIT = 20_000;
const SELECTED_DOCUMENT_TEXT_BYTE_LIMIT = 16 * 1024 * 1024;
const SELECTED_DOCUMENT_PRIORITY = 1000;
const SELECTED_DOCUMENT_WAIT_MS = 90_000;
declare global {
interface Jobs {
'nightly.cleanExpiredBackendRuntimeHousekeeping': {};
'backendRuntime.syncDocumentEmbedding': {
workspaceId: string;
docId: string;
};
'backendRuntime.reconcileDocumentEmbeddings': {
workspaceId: string;
};
}
}
@Injectable()
export class BackendRuntimeEmbeddingJob {
constructor(
private readonly rt: BackendRuntimeProvider,
private readonly queue: JobQueue,
private readonly models: Models
) {}
@OnEvent('doc.updated')
async onDocUpdated({ workspaceId, docId }: Events['doc.updated']) {
await this.queueDocument(workspaceId, docId);
}
@OnEvent('doc.snapshot.updated')
async onDocSnapshotUpdated({
workspaceId,
docId,
}: Events['doc.snapshot.updated']) {
if (workspaceId === docId) {
await this.queue.add(
'backendRuntime.reconcileDocumentEmbeddings',
{ workspaceId },
{ jobId: `reconcileDocumentEmbeddings/${workspaceId}` }
);
return;
}
await this.queueDocument(workspaceId, docId);
}
private async queueDocument(workspaceId: string, docId: string) {
if (
workspaceId === docId ||
docId.startsWith('db$') ||
docId.startsWith('userdata$')
) {
return;
}
await this.queue.add(
'backendRuntime.syncDocumentEmbedding',
{ workspaceId, docId },
{ jobId: `syncDocumentEmbedding/${workspaceId}/${docId}` }
);
}
@OnJob('backendRuntime.syncDocumentEmbedding')
async syncDocument({
workspaceId,
docId,
}: Jobs['backendRuntime.syncDocumentEmbedding']) {
await this.syncDocuments(workspaceId, [docId], true);
}
async prepareSelectedDocuments(workspaceId: string, docIds: string[]) {
const selectedDocIds = [...new Set(docIds)];
if (selectedDocIds.length > SELECTED_DOCUMENT_LIMIT) {
throw new CopilotSelectedSourcesLimitExceeded();
}
try {
await this.syncDocuments(workspaceId, selectedDocIds, false, {
priority: SELECTED_DOCUMENT_PRIORITY,
waitForReadyMs: SELECTED_DOCUMENT_WAIT_MS,
});
} catch (error) {
throw this.mapSelectedSourceError(error);
}
}
private mapSelectedSourceError(error: unknown) {
const message = error instanceof Error ? error.message : String(error);
if (message.includes('embedding_selected_sources_processing')) {
return new CopilotSelectedSourcesProcessing();
}
if (message.includes('embedding_selected_sources_failed')) {
return new CopilotSelectedSourcesFailed();
}
if (message.includes('embedding_selected_sources_unavailable')) {
return new CopilotSelectedSourcesUnavailable();
}
return error;
}
private async syncDocuments(
workspaceId: string,
docIds: string[],
reconcileDocuments: boolean,
scheduling?: { priority: number; waitForReadyMs: number }
) {
if (!(await this.rt.embeddingHealth()).enabled) {
if (scheduling) throw new CopilotSelectedSourcesUnavailable();
return;
}
const enabled = await this.models.workspace.allowEmbedding(workspaceId);
if (!enabled) {
if (scheduling) throw new CopilotSelectedSourcesUnavailable();
return;
}
const documents = [];
let unitCount = 0;
let textBytes = 0;
for (const docId of docIds) {
const snapshot = await this.models.doc.getSnapshot(workspaceId, docId);
if (!snapshot) {
if (scheduling) throw new CopilotSelectedSourcesUnavailable();
continue;
}
const revision = snapshot.updatedAt.getTime().toString();
const projection = projectDocSearch(snapshot.blob, docId, revision);
unitCount += projection.units.length;
for (const unit of projection.units) {
textBytes += Buffer.byteLength(unit.text);
}
if (
unitCount > SELECTED_DOCUMENT_UNIT_LIMIT ||
textBytes > SELECTED_DOCUMENT_TEXT_BYTE_LIMIT
) {
throw new CopilotSelectedSourcesLimitExceeded();
}
documents.push({
docId,
revision,
sourceHash: projection.sourceHash,
units: projection.units.map(unit => ({
unitId: unit.unitId,
visibility: unit.visibility,
text: unit.text,
blockId: unit.blockId,
elementId: unit.elementId,
frameId: unit.frameId,
})),
});
}
if (!documents.length && !reconcileDocuments) return;
await this.rt.syncEmbeddingState({
workspaceId,
enabled,
reconcileDocuments,
documents,
...scheduling,
});
}
@OnJob('backendRuntime.reconcileDocumentEmbeddings')
async reconcileDocuments({
workspaceId,
}: Jobs['backendRuntime.reconcileDocumentEmbeddings']) {
if (!(await this.rt.embeddingHealth()).enabled) return;
await this.rt.syncEmbeddingState({
workspaceId,
enabled: await this.models.workspace.allowEmbedding(workspaceId),
reconcileDocuments: true,
});
}
}
@@ -41,9 +218,13 @@ export class BackendRuntimeHousekeepingJob {
const rollingQuota = await this.cleanBatches(() =>
this.rt.cleanupExpiredRollingQuota(1000)
);
const artifacts = await this.cleanBatches(() =>
this.rt.cleanupUnreferencedArtifacts(1000)
);
const embeddingWorkspaces = await this.rt.reconcileEmbeddingWorkspaces();
this.logger.log(
`cleaned runtime housekeeping states=${states} gates=${gates} rollingQuota=${rollingQuota}`
`cleaned runtime housekeeping states=${states} gates=${gates} rollingQuota=${rollingQuota} artifacts=${artifacts} embeddingWorkspaces=${embeddingWorkspaces}`
);
}
packages/backend/server/src/core/backend-runtime/provider.ts
+132 -3
View File
@@ -1,4 +1,5 @@
import {
Inject,
Injectable,
Logger,
type OnApplicationBootstrap,
@@ -12,21 +13,35 @@ import {
BackendRuntime,
type BackendRuntimeHealth,
type ByokLocalLeaseOutput,
type ByokPolicyOutput,
type ByokProbeResultOutput,
type ByokProfileOutput,
type CompileScopeInput,
type CopilotExecuteInput,
type CopilotRouteCheckInput,
type CreateByokLocalLeaseInput,
type CreateByokProfileInput,
type EmbeddingHealth,
type EnsureWorkspaceBlobArtifactInput,
type MatchEmbeddingCandidatesInput,
type ProbeByokDraftInput,
type ProbeByokProfileInput,
type PutWorkspaceArtifactInput,
type ReadEmbeddingSourceContentInput,
type ReorderByokProfilesInput,
type ReplaceByokProfileInput,
type RotateByokCredentialInput,
type RuntimeTurnScopeSnapshot,
type RuntimeWorkspaceArtifact,
type SyncEmbeddingStateInput,
} from '../../native';
type RuntimeInstance = InstanceType<typeof BackendRuntime>;
export const BACKEND_RUNTIME_CONFIG_PATHS = Symbol(
'BACKEND_RUNTIME_CONFIG_PATHS'
);
class RuntimeEventStream<T> implements AsyncIterableIterator<T> {
private readonly values: T[] = [];
private readonly readers: Array<(result: IteratorResult<T>) => void> = [];
@@ -262,8 +277,16 @@ export class BackendRuntimeProvider
private readonly runtime: RuntimeInstance;
private migrationsStarted = false;
constructor(@Optional() private readonly config?: Config) {
this.runtime = new BackendRuntime(this.config?.crypto.privateKey);
constructor(
@Optional() private readonly config?: Config,
@Optional()
@Inject(BACKEND_RUNTIME_CONFIG_PATHS)
configPaths?: string[]
) {
this.runtime = new BackendRuntime(
this.config?.crypto.privateKey,
configPaths
);
}
async onApplicationBootstrap() {
@@ -288,7 +311,12 @@ export class BackendRuntimeProvider
@OnEvent('config.changed')
async onConfigChanged({ updates }: Events['config.changed']) {
if (!updates.copilot && !updates.crypto && !updates.db) {
if (
!updates.copilot &&
!updates.crypto &&
!updates.db &&
!updates.storages
) {
return;
}
await this.runtime.reloadConfig(this.config?.crypto.privateKey);
@@ -298,6 +326,101 @@ export class BackendRuntimeProvider
return await this.runtime.health();
}
async embeddingHealth(): Promise<EmbeddingHealth> {
return await this.measured('embeddingHealth', runtime =>
runtime.embeddingHealth()
);
}
async embeddingQueueCounts() {
return await this.measured('embeddingQueueCounts', runtime =>
runtime.embeddingQueueCounts()
);
}
async embeddingWorkspaceProgress(workspaceId: string) {
return await this.measured('embeddingWorkspaceProgress', runtime =>
runtime.embeddingWorkspaceProgress(workspaceId)
);
}
async reconcileEmbeddingWorkspaces() {
return await this.measured('reconcileEmbeddingWorkspaces', runtime =>
runtime.reconcileEmbeddingWorkspaces()
);
}
async compileTurnScope(
input: CompileScopeInput
): Promise<RuntimeTurnScopeSnapshot> {
return await this.measured('compileTurnScope', runtime =>
runtime.compileTurnScope(input)
);
}
async putWorkspaceArtifact(
input: PutWorkspaceArtifactInput,
body: Buffer
): Promise<RuntimeWorkspaceArtifact> {
return await this.measured('putWorkspaceArtifact', runtime =>
runtime.putWorkspaceArtifact(input, body)
);
}
async ensureWorkspaceBlobArtifact(
input: EnsureWorkspaceBlobArtifactInput
): Promise<RuntimeWorkspaceArtifact> {
return await this.measured('ensureWorkspaceBlobArtifact', runtime =>
runtime.ensureWorkspaceBlobArtifact(input)
);
}
async syncEmbeddingState(input: SyncEmbeddingStateInput) {
return await this.measured('syncEmbeddingState', runtime =>
runtime.syncEmbeddingState(input)
);
}
async readEmbeddingSourceContent(input: ReadEmbeddingSourceContentInput) {
return await this.measured('readEmbeddingSourceContent', runtime =>
runtime.readEmbeddingSourceContent(input)
);
}
async matchEmbeddingCandidates(input: MatchEmbeddingCandidatesInput) {
return await this.measured('matchEmbeddingCandidates', runtime =>
runtime.matchEmbeddingCandidates(input)
);
}
async cleanupUnreferencedArtifacts(limit: number) {
return await this.measured('cleanupUnreferencedArtifacts', runtime =>
runtime.cleanupUnreferencedArtifacts(limit)
);
}
async setArtifactLibraryOwned(
workspaceId: string,
artifactId: string,
libraryOwned: boolean,
displayName?: string
) {
return await this.measured('setArtifactLibraryOwned', runtime =>
runtime.setArtifactLibraryOwned(
workspaceId,
artifactId,
libraryOwned,
displayName
)
);
}
async cancelEmbeddingCandidateRequest(requestId: string) {
return await this.measured('cancelEmbeddingCandidateRequest', runtime =>
runtime.cancelEmbeddingCandidateRequest(requestId)
);
}
async cleanupExpiredSnapshotHistories(limit: number) {
return await this.measured('cleanupExpiredSnapshotHistories', rt =>
rt.cleanupExpiredSnapshotHistories(limit)
@@ -379,6 +502,12 @@ export class BackendRuntimeProvider
);
}
async getByokPolicy(): Promise<ByokPolicyOutput> {
return await this.measured('getByokPolicy', runtime =>
Promise.resolve(runtime.getByokPolicy())
);
}
async createByokProfile(
input: CreateByokProfileInput
): Promise<ByokProfileOutput> {
packages/backend/server/src/core/doc-renderer/__tests__/controller.spec.ts
+1
View File
@@ -150,6 +150,7 @@ const policyCases: Array<{
markdown: Sinon.stub(docReader, 'getDocMarkdown').resolves({
title: 'markdown-doc',
markdown: '# markdown-doc',
revision: '1',
knownUnsupportedBlocks: [],
unknownBlocks: [],
}),
packages/backend/server/src/core/doc-service/controller.ts
+14
View File
@@ -91,6 +91,20 @@ export class DocRpcController {
res.send(Buffer.concat([diff.missing, diff.state]));
}
@SkipThrottle()
@Internal()
@Get('/workspaces/:workspaceId/docs/:docId/canvas')
async getDocCanvas(
@Param('workspaceId') workspaceId: string,
@Param('docId') docId: string
) {
const projection = await this.docReader.getDocCanvas(workspaceId, docId);
if (!projection) {
throw new NotFound('Doc not found');
}
return projection;
}
@SkipThrottle()
@Internal()
@Get('/workspaces/:workspaceId/docs/:docId/content')
packages/backend/server/src/core/doc/__tests__/reader-from-database.spec.ts
+16 -1
View File
@@ -278,7 +278,21 @@ test('should return doc markdown success', async t => {
docSnapshot.id,
false
);
t.snapshot(result);
if (result) {
const { revision, ...markdown } = result;
t.truthy(revision);
t.snapshot(markdown);
}
const canvas = await docReader.getDocCanvas(workspace.id, docSnapshot.id);
t.is(canvas?.version, 1);
t.is(canvas?.docId, docSnapshot.id);
t.truthy(canvas?.revision);
t.deepEqual(canvas?.counts, {
connector: 6,
group: 6,
shape: 7,
text: 7,
});
});
test('should read markdown return null when doc not exists', async t => {
@@ -293,4 +307,5 @@ test('should read markdown return null when doc not exists', async t => {
false
);
t.is(result, null);
t.is(await docReader.getDocCanvas(workspace.id, randomUUID()), null);
});
packages/backend/server/src/core/doc/__tests__/reader-from-rpc.spec.ts
+16 -1
View File
@@ -397,7 +397,21 @@ test('should return doc markdown success', async t => {
docSnapshot.id,
false
);
t.snapshot(result);
if (result) {
const { revision, ...markdown } = result;
t.truthy(revision);
t.snapshot(markdown);
}
const canvas = await docReader.getDocCanvas(workspace.id, docSnapshot.id);
t.is(canvas?.version, 1);
t.is(canvas?.docId, docSnapshot.id);
t.truthy(canvas?.revision);
t.deepEqual(canvas?.counts, {
connector: 6,
group: 6,
shape: 7,
text: 7,
});
});
test('should read markdown return null when doc not exists', async t => {
@@ -414,4 +428,5 @@ test('should read markdown return null when doc not exists', async t => {
false
);
t.is(result, null);
t.is(await docReader.getDocCanvas(workspace.id, randomUUID()), null);
});
packages/backend/server/src/core/doc/reader.ts
+44 -1
View File
@@ -13,10 +13,13 @@ import {
import { Models } from '../../models';
import { WorkspaceBlobStorage } from '../storage';
import {
type CanvasProjectionV1,
type PageDocContent,
parseCanvasProjection,
parseDocToMarkdownFromDocSnapshot,
parsePageDoc,
parseWorkspaceDoc,
projectDocCanvas,
type WorkspaceDocContent,
} from '../utils/blocksuite';
import { PgWorkspaceDocStorageAdapter } from './adapters/workspace';
@@ -34,6 +37,7 @@ export interface WorkspaceDocInfo {
export interface DocMarkdown {
title: string;
markdown: string;
revision: string;
knownUnsupportedBlocks: string[];
unknownBlocks: string[];
}
@@ -68,6 +72,11 @@ export abstract class DocReader {
aiEditable: boolean
): Promise<DocMarkdown | null>;
abstract getDocCanvas(
workspaceId: string,
docId: string
): Promise<CanvasProjectionV1 | null>;
abstract getDocDiff(
spaceId: string,
docId: string,
@@ -205,13 +214,24 @@ export class DatabaseDocReader extends DocReader {
);
}
return markdown;
return { ...markdown, revision: doc.timestamp.toString() };
} catch (error) {
this.logger.error(`Failed to parse ${workspaceId}/${docId}.`, error);
throw error;
}
}
async getDocCanvas(
workspaceId: string,
docId: string
): Promise<CanvasProjectionV1 | null> {
const doc = await this.workspace.getDoc(workspaceId, docId);
if (!doc) {
return null;
}
return projectDocCanvas(doc.bin, docId, doc.timestamp.toString());
}
async getDocDiff(
spaceId: string,
docId: string,
@@ -387,6 +407,29 @@ export class RpcDocReader extends DatabaseDocReader {
}
}
override async getDocCanvas(
workspaceId: string,
docId: string
): Promise<CanvasProjectionV1 | null> {
const url = `${this.config.docService.endpoint}/rpc/workspaces/${workspaceId}/docs/${docId}/canvas`;
try {
const res = await this.fetch(url, 'GET');
if (!res) {
return null;
}
return parseCanvasProjection(await res.json());
} catch (e) {
if (e instanceof UserFriendlyError) {
throw e;
}
this.logger.error(
`Failed to fetch doc canvas ${url}, fallback to database doc reader`,
e as Error
);
return await super.getDocCanvas(workspaceId, docId);
}
}
override async getDocDiff(
workspaceId: string,
docId: string,
packages/backend/server/src/core/mail/__tests__/mailer.spec.ts
+7
View File
@@ -24,6 +24,7 @@ interface Context {
assertMailDeliveryQuotaV1: Sinon.SinonStub;
commitMailDeliveryQuotaV1: Sinon.SinonStub;
releaseMailDeliveryQuotaV1: Sinon.SinonStub;
embeddingHealth: Sinon.SinonStub;
};
}
@@ -34,6 +35,12 @@ test.before(async t => {
assertMailDeliveryQuotaV1: Sinon.stub(),
commitMailDeliveryQuotaV1: Sinon.stub(),
releaseMailDeliveryQuotaV1: Sinon.stub(),
embeddingHealth: Sinon.stub().resolves({
enabled: false,
state: 'disabled',
reason: 'test',
workerRunning: false,
}),
};
t.context.module = await createTestingModule({
tapModule: builder => {
packages/backend/server/src/core/realtime/__tests__/registry.spec.ts
+6 -26
View File
@@ -8,7 +8,7 @@ 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/realtime';
import { CopilotEmbeddingRealtimeProvider } from '../../../plugins/copilot/embedding/realtime';
import type { CopilotTranscriptionReader } from '../../../plugins/copilot/transcript/reader';
import { CopilotTranscriptRealtimeProvider } from '../../../plugins/copilot/transcript/realtime';
import type { CurrentUser } from '../../auth';
@@ -440,7 +440,6 @@ test('front and sync realtime gateway required handlers are registered by lightw
{} as never,
{} as never,
registry,
{} as never,
{} as never
).onModuleInit();
new CopilotTranscriptRealtimeProvider(
@@ -970,9 +969,8 @@ test('quota realtime provider exposes effective quota state snapshots', async t
);
});
test('copilot embedding realtime provider uses lightweight model reads', async t => {
test('copilot embedding realtime provider uses native health and progress', async t => {
const registry = new RealtimeRegistry();
const published: unknown[][] = [];
const assertions: unknown[] = [];
const ac = {
user(userId: string) {
@@ -990,25 +988,16 @@ test('copilot embedding realtime provider uses lightweight model reads', async t
};
},
} as unknown as PermissionAccess;
const models = {
copilotWorkspace: {
checkEmbeddingAvailable: async () => true,
getEmbeddingStatus: async () => ({ total: 5, embedded: 3 }),
},
copilotContext: {
getConfig: async () => ({ workspaceId: 'space' }),
},
const embedding = {
health: async () => ({ enabled: true }),
progress: async () => ({ total: 5, embedded: 3 }),
};
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,
embedding as never,
registry,
publisher,
config as never
);
provider.onModuleInit();
@@ -1035,18 +1024,9 @@ test('copilot embedding realtime provider uses lightweight model reads', async t
.room(user, { workspaceId: 'space' }),
realtimeWorkspaceEmbeddingProgressRoom('space')
);
await provider.onDocEmbedFinished({ contextId: 'context', docId: 'doc' });
t.deepEqual(assertions, [
{ userId: 'u1', workspaceId: 'space', action: 'Workspace.Copilot' },
]);
t.deepEqual(published[0], [
'workspace.embedding.progress.changed',
{ workspaceId: 'space' },
{ reason: 'finished' },
{ room: realtimeWorkspaceEmbeddingProgressRoom('space') },
]);
});
test('copilot transcript realtime provider registers task live query handlers', async t => {
packages/backend/server/src/core/realtime/gateway.ts
+22 -4
View File
@@ -4,7 +4,12 @@ import type {
RealtimeUnsubscribeEnvelope,
} from '@affine/realtime';
import { getRealtimeInputKey } from '@affine/realtime';
import { applyDecorators, Logger, UseInterceptors } from '@nestjs/common';
import {
applyDecorators,
Logger,
Optional,
UseInterceptors,
} from '@nestjs/common';
import {
ConnectedSocket,
MessageBody,
@@ -20,6 +25,7 @@ import type { Server, Socket } from 'socket.io';
import {
checkCanaryDateClientVersion,
EventBus,
GatewayErrorWrapper,
OnEvent,
UnsupportedClientVersion,
@@ -63,7 +69,8 @@ export class RealtimeGateway implements OnGatewayInit, OnGatewayDisconnect {
constructor(
private readonly registry: RealtimeRegistry,
private readonly publisher: RealtimePublisher
private readonly publisher: RealtimePublisher,
@Optional() private readonly event?: EventBus
) {}
afterInit(_server: Server) {
@@ -76,17 +83,28 @@ export class RealtimeGateway implements OnGatewayInit, OnGatewayDisconnect {
this.subscriptions.delete(subscriptionId);
}
}
this.event?.emit('realtime.connection.disconnected', {
connectionId: client.id,
});
this.event?.broadcast('realtime.connection.disconnected', {
connectionId: client.id,
});
}
@SubscribeMessage('realtime:request')
async onRequest(
@CurrentUser() user: CurrentUser,
@MessageBody() envelope: RealtimeRequestEnvelope
@MessageBody() envelope: RealtimeRequestEnvelope,
@ConnectedSocket() client?: Socket
) {
this.assertVersion(envelope.clientVersion);
const handler = this.registry.getRequest(envelope.op);
const input = handler.input.parse(envelope.input);
return { data: await handler.handle(user, input as never) };
return {
data: await handler.handle(user, input as never, {
connectionId: client?.id,
}),
};
}
@SubscribeMessage('realtime:subscribe')
packages/backend/server/src/core/realtime/required-handlers.ts
-2
View File
@@ -11,7 +11,6 @@ export const REALTIME_GATEWAY_REQUIRED_REQUESTS = [
'user.settings.get',
'notification.count.get',
'comment.changes.get',
'workspace.embedding.progress.get',
'copilot.transcript.task.get',
'user.quota-state.get',
'workspace.quota-state.get',
@@ -28,7 +27,6 @@ export const REALTIME_GATEWAY_REQUIRED_TOPICS = [
'user.settings.changed',
'notification.count.changed',
'comment.changed',
'workspace.embedding.progress.changed',
'copilot.transcript.task.changed',
'user.quota-state.changed',
'workspace.quota-state.changed',
packages/backend/server/src/core/realtime/types.ts
+7 -1
View File
@@ -10,9 +10,14 @@ import type { z } from 'zod';
import type { CurrentUser } from '../auth';
export type RealtimeRequestContext = {
connectionId?: string;
};
declare global {
interface Events {
'realtime.topic.changed': RealtimePublishPayload;
'realtime.connection.disconnected': { connectionId: string };
}
}
@@ -21,7 +26,8 @@ export type RealtimeRequestHandler<Op extends RealtimeRequestName> = {
input: z.ZodType<RealtimeRequestInputOf<Op>>;
handle(
user: CurrentUser,
input: RealtimeRequestInputOf<Op>
input: RealtimeRequestInputOf<Op>,
context?: RealtimeRequestContext
): Promise<RealtimeRequestOutputOf<Op>>;
};
packages/backend/server/src/core/storage/__tests__/blob-job.spec.ts
+1 -11
View File
@@ -289,7 +289,7 @@ test('storage reconciliation still refreshes document retention without object s
t.false(t.context.runtime.planUnreferencedWorkspaceBlobs.called);
});
test('document cleanup dispatches independent stable search and copilot effects', async t => {
test('document cleanup dispatches stable search effects', async t => {
t.context.runtime.executeDocumentCleanupCandidates.resolves({
scannedCandidates: 1,
serializationRetries: 0,
@@ -305,7 +305,6 @@ test('document cleanup dispatches independent stable search and copilot effects'
cleanupVersion: 'version-1',
commentObjectsDone: true,
searchDone: false,
copilotDone: false,
},
],
});
@@ -321,15 +320,6 @@ test('document cleanup dispatches independent stable search and copilot effects'
}
)
);
t.true(
t.context.queue.add.calledWith(
'copilot.embedding.reconcileDocumentCleanup',
Sinon.match({ docId: 'doc-1' }),
{
jobId: 'document-cleanup:copilot:workspace-1:doc-1:version-1',
}
)
);
t.true(
t.context.event.emitAsync.calledWith('workspace.blobs.updated', {
workspaceId: 'workspace-1',
packages/backend/server/src/core/storage/blob-job.ts
-9
View File
@@ -304,15 +304,6 @@ export class StorageBlobJob {
jobId: `document-cleanup:search:${effect.workspaceId}:${effect.docId}:${effect.cleanupVersion}`,
});
}
if (!effect.copilotDone) {
await this.queue.add(
'copilot.embedding.reconcileDocumentCleanup',
effect,
{
jobId: `document-cleanup:copilot:${effect.workspaceId}:${effect.docId}:${effect.cleanupVersion}`,
}
);
}
if (effect.commentObjectsDone) {
await this.event.emitAsync('workspace.blobs.updated', {
workspaceId: effect.workspaceId,
packages/backend/server/src/core/sync/gateway.ts
+19 -7
View File
@@ -24,7 +24,6 @@ import {
checkCanaryDateClientVersion,
DocNotFound,
DocUpdateBlocked,
EventBus,
GatewayErrorWrapper,
metrics,
NotInSpace,
@@ -32,6 +31,7 @@ import {
SpaceAccessDenied,
} from '../../base';
import { Models } from '../../models';
import { authorizeUserdataDocSubject } from '../../native';
import { CurrentUser } from '../auth';
import {
DocReader,
@@ -226,7 +226,6 @@ export class SpaceSyncGateway
constructor(
private readonly ac: PermissionAccess,
private readonly event: EventBus,
private readonly workspace: PgWorkspaceDocStorageAdapter,
private readonly userspace: PgUserspaceDocStorageAdapter,
private readonly docReader: DocReader,
@@ -332,6 +331,20 @@ export class SpaceSyncGateway
await this.ac.user(userId).doc(spaceId, docId).assert(action);
}
private assertUserdataSubject(
spaceType: SpaceType,
userId: string,
workspaceId: string,
docId: string
) {
if (
spaceType === SpaceType.Workspace &&
!authorizeUserdataDocSubject(userId, workspaceId, docId)
) {
throw new SpaceAccessDenied({ spaceId: workspaceId });
}
}
handleConnection(client: Socket) {
this.connectionCount++;
this.logger.debug(`New connection, total: ${this.connectionCount}`);
@@ -599,10 +612,6 @@ export class SpaceSyncGateway
return { data: { clientId: client.id, success: false } };
}
if (spaceType === SpaceType.Workspace) {
this.event.emit('workspace.embedding', { workspaceId: spaceId });
}
const adapter = this.selectAdapter(client, spaceType);
await adapter.join(user.id, spaceId);
@@ -644,6 +653,7 @@ export class SpaceSyncGateway
const id = new DocID(docId, spaceId);
const adapter = this.selectAdapter(client, spaceType);
adapter.assertIn(spaceId);
this.assertUserdataSubject(spaceType, user.id, spaceId, id.guid);
await this.assertDocActionAllowed(
spaceType,
user.id,
@@ -678,6 +688,7 @@ export class SpaceSyncGateway
@MessageBody() { spaceType, spaceId, docId }: DeleteDocMessage
): Promise<EventResponse<{ success: true }>> {
const adapter = this.selectAdapter(client, spaceType);
this.assertUserdataSubject(spaceType, user.id, spaceId, docId);
await this.assertDocActionAllowed(
spaceType,
user.id,
@@ -702,7 +713,8 @@ export class SpaceSyncGateway
const { spaceType, spaceId, docId, update } = message;
const adapter = this.selectAdapter(client, spaceType);
// Quota recovery mode is intentionally not applied to sync in this phase.
// Quota recovery mode is intentionally not applied to sync.
this.assertUserdataSubject(spaceType, user.id, spaceId, docId);
await this.assertDocActionAllowed(
spaceType,
user.id,
packages/backend/server/src/core/utils/__tests__/__snapshots__/blocksute.spec.ts.md
+1427 -1362
View File
File diff suppressed because it is too large Load Diff
packages/backend/server/src/core/utils/__tests__/__snapshots__/blocksute.spec.ts.snap
BIN
View File
Binary file not shown.
packages/backend/server/src/core/utils/__tests__/blocksute.spec.ts
+34 -41
View File
@@ -1,5 +1,4 @@
import test from 'ava';
import { omit } from 'lodash-es';
import * as Y from 'yjs';
import { createModule } from '../../../__tests__/create-module';
@@ -7,7 +6,8 @@ import { Mockers } from '../../../__tests__/mocks';
import { Models } from '../../../models';
import {
parseDocToMarkdownFromDocSnapshot,
readAllBlocksFromDocSnapshot,
projectDocCanvas,
projectDocSearch,
readAllDocIdsFromWorkspaceSnapshot,
} from '../blocksuite';
@@ -119,45 +119,6 @@ test('nested concurrent meta edits do not restore a deleted entry', t => {
);
});
test('can read all blocks from doc snapshot', async t => {
const rootDoc = await models.doc.get(workspace.id, workspace.id);
t.truthy(rootDoc);
const doc = await models.doc.get(workspace.id, docSnapshot.id);
t.truthy(doc);
const result = await readAllBlocksFromDocSnapshot('doc-0', docSnapshot.blob);
t.snapshot({
...result,
blocks: result!.blocks.map(block => omit(block, ['yblock'])),
});
});
test('can read blob filename from doc snapshot', async t => {
const docSnapshot = await module.create(Mockers.DocSnapshot, {
workspaceId: workspace.id,
user: owner,
snapshotFile: 'test-doc-with-blob.snapshot.bin',
});
const result = await readAllBlocksFromDocSnapshot('doc-0', docSnapshot.blob);
// NOTE: avoid snapshot result directly, because it will cause hanging
t.snapshot(JSON.parse(JSON.stringify(result)));
});
test('can read all blocks from doc snapshot without workspace snapshot', async t => {
const doc = await models.doc.get(workspace.id, docSnapshot.id);
t.truthy(doc);
const result = await readAllBlocksFromDocSnapshot('doc-0', docSnapshot.blob);
t.snapshot({
...result,
blocks: result!.blocks.map(block => omit(block, ['yblock'])),
});
});
test('can parse doc to markdown from doc snapshot', async t => {
const result = parseDocToMarkdownFromDocSnapshot(
workspace.id,
@@ -166,6 +127,38 @@ test('can parse doc to markdown from doc snapshot', async t => {
);
t.snapshot(result);
const canvas = projectDocCanvas(
docSnapshot.blob,
'fixture-doc',
'fixture-revision'
);
const search = projectDocSearch(
docSnapshot.blob,
'fixture-doc',
'fixture-revision'
);
t.snapshot(canvas, 'should export the exact canvas projection');
t.snapshot(search, 'should export the exact search projection');
t.deepEqual(
search,
projectDocSearch(docSnapshot.blob, 'fixture-doc', 'fixture-revision')
);
const blobSnapshot = await module.create(Mockers.DocSnapshot, {
workspaceId: workspace.id,
user: owner,
docId: 'fixture-blob-doc',
snapshotFile: 'test-doc-with-blob.snapshot.bin',
});
t.snapshot(
projectDocSearch(
blobSnapshot.blob,
blobSnapshot.id,
'fixture-blob-revision'
),
'should export attachment metadata from the blob fixture'
);
});
test('can parse doc to markdown from doc snapshot with ai editable', async t => {
packages/backend/server/src/core/utils/blocksuite.ts
+120 -23
View File
@@ -1,11 +1,115 @@
import { z } from 'zod';
import {
parsePageDocFromBinary,
parseWorkspaceDocFromBinary,
parseYDocFromBinary,
parseYDocToMarkdown,
projectDocCanvasFromBinary,
projectDocSearchFromBinary,
readAllDocIdsFromRootDoc,
} from '../../native';
const DocVisibilitySchema = z.enum(['page', 'edgeless', 'both']);
const DocBoundsSchema = z
.object({
x: z.number().finite(),
y: z.number().finite(),
width: z.number().finite().nonnegative(),
height: z.number().finite().nonnegative(),
})
.strict();
const ProjectionWarningSchema = z
.object({ code: z.string(), locator: z.string() })
.strict();
const CanvasProjectionBlockSchema = z
.object({
id: z.string(),
type: z.string(),
visibility: DocVisibilitySchema,
bounds: DocBoundsSchema.optional(),
text: z.string().optional(),
title: z.string().optional(),
childIds: z.array(z.string()),
})
.strict();
const CanvasProjectionElementSchema = z
.object({
id: z.string(),
type: z.string(),
bounds: DocBoundsSchema.optional(),
text: z.string().optional(),
title: z.string().optional(),
frameId: z.string().optional(),
childIds: z.array(z.string()),
sourceId: z.string().optional(),
targetId: z.string().optional(),
parentId: z.string().optional(),
index: z.string().optional(),
pointCount: z.number().int().nonnegative().optional(),
color: z.string().optional(),
lineWidth: z.number().finite().optional(),
})
.strict();
const CanvasProjectionV1Schema = z
.object({
version: z.literal(1),
docId: z.string(),
revision: z.string(),
title: z.string(),
surfaceBlockId: z.string().optional(),
bounds: DocBoundsSchema.optional(),
counts: z.record(z.string(), z.number().int().nonnegative()),
blocks: z.array(CanvasProjectionBlockSchema),
elements: z.array(CanvasProjectionElementSchema),
warnings: z.array(ProjectionWarningSchema),
})
.strict();
const DocumentSearchUnitV1Schema = z
.object({
unitId: z.string(),
source: z.enum(['page-block', 'canvas-block', 'surface-element']),
visibility: DocVisibilitySchema,
blockId: z.string().optional(),
elementId: z.string().optional(),
frameId: z.string().optional(),
blobId: z.string().optional(),
refDocIds: z.array(z.string()),
refs: z.array(z.string()),
parentFlavour: z.string().optional(),
parentBlockId: z.string().optional(),
additional: z.string().optional(),
type: z.string(),
text: z.string(),
})
.strict();
const DocumentSearchProjectionV1Schema = z
.object({
version: z.literal(1),
docId: z.string(),
revision: z.string(),
sourceHash: z.string(),
title: z.string(),
units: z.array(DocumentSearchUnitV1Schema),
warnings: z.array(ProjectionWarningSchema),
})
.strict();
export type DocVisibility = z.infer<typeof DocVisibilitySchema>;
export type DocBounds = z.infer<typeof DocBoundsSchema>;
export type ProjectionWarning = z.infer<typeof ProjectionWarningSchema>;
export type CanvasProjectionBlock = z.infer<typeof CanvasProjectionBlockSchema>;
export type CanvasProjectionElement = z.infer<
typeof CanvasProjectionElementSchema
>;
export type CanvasProjectionV1 = z.infer<typeof CanvasProjectionV1Schema>;
export type DocumentSearchUnitV1 = z.infer<typeof DocumentSearchUnitV1Schema>;
export type DocumentSearchProjectionV1 = z.infer<
typeof DocumentSearchProjectionV1Schema
>;
export const parseCanvasProjection = (value: unknown) =>
CanvasProjectionV1Schema.parse(value);
export interface PageDocContent {
title: string;
summary: string;
@@ -52,31 +156,24 @@ export function readAllDocIdsFromWorkspaceSnapshot(
return readAllDocIdsFromRootDoc(Buffer.from(snapshot), includeTrash);
}
function safeParseJson<T>(str: string): T | undefined {
try {
return JSON.parse(str) as T;
} catch {
return undefined;
}
export function projectDocCanvas(
docSnapshot: Uint8Array,
docId: string,
revision: string
): CanvasProjectionV1 {
return parseCanvasProjection(
projectDocCanvasFromBinary(Buffer.from(docSnapshot), docId, revision)
);
}
export async function readAllBlocksFromDocSnapshot(
export function projectDocSearch(
docSnapshot: Uint8Array,
docId: string,
docSnapshot: Uint8Array
) {
const result = parseYDocFromBinary(Buffer.from(docSnapshot), docId);
return {
...result,
blocks: result.blocks.map(block => ({
...block,
docId,
ref: block.refInfo,
additional: block.additional
? safeParseJson(block.additional)
: undefined,
})),
};
revision: string
): DocumentSearchProjectionV1 {
return DocumentSearchProjectionV1Schema.parse(
projectDocSearchFromBinary(Buffer.from(docSnapshot), docId, revision)
);
}
export function parseDocToMarkdownFromDocSnapshot(
packages/backend/server/src/models/common/copilot.ts
+22 -101
View File
@@ -1,6 +1,5 @@
import { AiJobStatus, AiJobType } from '@prisma/client';
import type { JsonValue } from '@prisma/client/runtime/library';
import { z } from 'zod';
export interface CopilotJob {
id?: string;
@@ -12,83 +11,6 @@ export interface CopilotJob {
payload?: JsonValue;
}
export interface CopilotContext {
id?: string;
sessionId: string;
config: JsonValue;
createdAt: Date;
updatedAt: Date;
}
export enum ContextEmbedStatus {
processing = 'processing',
finished = 'finished',
failed = 'failed',
}
export enum ContextCategories {
Tag = 'tag',
Collection = 'collection',
}
const ContextEmbedStatusSchema = z.enum([
ContextEmbedStatus.processing,
ContextEmbedStatus.finished,
ContextEmbedStatus.failed,
]);
const ContextBlobSchema = z.object({
id: z.string(),
createdAt: z.number(),
});
const ContextDocSchema = z.object({
id: z.string(),
createdAt: z.number(),
});
export const ContextFileSchema = z.object({
id: z.string(),
chunkSize: z.number(),
name: z.string(),
mimeType: z.string().optional(),
status: ContextEmbedStatusSchema,
error: z.string().nullable(),
blobId: z.string(),
createdAt: z.number(),
});
export const ContextCategorySchema = z.object({
id: z.string(),
type: z.enum([ContextCategories.Tag, ContextCategories.Collection]),
docs: ContextDocSchema.merge(
z.object({ status: ContextEmbedStatusSchema })
).array(),
createdAt: z.number(),
});
export const ContextConfigSchema = z.object({
workspaceId: z.string(),
blobs: ContextBlobSchema.merge(
z.object({ status: ContextEmbedStatusSchema.optional() })
).array(),
files: ContextFileSchema.array(),
docs: ContextDocSchema.merge(
z.object({ status: ContextEmbedStatusSchema.optional() })
).array(),
categories: ContextCategorySchema.array(),
});
export const MinimalContextConfigSchema = ContextConfigSchema.pick({
workspaceId: true,
});
export type ContextCategory = z.infer<typeof ContextCategorySchema>;
export type ContextConfig = z.infer<typeof ContextConfigSchema>;
export type ContextBlob = z.infer<typeof ContextConfigSchema>['blobs'][number];
export type ContextDoc = z.infer<typeof ContextConfigSchema>['docs'][number];
export type ContextFile = z.infer<typeof ContextConfigSchema>['files'][number];
// embeddings
export type Embedding = {
@@ -100,40 +22,39 @@ export type Embedding = {
embedding: Array<number>;
};
export type DocumentEmbedding = Embedding & {
projectionVersion: number;
sourceHash: string;
unitId: string;
visibility: 'page' | 'edgeless' | 'both';
blockId?: string;
elementId?: string;
frameId?: string;
};
export type ChunkSimilarity = {
chunk: number;
content: string;
distance: number | null;
};
export type FileChunkSimilarity = ChunkSimilarity & {
fileId: string;
blobId: string;
name: string;
mimeType: string;
};
export type BlobChunkSimilarity = ChunkSimilarity & {
blobId: string;
};
export type DocChunkSimilarity = ChunkSimilarity & {
docId: string;
unitId: string;
visibility: 'page' | 'edgeless' | 'both';
blockId?: string;
elementId?: string;
frameId?: string;
};
export const CopilotWorkspaceFileSchema = z.object({
fileName: z.string(),
blobId: z.string(),
mimeType: z.string(),
size: z.number(),
});
export type CopilotWorkspaceFileMetadata = z.infer<
typeof CopilotWorkspaceFileSchema
>;
export type CopilotWorkspaceFile = CopilotWorkspaceFileMetadata & {
export type CopilotWorkspaceArtifact = {
workspaceId: string;
fileId: string;
artifactId: string;
contentHash: string;
fileName: string;
embeddingStatus: 'processing' | 'ready' | 'failed';
mediaType: string;
size: number;
createdAt: Date;
};
packages/backend/server/src/models/copilot-context.ts
-378
View File
@@ -1,378 +0,0 @@
import { randomUUID } from 'node:crypto';
import { Injectable } from '@nestjs/common';
import { Prisma } from '@prisma/client';
import { CopilotSessionNotFound } from '../base';
import { BaseModel } from './base';
import {
clearEmbeddingContent,
ContextBlob,
ContextConfigSchema,
ContextDoc,
ContextEmbedStatus,
CopilotContext,
DocChunkSimilarity,
Embedding,
EMBEDDING_DIMENSIONS,
FileChunkSimilarity,
MinimalContextConfigSchema,
} from './common/copilot';
type UpdateCopilotContextInput = Pick<CopilotContext, 'config'>;
/**
* Copilot Job Model
*/
@Injectable()
export class CopilotContextModel extends BaseModel {
// ================ contexts ================
async create(sessionId: string) {
const session = await this.db.aiSession.findFirst({
where: { id: sessionId },
select: { workspaceId: true },
});
if (!session) {
throw new CopilotSessionNotFound();
}
const row = await this.db.aiContext.create({
data: {
sessionId,
config: {
workspaceId: session.workspaceId,
blobs: [],
docs: [],
files: [],
categories: [],
},
},
});
return row;
}
async get(id: string) {
const row = await this.db.aiContext.findFirst({
where: { id },
});
return row;
}
async getAccessInfo(id: string) {
return await this.db.aiContext.findFirst({
where: { id },
select: {
id: true,
sessionId: true,
session: {
select: {
userId: true,
workspaceId: true,
},
},
},
});
}
async getConfig(id: string) {
const row = await this.get(id);
if (row) {
const config = ContextConfigSchema.safeParse(row.config);
if (config.success) {
return config.data;
}
const minimalConfig = MinimalContextConfigSchema.safeParse(row.config);
if (minimalConfig.success) {
// fulfill the missing fields
return {
blobs: [],
docs: [],
files: [],
categories: [],
...minimalConfig.data,
};
}
}
return null;
}
async getBySessionId(sessionId: string) {
const row = await this.db.aiContext.findFirst({
where: { sessionId },
});
return row;
}
async mergeBlobStatus(
workspaceId: string,
blobs: ContextBlob[]
): Promise<ContextBlob[]> {
const canEmbedding = await this.checkEmbeddingAvailable();
const finishedBlobs = canEmbedding
? await this.listWorkspaceBlobEmbedding(
workspaceId,
Array.from(new Set(blobs.map(blob => blob.id)))
)
: [];
const finishedBlobSet = new Set(finishedBlobs);
for (const blob of blobs) {
const status = finishedBlobSet.has(blob.id)
? ContextEmbedStatus.finished
: undefined;
// NOTE: when the blob has not been synchronized to the server or is in the embedding queue
// the status will be empty, fallback to processing if no status is provided
blob.status = status || blob.status || ContextEmbedStatus.processing;
}
return blobs;
}
async mergeDocStatus(workspaceId: string, docs: ContextDoc[]) {
const canEmbedding = await this.checkEmbeddingAvailable();
const finishedDoc = canEmbedding
? await this.listWorkspaceDocEmbedding(
workspaceId,
Array.from(new Set(docs.map(doc => doc.id)))
)
: [];
const finishedDocSet = new Set(finishedDoc);
for (const doc of docs) {
const status = finishedDocSet.has(doc.id)
? ContextEmbedStatus.finished
: undefined;
// NOTE: when the document has not been synchronized to the server or is in the embedding queue
// the status will be empty, fallback to processing if no status is provided
doc.status = status || doc.status || ContextEmbedStatus.processing;
}
return docs;
}
async update(contextId: string, data: UpdateCopilotContextInput) {
const ret = await this.db.aiContext.updateMany({
where: {
id: contextId,
},
data: {
config: data.config || undefined,
},
});
return ret.count > 0;
}
// ================ embeddings ================
async checkEmbeddingAvailable(): Promise<boolean> {
const [{ count }] = await this.db.$queryRaw<
{ count: number }[]
>`SELECT count(1) FROM pg_tables WHERE tablename in ('ai_context_embeddings', 'ai_workspace_embeddings')`;
return Number(count) === 2;
}
async listWorkspaceBlobEmbedding(
workspaceId: string,
blobIds?: string[]
): Promise<string[]> {
const existsIds = await this.db.aiWorkspaceBlobEmbedding
.groupBy({
where: {
workspaceId,
blobId: blobIds ? { in: blobIds } : undefined,
},
by: ['blobId'],
})
.then(r => r.map(r => r.blobId));
return existsIds;
}
async listWorkspaceDocEmbedding(workspaceId: string, docIds?: string[]) {
const existsIds = await this.db.aiWorkspaceEmbedding
.groupBy({
where: {
workspaceId,
docId: docIds ? { in: docIds } : undefined,
},
by: ['docId'],
})
.then(r => r.map(r => r.docId));
return existsIds;
}
private processEmbeddings(
contextOrWorkspaceId: string,
fileOrDocId: string,
embeddings: Embedding[],
withId = true
) {
const groups = embeddings.map(e =>
[
withId ? randomUUID() : undefined,
contextOrWorkspaceId,
fileOrDocId,
e.index,
e.content,
Prisma.raw(`'[${e.embedding.join(',')}]'`),
new Date(),
].filter(v => v !== undefined)
);
return Prisma.join(groups.map(row => Prisma.sql`(${Prisma.join(row)})`));
}
async getFileContent(
contextId: string,
fileId: string,
chunk?: number
): Promise<string | undefined> {
const file = await this.db.aiContextEmbedding.findMany({
where: { contextId, fileId, chunk },
select: { content: true },
orderBy: { chunk: 'asc' },
});
return file?.map(f => clearEmbeddingContent(f.content)).join('\n');
}
async insertFileEmbedding(
contextId: string,
fileId: string,
embeddings: Embedding[]
) {
if (embeddings.length === 0) {
this.logger.warn(
`No embeddings provided for contextId: ${contextId}, fileId: ${fileId}. Skipping insertion.`
);
return;
}
const values = this.processEmbeddings(contextId, fileId, embeddings);
await this.db.$executeRaw`
INSERT INTO "ai_context_embeddings"
("id", "context_id", "file_id", "chunk", "content", "embedding", "updated_at") VALUES ${values}
ON CONFLICT (context_id, file_id, chunk) DO UPDATE SET
content = EXCLUDED.content, embedding = EXCLUDED.embedding, updated_at = excluded.updated_at;
`;
}
async deleteFileEmbedding(contextId: string, fileId: string) {
await this.db.aiContextEmbedding.deleteMany({
where: { contextId, fileId },
});
}
async matchFileEmbedding(
embedding: number[],
contextId: string,
topK: number,
threshold: number
): Promise<Omit<FileChunkSimilarity, 'blobId' | 'name' | 'mimeType'>[]> {
const similarityChunks = await this.db.$queryRaw<
Array<Omit<FileChunkSimilarity, 'blobId' | 'name' | 'mimeType'>>
>`
SELECT "file_id" as "fileId", "chunk", "content", "embedding" <=> ${embedding}::vector as "distance"
FROM "ai_context_embeddings"
WHERE context_id = ${contextId}
ORDER BY "distance" ASC
LIMIT ${topK};
`;
return similarityChunks.filter(c => Number(c.distance) <= threshold);
}
async getWorkspaceContent(
workspaceId: string,
docId: string,
chunk?: number
): Promise<string | undefined> {
const file = await this.db.aiWorkspaceEmbedding.findMany({
where: { workspaceId, docId, chunk },
select: { content: true },
orderBy: { chunk: 'asc' },
});
return file?.map(f => clearEmbeddingContent(f.content)).join('\n');
}
async insertWorkspaceEmbedding(
workspaceId: string,
docId: string,
embeddings: Embedding[]
) {
if (embeddings.length === 0) {
this.logger.warn(
`No embeddings provided for workspaceId: ${workspaceId}, docId: ${docId}. Skipping insertion.`
);
return;
}
const values = this.processEmbeddings(
workspaceId,
docId,
embeddings,
false
);
await this.db.$executeRaw`
INSERT INTO "ai_workspace_embeddings"
("workspace_id", "doc_id", "chunk", "content", "embedding", "updated_at")
VALUES ${values}
ON CONFLICT (workspace_id, doc_id, chunk)
DO UPDATE SET
content = EXCLUDED.content,
embedding = EXCLUDED.embedding,
updated_at = excluded.updated_at;
`;
}
async fulfillEmptyEmbedding(workspaceId: string, docId: string) {
const emptyEmbedding = {
index: 0,
content: '',
embedding: Array.from({ length: EMBEDDING_DIMENSIONS }, () => 0),
};
await this.models.copilotContext.insertWorkspaceEmbedding(
workspaceId,
docId,
[emptyEmbedding]
);
}
async deleteWorkspaceEmbedding(workspaceId: string, docId: string) {
await this.purgeWorkspaceEmbedding(workspaceId, docId);
await this.fulfillEmptyEmbedding(workspaceId, docId);
}
async purgeWorkspaceEmbedding(workspaceId: string, docId: string) {
await this.db.aiWorkspaceEmbedding.deleteMany({
where: { workspaceId, docId },
});
}
async matchWorkspaceEmbedding(
embedding: number[],
workspaceId: string,
topK: number,
threshold: number,
matchDocIds?: string[]
): Promise<DocChunkSimilarity[]> {
const similarityChunks = await this.db.$queryRaw<Array<DocChunkSimilarity>>`
SELECT
w."doc_id" as "docId",
w."chunk",
w."content",
w."embedding" <=> ${embedding}::vector as "distance"
FROM "ai_workspace_embeddings" w
LEFT JOIN "ai_workspace_ignored_docs" i
ON i."workspace_id" = w."workspace_id"
AND i."doc_id" = w."doc_id"
${matchDocIds?.length ? Prisma.sql`AND w."doc_id" NOT IN (${Prisma.join(matchDocIds)})` : Prisma.empty}
WHERE
w."workspace_id" = ${workspaceId}
AND i."doc_id" IS NULL
AND (w."embedding" <=> ${embedding}::vector) <= ${threshold}
ORDER BY "distance" ASC
LIMIT ${topK};
`;
return similarityChunks;
}
}
packages/backend/server/src/models/copilot-session.ts
+58 -2
View File
@@ -10,6 +10,10 @@ import {
CopilotSessionNotFound,
} from '../base';
import type { PromptAttachment } from '../plugins/copilot/providers/types';
import type {
SessionFocus,
TurnScopeSnapshot,
} from '../plugins/copilot/runtime/contracts/shared';
import {
type ChatMessage as CopilotChatMessage,
ChatMessageSchema,
@@ -48,6 +52,7 @@ type ChatMessage = {
content: string;
attachments?: ChatAttachment[] | null;
params?: Record<string, any> | null;
scopeSnapshot?: TurnScopeSnapshot | null;
streamObjects?: ChatStreamObject[] | null;
createdAt: Date;
};
@@ -61,6 +66,7 @@ type StoredChatMessage = Prisma.AiSessionMessageGetPayload<{
attachments: true;
streamObjects: true;
params: true;
scopeSnapshot: true;
createdAt: true;
};
}>;
@@ -317,6 +323,7 @@ export class CopilotSessionModel extends BaseModel {
params: this.sanitizeJsonValue(
omit(message.params, ['docs']) || undefined
),
scopeSnapshot: this.sanitizeJsonValue(message.scopeSnapshot),
streamObjects: message.streamObjects?.map(o =>
this.sanitizeStreamObject(o)
),
@@ -488,6 +495,7 @@ export class CopilotSessionModel extends BaseModel {
parentSessionId: true,
pinned: true,
title: true,
focus: true,
promptName: true,
createdAt: true,
updatedAt: true,
@@ -499,6 +507,7 @@ export class CopilotSessionModel extends BaseModel {
attachments: true,
streamObjects: true,
params: true,
scopeSnapshot: true,
createdAt: true,
},
orderBy: { createdAt: 'asc' },
@@ -516,6 +525,7 @@ export class CopilotSessionModel extends BaseModel {
parentSessionId: true,
pinned: true,
title: true,
focus: true,
promptName: true,
createdAt: true,
updatedAt: true,
@@ -584,6 +594,7 @@ export class CopilotSessionModel extends BaseModel {
parentSessionId: true,
pinned: true,
title: true,
focus: true,
promptName: true,
createdAt: true,
updatedAt: true,
@@ -596,6 +607,7 @@ export class CopilotSessionModel extends BaseModel {
attachments: true,
streamObjects: true,
params: true,
scopeSnapshot: true,
createdAt: true,
},
orderBy: {
@@ -744,6 +756,7 @@ export class CopilotSessionModel extends BaseModel {
attachments: true,
streamObjects: true,
params: true,
scopeSnapshot: true,
createdAt: true,
},
});
@@ -766,6 +779,7 @@ export class CopilotSessionModel extends BaseModel {
attachments: true,
streamObjects: true,
params: true,
scopeSnapshot: true,
createdAt: true,
},
orderBy: { createdAt: 'asc' },
@@ -791,6 +805,7 @@ export class CopilotSessionModel extends BaseModel {
content: m.content,
attachments: m.attachments || undefined,
params: m.params || undefined,
scopeSnapshot: m.scopeSnapshot || undefined,
streamObjects: m.streamObjects || undefined,
createdAt: m.createdAt,
sessionId,
@@ -813,13 +828,32 @@ export class CopilotSessionModel extends BaseModel {
sessionId: string;
userId: string;
message: ChatMessage;
focus?: SessionFocus;
artifacts?: Array<{
artifactId: string;
role: string;
displayName?: string;
metadata?: Record<string, unknown>;
}>;
}) {
const haveSession = await this.has(state.sessionId, state.userId);
if (!haveSession) {
const session = await this.getExists(
state.sessionId,
{ id: true, workspaceId: true },
{ userId: state.userId }
);
if (!session) {
throw new CopilotSessionNotFound();
}
const message = this.sanitizeMessage(state.message);
const artifacts = [];
const artifactKeys = new Set<string>();
for (const artifact of state.artifacts ?? []) {
const key = `${artifact.artifactId}:${artifact.role}`;
if (artifactKeys.has(key)) continue;
artifactKeys.add(key);
artifacts.push(artifact);
}
const created = await this.db.aiSessionMessage.create({
data: {
sessionId: state.sessionId,
@@ -828,8 +862,28 @@ export class CopilotSessionModel extends BaseModel {
content: message.content,
attachments: message.attachments || undefined,
params: message.params || undefined,
scopeSnapshot: message.scopeSnapshot || undefined,
streamObjects: message.streamObjects || undefined,
createdAt: message.createdAt,
artifacts: artifacts.length
? {
create: artifacts.map(artifact => ({
role: artifact.role,
displayName: this.sanitizeString(artifact.displayName),
metadata: this.sanitizeJsonValue(artifact.metadata) as
| Prisma.InputJsonObject
| undefined,
artifact: {
connect: {
workspaceId_id: {
workspaceId: session.workspaceId,
id: artifact.artifactId,
},
},
},
})),
}
: undefined,
},
select: {
id: true,
@@ -839,6 +893,7 @@ export class CopilotSessionModel extends BaseModel {
attachments: true,
streamObjects: true,
params: true,
scopeSnapshot: true,
createdAt: true,
},
});
@@ -850,6 +905,7 @@ export class CopilotSessionModel extends BaseModel {
message.role === AiSessionMessageRole.user
? { increment: 1 }
: undefined,
focus: state.focus,
},
});
packages/backend/server/src/models/copilot-workspace.ts
+21 -439
View File
@@ -1,75 +1,26 @@
import { randomUUID } from 'node:crypto';
import { Injectable } from '@nestjs/common';
import { Transactional } from '@nestjs-cls/transactional';
import { Prisma, PrismaClient } from '@prisma/client';
import { PaginationInput } from '../base';
import { BaseModel } from './base';
import {
type BlobChunkSimilarity,
clearEmbeddingContent,
type CopilotWorkspaceFile,
type CopilotWorkspaceFileMetadata,
type Embedding,
type FileChunkSimilarity,
type IgnoredDoc,
} from './common';
import type { IgnoredDoc } from './common';
@Injectable()
export class CopilotWorkspaceConfigModel extends BaseModel {
constructor(private readonly database: PrismaClient) {
super();
}
@Transactional()
private async listIgnoredDocIds(
workspaceId: string,
options?: PaginationInput
) {
return await this.db.aiWorkspaceIgnoredDocs.findMany({
where: {
workspaceId,
},
select: {
docId: true,
createdAt: true,
},
where: { workspaceId },
select: { docId: true, createdAt: true },
orderBy: { createdAt: 'desc' },
skip: options?.offset,
take: options?.first,
});
}
/**
* find docs to embed, excluding ignored and already embedded docs
* newer docs will be list first
* @param workspaceId id of the workspace
* @returns docIds
*/
async findDocsToEmbed(workspaceId: string): Promise<string[]> {
// NOTE: for unknown reason, the transaction will timeout if call from event handler
// so we use an independent client here
const docIds = await this.database.$queryRaw<{ id: string }[]>`
SELECT s.guid as id
FROM snapshots AS s
LEFT JOIN ai_workspace_embeddings e
ON e.workspace_id = s.workspace_id
AND e.doc_id = s.guid
LEFT JOIN ai_workspace_ignored_docs id
ON id.workspace_id = s.workspace_id
AND id.doc_id = s.guid
WHERE s.workspace_id = ${workspaceId}
AND s.guid <> s.workspace_id
AND s.guid NOT LIKE '%$%'
AND s.guid NOT LIKE '%:settings:%'
AND e.doc_id IS NULL
AND id.doc_id IS NULL
AND s.blob <> E'\\\\x0000';`;
return docIds.map(r => r.id);
}
@Transactional()
async updateIgnoredDocs(
workspaceId: string,
@@ -78,28 +29,17 @@ export class CopilotWorkspaceConfigModel extends BaseModel {
) {
const removed = new Set(remove);
const ignored = await this.listIgnoredDocIds(workspaceId).then(
r => new Set(r.map(r => r.docId).filter(id => !removed.has(id)))
rows => new Set(rows.map(row => row.docId).filter(id => !removed.has(id)))
);
const added = add.filter(id => !ignored.has(id));
const { count: addedCount } =
await this.db.aiWorkspaceIgnoredDocs.createMany({
data: added.map(docId => ({
workspaceId,
docId,
})),
data: added.map(docId => ({ workspaceId, docId })),
});
const { count: removedCount } =
await this.db.aiWorkspaceIgnoredDocs.deleteMany({
where: {
workspaceId,
docId: {
in: Array.from(removed),
},
},
where: { workspaceId, docId: { in: Array.from(removed) } },
});
return addedCount + removedCount;
}
@@ -108,22 +48,25 @@ export class CopilotWorkspaceConfigModel extends BaseModel {
workspaceId: string,
options?: PaginationInput
): Promise<IgnoredDoc[]> {
const row = await this.listIgnoredDocIds(workspaceId, options);
const ids = row.map(r => ({ workspaceId, docId: r.docId }));
const rows = await this.listIgnoredDocIds(workspaceId, options);
const ids = rows.map(row => ({ workspaceId, docId: row.docId }));
const docs = await this.models.doc.findMetas(ids);
const docsMap = new Map(
docs.filter(r => !!r).map(r => [`${r.workspaceId}-${r.docId}`, r])
docs.flatMap(doc =>
doc ? [[`${doc.workspaceId}-${doc.docId}`, doc] as const] : []
)
);
const authors = await this.models.doc.findAuthors(ids);
const authorsMap = new Map(
authors.filter(r => !!r).map(r => [`${r.workspaceId}-${r.id}`, r])
authors.flatMap(author =>
author ? [[`${author.workspaceId}-${author.id}`, author] as const] : []
)
);
return row.map(r => {
const docMeta = docsMap.get(`${workspaceId}-${r.docId}`);
const docAuthor = authorsMap.get(`${workspaceId}-${r.docId}`);
return rows.map(row => {
const docMeta = docsMap.get(`${workspaceId}-${row.docId}`);
const docAuthor = authorsMap.get(`${workspaceId}-${row.docId}`);
return {
...r,
...row,
docCreatedAt: docAuthor?.createdAt,
docUpdatedAt: docAuthor?.updatedAt,
title: docMeta?.title || undefined,
@@ -136,377 +79,16 @@ export class CopilotWorkspaceConfigModel extends BaseModel {
@Transactional()
async countIgnoredDocs(workspaceId: string): Promise<number> {
const count = await this.db.aiWorkspaceIgnoredDocs.count({
where: {
workspaceId,
},
return await this.db.aiWorkspaceIgnoredDocs.count({
where: { workspaceId },
});
return count;
}
@Transactional()
async checkIgnoredDocs(workspaceId: string, docIds: string[]) {
const ignored = await this.listIgnoredDocIds(workspaceId).then(
r => new Set(r.map(r => r.docId))
rows => new Set(rows.map(row => row.docId))
);
return docIds.filter(id => ignored.has(id));
}
// check if a docId has only placeholder embeddings
@Transactional()
async hasPlaceholder(workspaceId: string, docId: string): Promise<boolean> {
const [total, nonPlaceholder] = await Promise.all([
this.db.aiWorkspaceEmbedding.count({ where: { workspaceId, docId } }),
this.db.aiWorkspaceEmbedding.count({
where: {
workspaceId,
docId,
NOT: { AND: [{ chunk: 0 }, { content: '' }] },
},
}),
]);
return total > 0 && nonPlaceholder === 0;
}
private getEmbeddableCondition(
workspaceId: string,
ignoredDocIds?: string[]
): Prisma.SnapshotWhereInput {
const condition: Prisma.SnapshotWhereInput['AND'] = [
{ id: { not: workspaceId } },
{ id: { not: { contains: '$' } } },
{ id: { not: { contains: ':settings:' } } },
{ blob: { not: new Uint8Array([0, 0]) } },
];
if (ignoredDocIds && ignoredDocIds.length > 0) {
condition.push({ id: { notIn: ignoredDocIds } });
}
return { workspaceId, AND: condition };
}
async listEmbeddableDocIds(workspaceId: string) {
const condition = this.getEmbeddableCondition(workspaceId);
const rows = await this.db.snapshot.findMany({
where: condition,
select: { id: true },
});
return rows.map(r => r.id);
}
@Transactional()
async getEmbeddingStatus(workspaceId: string) {
const ignoredDocIds = (await this.listIgnoredDocIds(workspaceId)).map(
d => d.docId
);
const snapshotCondition = this.getEmbeddableCondition(
workspaceId,
ignoredDocIds
);
const [docTotal, docEmbedded, fileTotal, fileEmbedded] = await Promise.all([
this.db.snapshot.findMany({
where: snapshotCondition,
select: { id: true },
}),
this.db.snapshot.findMany({
where: { ...snapshotCondition, embedding: { some: {} } },
select: { id: true },
}),
this.db.aiWorkspaceFiles.count({ where: { workspaceId } }),
this.db.aiWorkspaceFiles.count({
where: { workspaceId, embeddings: { some: {} } },
}),
]);
const docTotalIds = docTotal.map(d => d.id);
const docTotalSet = new Set(docTotalIds);
const outdatedDocPrefix = `${workspaceId}:space:`;
const duplicateOutdatedDocSet = new Set(
docTotalIds
.filter(id => id.startsWith(outdatedDocPrefix))
.filter(id => docTotalSet.has(id.slice(outdatedDocPrefix.length)))
);
return {
total:
docTotalIds.filter(id => !duplicateOutdatedDocSet.has(id)).length +
fileTotal,
embedded:
docEmbedded
.map(d => d.id)
.filter(id => !duplicateOutdatedDocSet.has(id)).length + fileEmbedded,
};
}
@Transactional()
async checkDocNeedEmbedded(workspaceId: string, docId: string) {
// NOTE: check if the document needs re-embedding.
// 1. first-time embedding when no embedding exists
// 2. re-embedding only when the doc has updates newer than the last embedding
// AND the last embedding is older than 10 minutes (avoid frequent updates)
const result = await this.db.$queryRaw<{ needs_embedding: boolean }[]>`
SELECT
EXISTS (
WITH docs AS (
SELECT
s.workspace_id,
s.guid AS doc_id,
s.updated_at
FROM
snapshots s
WHERE
s.workspace_id = ${workspaceId}
AND s.guid = ${docId}
UNION
ALL
SELECT
u.workspace_id,
u.guid AS doc_id,
u.created_at AS updated_at
FROM
"updates" u
WHERE
u.workspace_id = ${workspaceId}
AND u.guid = ${docId}
)
SELECT
1
FROM
docs
LEFT JOIN ai_workspace_embeddings e
ON e.workspace_id = docs.workspace_id
AND e.doc_id = docs.doc_id
WHERE
e.updated_at IS NULL
OR (docs.updated_at > e.updated_at AND e.updated_at < NOW() - INTERVAL '10 minutes')
) AS needs_embedding;
`;
return result[0]?.needs_embedding ?? false;
}
// ================ embeddings ================
async checkEmbeddingAvailable(): Promise<boolean> {
const [{ count }] = await this.db.$queryRaw<
{ count: number }[]
>`SELECT count(1) FROM pg_tables WHERE tablename in ('ai_workspace_embeddings', 'ai_workspace_file_embeddings', 'ai_workspace_blob_embeddings')`;
return Number(count) === 3;
}
private processEmbeddings(
workspaceId: string,
fileOrBlobId: string,
embeddings: Embedding[]
) {
const groups = embeddings.map(e =>
[
workspaceId,
fileOrBlobId,
e.index,
e.content,
Prisma.raw(`'[${e.embedding.join(',')}]'`),
].filter(v => v !== undefined)
);
return Prisma.join(groups.map(row => Prisma.sql`(${Prisma.join(row)})`));
}
async addFile(
workspaceId: string,
file: CopilotWorkspaceFileMetadata
): Promise<CopilotWorkspaceFile> {
const fileId = randomUUID();
const row = await this.db.aiWorkspaceFiles.create({
data: { ...file, workspaceId, fileId },
});
return row;
}
async getFile(workspaceId: string, fileId: string) {
const file = await this.db.aiWorkspaceFiles.findFirst({
where: {
workspaceId,
fileId,
},
});
return file;
}
@Transactional()
async insertFileEmbeddings(
workspaceId: string,
fileId: string,
embeddings: Embedding[]
) {
if (embeddings.length === 0) {
this.logger.warn(
`No embeddings provided for workspaceId: ${workspaceId}, fileId: ${fileId}. Skipping insertion.`
);
return;
}
const values = this.processEmbeddings(workspaceId, fileId, embeddings);
await this.db.$executeRaw`
INSERT INTO "ai_workspace_file_embeddings"
("workspace_id", "file_id", "chunk", "content", "embedding") VALUES ${values}
ON CONFLICT (workspace_id, file_id, chunk) DO NOTHING;
`;
}
async listFiles(
workspaceId: string,
options?: {
includeRead?: boolean;
} & PaginationInput
): Promise<CopilotWorkspaceFile[]> {
const files = await this.db.aiWorkspaceFiles.findMany({
where: {
workspaceId,
},
orderBy: { createdAt: 'desc' },
skip: options?.offset,
take: options?.first,
});
return files;
}
async countFiles(workspaceId: string): Promise<number> {
const count = await this.db.aiWorkspaceFiles.count({
where: {
workspaceId,
},
});
return count;
}
async matchFileEmbedding(
workspaceId: string,
embedding: number[],
topK: number,
threshold: number
): Promise<FileChunkSimilarity[]> {
if (!(await this.allowEmbedding(workspaceId))) {
return [];
}
const similarityChunks = await this.db.$queryRaw<
Array<FileChunkSimilarity>
>`
SELECT
e."file_id" as "fileId",
f."file_name" as "name",
f."blob_id" as "blobId",
f."mime_type" as "mimeType",
e."chunk",
e."content",
e."embedding" <=> ${embedding}::vector as "distance"
FROM "ai_workspace_file_embeddings" e
JOIN "ai_workspace_files" f
ON e."workspace_id" = f."workspace_id"
AND e."file_id" = f."file_id"
WHERE e.workspace_id = ${workspaceId}
ORDER BY "distance" ASC
LIMIT ${topK};
`;
return similarityChunks.filter(c => Number(c.distance) <= threshold);
}
async getBlobContent(
workspaceId: string,
blobId: string,
chunk?: number
): Promise<string | undefined> {
const blob = await this.db.aiWorkspaceBlobEmbedding.findMany({
where: { workspaceId, blobId, chunk },
select: { content: true },
orderBy: { chunk: 'asc' },
});
return blob?.map(f => clearEmbeddingContent(f.content)).join('\n');
}
async getBlobChunkSizes(workspaceId: string, blobIds: string[]) {
const sizes = await this.db.aiWorkspaceBlobEmbedding.groupBy({
by: ['blobId'],
_count: { chunk: true },
where: { workspaceId, blobId: { in: blobIds } },
});
return sizes.reduce((acc, cur) => {
if (cur._count.chunk) {
acc.set(cur.blobId, cur._count.chunk);
}
return acc;
}, new Map<string, number>());
}
@Transactional()
async insertBlobEmbeddings(
workspaceId: string,
blobId: string,
embeddings: Embedding[]
) {
if (embeddings.length === 0) {
this.logger.warn(
`No embeddings provided for workspaceId: ${workspaceId}, blobId: ${blobId}. Skipping insertion.`
);
return;
}
const values = this.processEmbeddings(workspaceId, blobId, embeddings);
await this.db.$executeRaw`
INSERT INTO "ai_workspace_blob_embeddings"
("workspace_id", "blob_id", "chunk", "content", "embedding") VALUES ${values}
ON CONFLICT (workspace_id, blob_id, chunk) DO NOTHING;
`;
}
async matchBlobEmbedding(
workspaceId: string,
embedding: number[],
topK: number,
threshold: number
): Promise<BlobChunkSimilarity[]> {
if (!(await this.allowEmbedding(workspaceId))) {
return [];
}
const similarityChunks = await this.db.$queryRaw<
Array<BlobChunkSimilarity>
>`
SELECT
e."blob_id" as "blobId",
e."chunk",
e."content",
e."embedding" <=> ${embedding}::vector as "distance"
FROM "ai_workspace_blob_embeddings" e
WHERE e.workspace_id = ${workspaceId}
ORDER BY "distance" ASC
LIMIT ${topK};
`;
return similarityChunks.filter(c => Number(c.distance) <= threshold);
}
async removeBlob(workspaceId: string, blobId: string) {
await this.db.$executeRaw`
DELETE FROM "ai_workspace_blob_embeddings"
WHERE workspace_id = ${workspaceId} AND blob_id = ${blobId};
`;
return true;
}
async removeFile(workspaceId: string, fileId: string) {
// embeddings will be removed by foreign key constraint
await this.db.aiWorkspaceFiles.deleteMany({
where: {
workspaceId,
fileId,
},
});
return true;
}
private allowEmbedding(workspaceId: string) {
return this.models.workspace.allowEmbedding(workspaceId);
}
}
packages/backend/server/src/models/index.ts
-3
View File
@@ -18,7 +18,6 @@ import { CommentAttachmentModel } from './comment-attachment';
import { AppConfigModel } from './config';
import { CopilotActionRunModel } from './copilot-action-run';
import { CopilotWorkspaceByokConfigModel } from './copilot-byok';
import { CopilotContextModel } from './copilot-context';
import { CopilotJobModel } from './copilot-job';
import { CopilotSessionModel } from './copilot-session';
import { CopilotTranscriptTaskModel } from './copilot-transcript-task';
@@ -77,7 +76,6 @@ const MODELS = {
copilotUsage: CopilotUsageModel,
copilotTranscriptTask: CopilotTranscriptTaskModel,
copilotActionRun: CopilotActionRunModel,
copilotContext: CopilotContextModel,
copilotWorkspace: CopilotWorkspaceConfigModel,
copilotWorkspaceByokConfig: CopilotWorkspaceByokConfigModel,
copilotJob: CopilotJobModel,
@@ -153,7 +151,6 @@ export * from './comment';
export * from './comment-attachment';
export * from './common';
export * from './copilot-byok';
export * from './copilot-context';
export * from './copilot-job';
export * from './copilot-session';
export * from './copilot-transcript-task';
packages/backend/server/src/native.ts
+35 -1
View File
@@ -10,8 +10,12 @@ import serverNativeModule, {
type CapabilityAttachmentContract,
type CapabilityModelCapability,
type CommandResponse,
type CompileScopeInput,
type ContentPolicyScanInput,
type ContentPolicyScanResult,
type DocumentEmbeddingProjectionInput,
type EmbeddingHealth,
type EnsureWorkspaceBlobArtifactInput,
type ImageInspection,
type ImageInspectionOptions,
type LicenseError,
@@ -27,12 +31,15 @@ import serverNativeModule, {
type LlmRequestContract,
type LlmRerankRequestContract,
type LlmStructuredRequestContract,
type MatchEmbeddingCandidatesInput,
type ModelConditionsContract,
type PortalResponse,
type PromptMessageContract,
type PromptRenderResult,
type PromptSessionResult,
type PromptStructuredResponseContract,
type PutWorkspaceArtifactInput,
type ReadEmbeddingSourceContentInput,
type RemoteAttachmentFetchRequest,
type RemoteAttachmentFetchResponse,
type RemoteMimeTypeRequest,
@@ -45,6 +52,9 @@ import serverNativeModule, {
type RuntimeBlobMetadataBackfillResult,
type RuntimeDocBlobRefsResult,
type RuntimeDocCompactionResult,
type RuntimeEmbeddingCandidate,
type RuntimeEmbeddingSourceContent,
type RuntimeEmbeddingWorkspaceState,
type RuntimeMagicLinkOtpConsumeResult,
type RuntimeMultipartUploadInit,
type RuntimeMultipartUploadPart,
@@ -53,13 +63,17 @@ import serverNativeModule, {
type RuntimeObjectMetadata,
type RuntimeObjectStoragePutOptions,
type RuntimePresignedObjectRequest,
type RuntimeRetrievalScope,
type RuntimeTurnScopeSnapshot,
type RuntimeVerificationTokenRecord,
type RuntimeWorkspaceArtifact,
type RuntimeWorkspaceInviteLinkRecord,
type RuntimeWorkspaceStatsDailyRecalibrationResult,
type SafeFetchRequest,
type SafeFetchResponse,
type StorageProviderCapabilities,
type StorageRuntimeHealth,
type SyncEmbeddingStateInput,
type Tokenizer,
} from '@affine/server-native';
@@ -76,6 +90,7 @@ export type {
ByokModelDeclarationInput,
ByokModelProbeCheckOutput,
ByokModelProbeOutput,
ByokPolicyOutput,
ByokProbeCheckInput,
ByokProbeResultOutput,
ByokProbeStatusOutput,
@@ -102,8 +117,12 @@ export type {
CapabilityAttachmentContract,
CapabilityModelCapability,
CommandResponse,
CompileScopeInput,
ContentPolicyScanInput,
ContentPolicyScanResult,
DocumentEmbeddingProjectionInput,
EmbeddingHealth,
EnsureWorkspaceBlobArtifactInput,
ImageInspection,
ImageInspectionOptions,
LicenseError,
@@ -113,10 +132,13 @@ export type {
LicenseRecurringRequest,
LicenseResponse,
LicenseSeatsRequest,
MatchEmbeddingCandidatesInput,
ModelConditionsContract,
PortalResponse,
PromptMessageContract,
PromptStructuredResponseContract,
PutWorkspaceArtifactInput,
ReadEmbeddingSourceContentInput,
RemoteAttachmentFetchRequest,
RemoteAttachmentFetchResponse,
RemoteMimeTypeRequest,
@@ -129,6 +151,9 @@ export type {
RuntimeBlobMetadataBackfillResult,
RuntimeDocBlobRefsResult,
RuntimeDocCompactionResult,
RuntimeEmbeddingCandidate,
RuntimeEmbeddingSourceContent,
RuntimeEmbeddingWorkspaceState,
RuntimeMagicLinkOtpConsumeResult,
RuntimeMultipartUploadInit,
RuntimeMultipartUploadPart,
@@ -137,13 +162,17 @@ export type {
RuntimeObjectMetadata,
RuntimeObjectStoragePutOptions,
RuntimePresignedObjectRequest,
RuntimeRetrievalScope,
RuntimeTurnScopeSnapshot,
RuntimeVerificationTokenRecord,
RuntimeWorkspaceArtifact,
RuntimeWorkspaceInviteLinkRecord,
RuntimeWorkspaceStatsDailyRecalibrationResult,
SafeFetchRequest,
SafeFetchResponse,
StorageProviderCapabilities,
StorageRuntimeHealth,
SyncEmbeddingStateInput,
};
export type ActionEventType =
@@ -198,6 +227,8 @@ import type {
} from './plugins/copilot/runtime/contracts/tool-contract';
export const mergeUpdatesInApplyWay = serverNativeModule.mergeUpdatesInApplyWay;
export const authorizeUserdataDocSubject =
serverNativeModule.authorizeUserdataDocSubject;
export const authSessionAccessTokenKeyId =
serverNativeModule.authSessionAccessTokenKeyId;
export const createAuthSessionRefreshToken =
@@ -312,7 +343,10 @@ export const updateLicenseSeats = serverNativeModule.updateLicenseSeats;
export const parseDoc = serverNativeModule.parseDoc;
export const htmlSanitize = serverNativeModule.htmlSanitize;
export const processImage = serverNativeModule.processImage;
export const parseYDocFromBinary = serverNativeModule.parseDocFromBinary;
export const projectDocCanvasFromBinary =
serverNativeModule.projectDocCanvasFromBinary;
export const projectDocSearchFromBinary =
serverNativeModule.projectDocSearchFromBinary;
export const parseYDocToMarkdown = serverNativeModule.parseDocToMarkdown;
export const parsePageDocFromBinary = serverNativeModule.parsePageDoc;
export const parseWorkspaceDocFromBinary = serverNativeModule.parseWorkspaceDoc;
packages/backend/server/src/plugins/copilot/byok/resolver.ts
+79 -59
View File
@@ -12,7 +12,7 @@ import {
} from '@nestjs/graphql';
import { SafeIntResolver } from 'graphql-scalars';
import { Config, Throttle } from '../../../base';
import { Throttle } from '../../../base';
import { CurrentUser } from '../../../core/auth';
import { BackendRuntimeProvider } from '../../../core/backend-runtime';
import { PermissionAccess } from '../../../core/permission';
@@ -22,27 +22,36 @@ import { llmGetByokCatalog } from '../../../native';
import { CopilotEnabled } from '../feature';
import { ByokEntitlementPolicy } from './policy';
import {
BYOK_ALLOWED_PROVIDERS,
ByokAttachmentKind,
ByokAttachmentSource,
ByokCustomEndpointMode,
ByokEndpointKind,
ByokModelFeature,
ByokModelInput,
ByokModelOutput,
ByokOpenAiDialect,
ByokProbeOperation,
ByokProbeStatusKind,
ByokProvider,
ByokProviderSource,
} from './types';
@ObjectType()
class WorkspaceByokCapabilityType {
@Field(() => [String])
input!: string[];
@Field(() => [ByokModelInput])
input!: ByokModelInput[];
@Field(() => [String])
output!: string[];
@Field(() => [ByokModelOutput])
output!: ByokModelOutput[];
@Field(() => [String])
features!: string[];
@Field(() => [ByokModelFeature])
features!: ByokModelFeature[];
@Field(() => [String])
attachmentKinds!: string[];
@Field(() => [ByokAttachmentKind])
attachmentKinds!: ByokAttachmentKind[];
@Field(() => [String])
attachmentSources!: string[];
@Field(() => [ByokAttachmentSource])
attachmentSources!: ByokAttachmentSource[];
}
@ObjectType()
@@ -59,18 +68,18 @@ class WorkspaceByokModelDeclarationType {
@ObjectType()
class WorkspaceByokEndpointType {
@Field(() => String)
kind!: string;
@Field(() => ByokEndpointKind)
kind!: ByokEndpointKind;
@Field(() => String, { nullable: true })
url!: string | null;
@Field(() => ByokOpenAiDialect, { nullable: true })
dialect!: ByokOpenAiDialect | null;
}
@ObjectType()
class WorkspaceByokProfileDefinitionType {
@Field(() => SafeIntResolver)
version!: number;
@Field(() => WorkspaceByokEndpointType)
endpoint!: WorkspaceByokEndpointType;
@@ -80,8 +89,8 @@ class WorkspaceByokProfileDefinitionType {
@ObjectType()
class WorkspaceByokProbeStatusType {
@Field(() => String)
kind!: string;
@Field(() => ByokProbeStatusKind)
kind!: ByokProbeStatusKind;
@Field(() => Date, { nullable: true })
testedAt!: Date | null;
@@ -92,8 +101,8 @@ class WorkspaceByokProbeStatusType {
@ObjectType()
class WorkspaceByokModelProbeCheckType {
@Field(() => String)
operation!: string;
@Field(() => ByokProbeOperation)
operation!: ByokProbeOperation;
@Field(() => WorkspaceByokProbeStatusType)
status!: WorkspaceByokProbeStatusType;
@@ -204,6 +213,21 @@ class WorkspaceByokCatalogType {
providers!: WorkspaceByokCatalogProviderType[];
}
@ObjectType()
class WorkspaceByokPolicyType {
@Field(() => Boolean)
enabled!: boolean;
@Field(() => [ByokProvider])
allowedProviders!: ByokProvider[];
@Field(() => ByokCustomEndpointMode)
customEndpointMode!: ByokCustomEndpointMode;
@Field(() => Boolean)
privateEndpointSupported!: boolean;
}
@ObjectType()
class WorkspaceByokSettingsType {
@Field(() => String)
@@ -221,14 +245,8 @@ class WorkspaceByokSettingsType {
@Field(() => [WorkspaceByokProfileType])
profiles!: WorkspaceByokProfileType[];
@Field(() => [ByokProvider])
allowedProviders!: ByokProvider[];
@Field(() => Boolean)
customEndpointSupported!: boolean;
@Field(() => Boolean)
privateEndpointSupported!: boolean;
@Field(() => WorkspaceByokPolicyType)
policy!: WorkspaceByokPolicyType;
@Field(() => WorkspaceByokCatalogType)
catalog!: WorkspaceByokCatalogType;
@@ -257,20 +275,20 @@ class CreateWorkspaceByokLocalLeaseResultType {
@InputType()
class WorkspaceByokCapabilityInput {
@Field(() => [String])
input!: string[];
@Field(() => [ByokModelInput])
input!: ByokModelInput[];
@Field(() => [String])
output!: string[];
@Field(() => [ByokModelOutput])
output!: ByokModelOutput[];
@Field(() => [String])
features!: string[];
@Field(() => [ByokModelFeature])
features!: ByokModelFeature[];
@Field(() => [String])
attachmentKinds!: string[];
@Field(() => [ByokAttachmentKind])
attachmentKinds!: ByokAttachmentKind[];
@Field(() => [String])
attachmentSources!: string[];
@Field(() => [ByokAttachmentSource])
attachmentSources!: ByokAttachmentSource[];
}
@InputType()
@@ -287,18 +305,18 @@ class WorkspaceByokModelDeclarationInput {
@InputType()
class WorkspaceByokEndpointInput {
@Field(() => String)
kind!: string;
@Field(() => ByokEndpointKind)
kind!: ByokEndpointKind;
@Field(() => String, { nullable: true })
url!: string | null;
@Field(() => ByokOpenAiDialect, { nullable: true })
dialect!: ByokOpenAiDialect | null;
}
@InputType()
class WorkspaceByokProfileDefinitionInput {
@Field(() => SafeIntResolver)
version!: number;
@Field(() => WorkspaceByokEndpointInput)
endpoint!: WorkspaceByokEndpointInput;
@@ -377,8 +395,8 @@ class WorkspaceByokProbeCheckInput {
@Field(() => String)
modelId!: string;
@Field(() => String)
operation!: string;
@Field(() => ByokProbeOperation)
operation!: ByokProbeOperation;
}
@InputType()
@@ -472,8 +490,7 @@ export class WorkspaceByokResolver {
private readonly ac: PermissionAccess,
private readonly entitlement: ByokEntitlementPolicy,
private readonly runtime: BackendRuntimeProvider,
private readonly models: Models,
private readonly config: Config
private readonly models: Models
) {}
@ResolveField(() => WorkspaceByokSettingsType, {
@@ -491,8 +508,8 @@ export class WorkspaceByokResolver {
const profiles = serverEntitled
? await this.runtime.listByokProfiles(workspace.id)
: [];
const customEndpointSupported =
this.config.copilot.byok.allowCustomEndpoint;
const policy = await this.runtime.getByokPolicy();
const allowedProviders = new Set(policy.allowedProviders);
const catalog = llmGetByokCatalog();
return {
workspaceId: workspace.id,
@@ -500,17 +517,19 @@ export class WorkspaceByokResolver {
serverEntitled,
localEntitled,
profiles: profiles.map(profile => projectProfile(profile)),
allowedProviders: [...BYOK_ALLOWED_PROVIDERS],
customEndpointSupported,
privateEndpointSupported:
customEndpointSupported &&
this.config.copilot.byok.allowPrivateEndpoint,
policy: {
...policy,
allowedProviders: policy.allowedProviders as ByokProvider[],
customEndpointMode: policy.customEndpointMode as ByokCustomEndpointMode,
},
catalog: {
...catalog,
providers: catalog.providers.map(provider => ({
...provider,
provider: provider.provider as ByokProvider,
})),
providers: catalog.providers
.filter(provider => allowedProviders.has(provider.provider))
.map(provider => ({
...provider,
provider: provider.provider as ByokProvider,
})),
},
};
}
@@ -711,6 +730,7 @@ function nativeDefinition(input: WorkspaceByokProfileDefinitionInput) {
endpoint: {
...input.endpoint,
url: input.endpoint.url ?? undefined,
dialect: input.endpoint.dialect ?? undefined,
},
};
}
packages/backend/server/src/plugins/copilot/byok/types.ts
+89 -8
View File
@@ -26,6 +26,74 @@ export enum ByokProviderSource {
AffinePlan = 'affine_plan',
}
export enum ByokEndpointKind {
provider_default = 'provider_default',
openai_compatible = 'openai_compatible',
}
export enum ByokOpenAiDialect {
responses = 'responses',
chat_completions = 'chat_completions',
}
export enum ByokModelInput {
text = 'text',
image = 'image',
audio = 'audio',
file = 'file',
}
export enum ByokModelOutput {
text = 'text',
object = 'object',
structured = 'structured',
embedding = 'embedding',
rerank = 'rerank',
image = 'image',
}
export enum ByokModelFeature {
tool_calling = 'tool_calling',
reasoning = 'reasoning',
web_search = 'web_search',
}
export enum ByokAttachmentKind {
image = 'image',
audio = 'audio',
file = 'file',
}
export enum ByokAttachmentSource {
url = 'url',
data = 'data',
bytes = 'bytes',
file_handle = 'file_handle',
}
export enum ByokProbeOperation {
chat = 'chat',
structured = 'structured',
tool_calling = 'tool_calling',
vision = 'vision',
embedding = 'embedding',
rerank = 'rerank',
image = 'image',
transcript = 'transcript',
}
export enum ByokProbeStatusKind {
verified = 'verified',
failed = 'failed',
not_tested = 'not_tested',
}
export enum ByokCustomEndpointMode {
unavailable = 'unavailable',
disabled = 'disabled',
enabled = 'enabled',
}
export type ByokFeatureKind =
| 'chat'
| 'action'
@@ -35,13 +103,6 @@ export type ByokFeatureKind =
| 'transcript'
| 'workspace_indexing';
export const BYOK_ALLOWED_PROVIDERS = [
ByokProvider.openai,
ByokProvider.anthropic,
ByokProvider.gemini,
ByokProvider.fal,
] as const;
export function byokProviderToCopilotType(provider: ByokProvider) {
switch (provider) {
case ByokProvider.openai:
@@ -71,9 +132,29 @@ export function copilotTypeToByokProvider(type: CopilotProviderType) {
}
export function isByokProvider(value: string): value is ByokProvider {
return (BYOK_ALLOWED_PROVIDERS as readonly string[]).includes(value);
switch (value) {
case ByokProvider.openai:
case ByokProvider.anthropic:
case ByokProvider.gemini:
case ByokProvider.fal:
return true;
default:
return false;
}
}
registerEnumType(ByokProvider, { name: 'ByokProvider' });
registerEnumType(ByokKeyStorage, { name: 'ByokKeyStorage' });
registerEnumType(ByokKeyTestStatus, { name: 'ByokKeyTestStatus' });
registerEnumType(ByokEndpointKind, { name: 'ByokEndpointKind' });
registerEnumType(ByokOpenAiDialect, { name: 'ByokOpenAiDialect' });
registerEnumType(ByokModelInput, { name: 'ByokModelInput' });
registerEnumType(ByokModelOutput, { name: 'ByokModelOutput' });
registerEnumType(ByokModelFeature, { name: 'ByokModelFeature' });
registerEnumType(ByokAttachmentKind, { name: 'ByokAttachmentKind' });
registerEnumType(ByokAttachmentSource, { name: 'ByokAttachmentSource' });
registerEnumType(ByokProbeOperation, { name: 'ByokProbeOperation' });
registerEnumType(ByokProbeStatusKind, { name: 'ByokProbeStatusKind' });
registerEnumType(ByokCustomEndpointMode, {
name: 'ByokCustomEndpointMode',
});
packages/backend/server/src/plugins/copilot/config.ts
+45 -98
View File
@@ -1,7 +1,7 @@
import { z } from 'zod';
import serverNativeModule from '@affine/server-native';
import {
defineModuleConfig,
defineNativeModuleConfig,
StorageJSONSchema,
StorageProviderConfig,
} from '../../base';
@@ -9,28 +9,21 @@ import { CopilotProviderType } from './providers/types';
export type ProviderSpecificConfig = Record<string, unknown>;
export const RustRequestMiddlewareValues = [
'normalize_messages',
'clamp_max_tokens',
'tool_schema_rewrite',
'openai_request_compat',
'omit_tool_choice',
] as const;
export type RustRequestMiddleware =
(typeof RustRequestMiddlewareValues)[number];
| 'normalize_messages'
| 'clamp_max_tokens'
| 'tool_schema_rewrite'
| 'openai_request_compat'
| 'omit_tool_choice';
export const RustStreamMiddlewareValues = [
'stream_event_normalize',
'citation_indexing',
] as const;
export type RustStreamMiddleware = (typeof RustStreamMiddlewareValues)[number];
export type RustStreamMiddleware =
| 'stream_event_normalize'
| 'citation_indexing';
export const NodeTextMiddlewareValues = [
'citation_footnote',
'callout',
'thinking_format',
] as const;
export type NodeTextMiddleware = (typeof NodeTextMiddlewareValues)[number];
export type NodeTextMiddleware =
| 'citation_footnote'
| 'callout'
| 'thinking_format';
export type ProviderMiddlewareConfig = {
rust?: { request?: RustRequestMiddleware[]; stream?: RustStreamMiddleware[] };
@@ -51,32 +44,6 @@ export type CopilotProviderProfile = CopilotProviderProfileCommon & {
config: ProviderSpecificConfig;
};
const CopilotProviderProfileBaseShape = z.object({
id: z.string().regex(/^[a-zA-Z0-9-_]+$/),
displayName: z.string().optional(),
priority: z.number().optional(),
enabled: z.boolean().optional(),
models: z.array(z.string().min(1)).min(1),
middleware: z
.object({
rust: z
.object({
request: z.array(z.enum(RustRequestMiddlewareValues)).optional(),
stream: z.array(z.enum(RustStreamMiddlewareValues)).optional(),
})
.optional(),
node: z
.object({ text: z.array(z.enum(NodeTextMiddlewareValues)).optional() })
.optional(),
})
.optional(),
});
const CopilotProviderProfileShape = CopilotProviderProfileBaseShape.extend({
type: z.nativeEnum(CopilotProviderType),
config: z.record(z.string(), z.unknown()),
});
declare global {
interface AppConfigSchema {
copilot: {
@@ -103,57 +70,37 @@ declare global {
}
}
defineModuleConfig('copilot', {
enabled: {
desc: 'Enable AI features. Workspace owners configure provider keys in Workspace Settings → Integrations → AI BYOK.',
default: false,
},
'byok.enabled': {
desc: 'Allow workspace owners and admins to configure AI provider keys through AI BYOK.',
default: true,
shape: z.boolean(),
},
'byok.allowedProviders': {
desc: 'AI providers that workspace owners and admins may add through AI BYOK.',
default: ['openai', 'anthropic', 'gemini', 'fal'],
shape: z.array(z.enum(['openai', 'anthropic', 'gemini', 'fal'])),
},
'byok.allowCustomEndpoint': {
desc: 'Allow AI BYOK keys to use a custom provider endpoint.',
default: false,
shape: z.boolean(),
},
'byok.allowPrivateEndpoint': {
desc: 'Whether workspace BYOK custom endpoints may resolve to private network targets. Enabling this allows workspace owners and admins to send provider probe requests to the private network.',
default: false,
shape: z.boolean(),
},
'providers.profiles': {
desc: 'The profile list for copilot providers.',
default: [],
shape: z.array(CopilotProviderProfileShape),
},
unsplash: {
desc: 'The config for the unsplash key.',
default: {
key: '',
defineNativeModuleConfig(
'copilot',
serverNativeModule.appConfigDescriptors('copilot'),
serverNativeModule.validateAppConfigValue,
{
enabled: {
desc: 'Enable AI features. Workspace owners configure provider keys in Workspace Settings → Integrations → AI BYOK.',
default: false,
},
},
exa: {
desc: 'The config for the exa web search key.',
default: {
key: '',
},
},
storage: {
desc: 'The config for the storage provider.',
default: {
provider: 'fs',
bucket: 'copilot',
config: {
path: '~/.affine/storage',
unsplash: {
desc: 'The config for the unsplash key.',
default: {
key: '',
},
},
schema: StorageJSONSchema,
},
});
exa: {
desc: 'The config for the exa web search key.',
default: {
key: '',
},
},
storage: {
desc: 'The config for the storage provider.',
default: {
provider: 'fs',
bucket: 'copilot',
config: {
path: '~/.affine/storage',
},
},
schema: StorageJSONSchema,
},
}
);
packages/backend/server/src/plugins/copilot/context/index.ts
-3
View File
@@ -1,3 +0,0 @@
export { CopilotEmbeddingRealtimeProvider } from './realtime';
export { CopilotContextResolver, CopilotContextRootResolver } from './resolver';
export { CopilotContextService } from './service';
packages/backend/server/src/plugins/copilot/context/realtime.ts
-131
View File
@@ -1,131 +0,0 @@
import { Injectable, OnModuleInit } from '@nestjs/common';
import { z } from 'zod';
import { Config } from '../../../base/config';
import { OnEvent } from '../../../base/event';
import { PermissionAccess } from '../../../core/permission';
import {
RealtimePublisher,
RealtimeRegistry,
realtimeWorkspaceEmbeddingProgressRoom,
registerRealtimeLiveQuery,
} from '../../../core/realtime';
import { Models } from '../../../models';
import { assertCopilotEnabled } from '../availability';
export function workspaceEmbeddingRoom(workspaceId: string) {
return realtimeWorkspaceEmbeddingProgressRoom(workspaceId);
}
@Injectable()
export class CopilotEmbeddingRealtimeProvider implements OnModuleInit {
constructor(
private readonly ac: PermissionAccess,
private readonly models: Models,
private readonly registry: RealtimeRegistry,
private readonly publisher: RealtimePublisher,
private readonly config: Config
) {}
onModuleInit() {
const input = z.object({ workspaceId: z.string() });
registerRealtimeLiveQuery(this.registry, {
request: {
name: 'workspace.embedding.progress.get',
input,
handle: async (user, payload) => {
await this.assertCopilot(user.id, payload.workspaceId);
const canEmbedding =
await this.models.copilotWorkspace.checkEmbeddingAvailable();
if (!canEmbedding) {
return { total: 0, embedded: 0 };
}
return await this.models.copilotWorkspace.getEmbeddingStatus(
payload.workspaceId
);
},
},
topic: {
name: 'workspace.embedding.progress.changed',
input,
authorize: async (user, payload) => {
await this.assertCopilot(user.id, payload.workspaceId);
},
room: (_user, payload) => workspaceEmbeddingRoom(payload.workspaceId),
},
});
}
@OnEvent('workspace.doc.embed.finished', { suppressError: true })
async onDocEmbedFinished(payload: Events['workspace.doc.embed.finished']) {
await this.publishContext(payload.contextId, 'finished');
}
@OnEvent('workspace.doc.embed.failed', { suppressError: true })
async onDocEmbedFailed(payload: Events['workspace.doc.embed.failed']) {
await this.publishContext(payload.contextId, 'failed');
}
@OnEvent('workspace.file.embed.finished', { suppressError: true })
async onFileEmbedFinished(payload: Events['workspace.file.embed.finished']) {
await this.publishEmbeddingProgress(payload, 'finished');
}
@OnEvent('workspace.file.embed.failed', { suppressError: true })
async onFileEmbedFailed(payload: Events['workspace.file.embed.failed']) {
await this.publishEmbeddingProgress(payload, 'failed');
}
@OnEvent('workspace.blob.embed.finished', { suppressError: true })
async onBlobEmbedFinished(payload: Events['workspace.blob.embed.finished']) {
await this.publishContext(payload.contextId, 'finished');
}
@OnEvent('workspace.blob.embed.failed', { suppressError: true })
async onBlobEmbedFailed(payload: Events['workspace.blob.embed.failed']) {
await this.publishContext(payload.contextId, 'failed');
}
private async publishContext(
contextId: string,
reason: 'finished' | 'failed'
) {
if (!this.publisher) return;
const context = await this.models.copilotContext.getConfig(contextId);
if (!context) return;
this.publishWorkspace(context.workspaceId, reason);
}
private async publishEmbeddingProgress(
payload:
| Events['workspace.file.embed.finished']
| Events['workspace.file.embed.failed'],
reason: 'finished' | 'failed'
) {
if (!this.publisher) return;
if (payload.contextId) {
await this.publishContext(payload.contextId, reason);
return;
}
this.publishWorkspace(payload.workspaceId, reason);
}
private publishWorkspace(workspaceId: string, reason: 'finished' | 'failed') {
this.publisher.publish(
'workspace.embedding.progress.changed',
{ workspaceId },
{ reason },
{ room: workspaceEmbeddingRoom(workspaceId) }
);
}
private async assertCopilot(userId: string, workspaceId: string) {
assertCopilotEnabled(this.config);
await this.ac
.user(userId)
.workspace(workspaceId)
.allowLocal()
.assert('Workspace.Copilot');
}
}
packages/backend/server/src/plugins/copilot/context/resolver.ts
-1063
View File
File diff suppressed because it is too large Load Diff
packages/backend/server/src/plugins/copilot/context/service.ts
-381
View File
@@ -1,381 +0,0 @@
/* oxlint-disable import/no-cycle -- Context embedding reuses the shared capability runtime. */
import { Injectable, OnApplicationBootstrap } from '@nestjs/common';
import {
Cache,
CopilotInvalidContext,
NoCopilotProviderAvailable,
OnEvent,
} from '../../../base';
import {
ContextConfig,
ContextConfigSchema,
ContextDoc,
ContextEmbedStatus,
ContextFile,
Models,
} from '../../../models';
import { CopilotEmbeddingClientService } from '../embedding/client';
import type {
EmbeddingCallOptions,
EmbeddingClient,
EmbeddingRouteContext,
} from '../embedding/types';
import { ContextSession } from './session';
const CONTEXT_SESSION_KEY = 'context-session';
@Injectable()
export class CopilotContextService implements OnApplicationBootstrap {
private supportEmbedding = false;
private client: EmbeddingClient | undefined;
constructor(
private readonly embeddingClients: CopilotEmbeddingClientService,
private readonly cache: Cache,
private readonly models: Models
) {}
@OnEvent('config.init')
async onConfigInit() {
await this.setup();
}
@OnEvent('config.changed')
async onConfigChanged() {
await this.setup();
}
private async setup() {
this.client = await this.embeddingClients.refresh();
}
async onApplicationBootstrap() {
const supportEmbedding =
await this.models.copilotContext.checkEmbeddingAvailable();
if (supportEmbedding) {
this.supportEmbedding = true;
}
}
get canEmbedding() {
return this.supportEmbedding;
}
// public this client to allow overriding in tests
get embeddingClient(): EmbeddingClient | undefined {
return this.client ?? this.embeddingClients.getClient();
}
private embeddingOptions(
workspaceId: string,
signal?: AbortSignal,
routeContext: EmbeddingRouteContext = {}
): EmbeddingCallOptions {
return { workspaceId, signal, ...routeContext, featureKind: 'embedding' };
}
private async saveConfig(
contextId: string,
config: ContextConfig,
refreshCache = false
): Promise<void> {
if (!refreshCache) {
await this.models.copilotContext.update(contextId, { config });
}
await this.cache.set(`${CONTEXT_SESSION_KEY}:${contextId}`, config);
}
private async getCachedSession(
contextId: string
): Promise<ContextSession | undefined> {
const cachedSession = await this.cache.get(
`${CONTEXT_SESSION_KEY}:${contextId}`
);
if (cachedSession) {
const config = ContextConfigSchema.safeParse(cachedSession);
if (config.success) {
return new ContextSession(
this.embeddingClient,
contextId,
config.data,
this.models,
this.saveConfig.bind(this, contextId)
);
}
}
return undefined;
}
// NOTE: we only cache config to avoid frequent database queries
// but we do not need to cache session instances because a distributed
// lock is already apply to mutation operation for the same context in
// the resolver, so there will be no simultaneous writing to the config
private async cacheSession(
contextId: string,
config: ContextConfig
): Promise<ContextSession> {
const dispatcher = this.saveConfig.bind(this, contextId);
await dispatcher(config, true);
return new ContextSession(
this.embeddingClient,
contextId,
config,
this.models,
dispatcher
);
}
async create(sessionId: string): Promise<ContextSession> {
// keep the context unique per session
const existsContext = await this.getBySessionId(sessionId);
if (existsContext) return existsContext;
const context = await this.models.copilotContext.create(sessionId);
const config = ContextConfigSchema.parse(context.config);
return await this.cacheSession(context.id, config);
}
async get(id: string): Promise<ContextSession> {
if (!this.embeddingClient) {
throw new NoCopilotProviderAvailable(
{ modelId: 'embedding' },
'embedding client not configured'
);
}
const context = await this.getCachedSession(id);
if (context) return context;
const config = await this.models.copilotContext.getConfig(id);
if (config) {
return this.cacheSession(id, config);
}
throw new CopilotInvalidContext({ contextId: id });
}
async getOwnedContext(
userId: string,
contextId: string,
options: { workspaceId?: string; sessionId?: string } = {}
): Promise<ContextSession> {
const accessInfo =
await this.models.copilotContext.getAccessInfo(contextId);
if (
!accessInfo ||
accessInfo.session.userId !== userId ||
(options.workspaceId &&
accessInfo.session.workspaceId !== options.workspaceId) ||
(options.sessionId && accessInfo.sessionId !== options.sessionId)
) {
throw new CopilotInvalidContext({ contextId });
}
return await this.get(contextId);
}
async getBySessionId(sessionId: string): Promise<ContextSession | null> {
const existsContext =
await this.models.copilotContext.getBySessionId(sessionId);
if (existsContext) return this.get(existsContext.id);
return null;
}
async matchWorkspaceBlobs(
workspaceId: string,
content: string,
topK: number = 5,
signal?: AbortSignal,
threshold: number = 0.5,
routeContext?: EmbeddingRouteContext
) {
const client = this.embeddingClient;
if (!client) return [];
const options = this.embeddingOptions(workspaceId, signal, routeContext);
const embedding = await client.getEmbedding(content, options);
if (!embedding) return [];
const blobChunks = await this.models.copilotWorkspace.matchBlobEmbedding(
workspaceId,
embedding,
topK * 2,
threshold
);
if (!blobChunks.length) return [];
return await client.reRank(content, blobChunks, topK, options);
}
async matchWorkspaceFiles(
workspaceId: string,
content: string,
topK: number = 5,
signal?: AbortSignal,
threshold: number = 0.5,
routeContext?: EmbeddingRouteContext
) {
const client = this.embeddingClient;
if (!client) return [];
const options = this.embeddingOptions(workspaceId, signal, routeContext);
const embedding = await client.getEmbedding(content, options);
if (!embedding) return [];
const fileChunks = await this.models.copilotWorkspace.matchFileEmbedding(
workspaceId,
embedding,
topK * 2,
threshold
);
if (!fileChunks.length) return [];
return await client.reRank(content, fileChunks, topK, options);
}
async matchWorkspaceDocs(
workspaceId: string,
content: string,
topK: number = 5,
signal?: AbortSignal,
threshold: number = 0.5,
routeContext?: EmbeddingRouteContext
) {
const client = this.embeddingClient;
if (!client) return [];
const options = this.embeddingOptions(workspaceId, signal, routeContext);
const embedding = await client.getEmbedding(content, options);
if (!embedding) return [];
const workspaceChunks =
await this.models.copilotContext.matchWorkspaceEmbedding(
embedding,
workspaceId,
topK * 2,
threshold
);
if (!workspaceChunks.length) return [];
return await client.reRank(content, workspaceChunks, topK, options);
}
async matchWorkspaceAll(
workspaceId: string,
content: string,
topK: number,
signal?: AbortSignal,
threshold: number = 0.8,
docIds?: string[],
scopedThreshold: number = 0.85,
routeContext?: EmbeddingRouteContext
) {
const client = this.embeddingClient;
if (!client) return [];
const options = this.embeddingOptions(workspaceId, signal, routeContext);
const embedding = await client.getEmbedding(content, options);
if (!embedding) return [];
const [fileChunks, blobChunks, workspaceChunks, scopedWorkspaceChunks] =
await Promise.all([
this.models.copilotWorkspace.matchFileEmbedding(
workspaceId,
embedding,
topK * 2,
threshold
),
this.models.copilotWorkspace.matchBlobEmbedding(
workspaceId,
embedding,
topK * 2,
threshold
),
this.models.copilotContext.matchWorkspaceEmbedding(
embedding,
workspaceId,
topK * 2,
threshold
),
docIds
? this.models.copilotContext.matchWorkspaceEmbedding(
embedding,
workspaceId,
topK * 2,
scopedThreshold,
docIds
)
: null,
]);
if (
!fileChunks.length &&
!blobChunks.length &&
!workspaceChunks.length &&
!scopedWorkspaceChunks?.length
) {
return [];
}
return await client.reRank(
content,
[
...fileChunks,
...blobChunks,
...workspaceChunks,
...(scopedWorkspaceChunks || []),
],
topK,
options
);
}
@OnEvent('workspace.doc.embed.failed')
async onDocEmbedFailed({
contextId,
docId,
}: Events['workspace.doc.embed.failed']) {
const context = await this.get(contextId);
await context.saveDocRecord(docId, doc => ({
...(doc as ContextDoc),
status: ContextEmbedStatus.failed,
}));
}
@OnEvent('workspace.doc.embed.finished')
async onDocEmbedFinished({
contextId,
docId,
}: Events['workspace.doc.embed.finished']) {
const context = await this.get(contextId);
await context.saveDocRecord(docId, doc => ({
...(doc as ContextDoc),
status: ContextEmbedStatus.finished,
}));
}
@OnEvent('workspace.file.embed.finished')
async onFileEmbedFinish({
contextId,
fileId,
chunkSize,
}: Events['workspace.file.embed.finished']) {
if (!contextId) return;
const context = await this.get(contextId);
await context.saveFileRecord(fileId, file => ({
...(file as ContextFile),
chunkSize,
status: ContextEmbedStatus.finished,
}));
}
@OnEvent('workspace.file.embed.failed')
async onFileEmbedFailed({
contextId,
fileId,
error,
}: Events['workspace.file.embed.failed']) {
if (!contextId) return;
const context = await this.get(contextId);
await context.saveFileRecord(fileId, file => ({
...(file as ContextFile),
error,
status: ContextEmbedStatus.failed,
}));
}
}
packages/backend/server/src/plugins/copilot/context/session.ts
-426
View File
@@ -1,426 +0,0 @@
import { nanoid } from 'nanoid';
import {
ContextBlob,
ContextCategories,
ContextCategory,
ContextConfig,
ContextDoc,
ContextEmbedStatus,
ContextFile,
FileChunkSimilarity,
Models,
} from '../../../models';
import type {
EmbeddingCallOptions,
EmbeddingClient,
EmbeddingRouteContext,
} from '../embedding/types';
export class ContextSession implements AsyncDisposable {
constructor(
private readonly client: EmbeddingClient | undefined,
private readonly contextId: string,
private readonly config: ContextConfig,
private readonly models: Models,
private readonly dispatcher?: (config: ContextConfig) => Promise<void>
) {}
get id() {
return this.contextId;
}
get workspaceId() {
return this.config.workspaceId;
}
get categories(): ContextCategory[] {
return this.config.categories.map(c => ({
...c,
docs: c.docs.map(d => ({ ...d })),
}));
}
get tags() {
const categories = this.config.categories;
return categories.filter(c => c.type === ContextCategories.Tag);
}
get collections() {
const categories = this.config.categories;
return categories.filter(c => c.type === ContextCategories.Collection);
}
get blobs(): ContextBlob[] {
return this.config.blobs.map(d => ({ ...d }));
}
get docs(): ContextDoc[] {
return this.config.docs.map(d => ({ ...d }));
}
get files(): Required<ContextFile>[] {
return this.config.files.map(f => this.fulfillFile(f));
}
get docIds() {
return Array.from(
new Set(
[this.config.docs, this.config.categories.flatMap(c => c.docs)]
.flat()
.map(d => d.id)
)
);
}
private embeddingOptions(
signal?: AbortSignal,
routeContext: EmbeddingRouteContext = {}
): EmbeddingCallOptions {
return {
workspaceId: this.workspaceId,
signal,
...routeContext,
featureKind: 'embedding',
};
}
async addCategoryRecord(type: ContextCategories, id: string, docs: string[]) {
const category = this.config.categories.find(
c => c.type === type && c.id === id
);
if (category) {
const missingDocs = docs.filter(
docId => !category.docs.some(d => d.id === docId)
);
if (missingDocs.length) {
category.docs.push(
...missingDocs.map(id => ({
id,
createdAt: Date.now(),
status: ContextEmbedStatus.processing,
}))
);
await this.save();
}
return category;
}
const createdAt = Date.now();
const record = {
id,
type,
docs: docs.map(id => ({
id,
createdAt,
status: ContextEmbedStatus.processing,
})),
createdAt,
};
this.config.categories.push(record);
await this.save();
return record;
}
async removeCategoryRecord(type: ContextCategories, id: string) {
const index = this.config.categories.findIndex(
c => c.type === type && c.id === id
);
if (index >= 0) {
this.config.categories.splice(index, 1);
await this.save();
}
return true;
}
async addBlobRecord(blobId: string): Promise<ContextBlob | null> {
const existsBlob = this.config.blobs.find(b => b.id === blobId);
if (existsBlob) {
return existsBlob;
}
const blob = await this.models.blob.get(this.config.workspaceId, blobId);
if (!blob) return null;
const record: ContextBlob = {
id: blobId,
createdAt: Date.now(),
status: ContextEmbedStatus.processing,
};
this.config.blobs.push(record);
await this.save();
return record;
}
async getBlobMetadata() {
const blobIds = this.blobs.map(b => b.id);
const blobs = await this.models.blob.list(this.config.workspaceId, {
where: { key: { in: blobIds } },
select: { key: true, mime: true },
});
const blobChunkSizes = await this.models.copilotWorkspace.getBlobChunkSizes(
this.config.workspaceId,
blobIds
);
return blobs
.filter(b => !!blobChunkSizes.get(b.key))
.map(b => ({
id: b.key,
mimeType: b.mime,
chunkSize: blobChunkSizes.get(b.key),
}));
}
async getBlobContent(
blobId: string,
chunk?: number
): Promise<string | undefined> {
return this.models.copilotWorkspace.getBlobContent(
this.config.workspaceId,
blobId,
chunk
);
}
async removeBlobRecord(blobId: string): Promise<boolean> {
const index = this.config.blobs.findIndex(b => b.id === blobId);
if (index >= 0) {
this.config.blobs.splice(index, 1);
await this.save();
}
return true;
}
async addDocRecord(docId: string): Promise<ContextDoc> {
const doc = this.config.docs.find(f => f.id === docId);
if (doc) {
return doc;
}
const record = { id: docId, createdAt: Date.now() };
this.config.docs.push(record);
await this.save();
return record;
}
async removeDocRecord(docId: string): Promise<boolean> {
const index = this.config.docs.findIndex(f => f.id === docId);
if (index >= 0) {
this.config.docs.splice(index, 1);
await this.save();
}
return true;
}
private fulfillFile(file: ContextFile): Required<ContextFile> {
return {
...file,
mimeType: file.mimeType || 'application/octet-stream',
};
}
async addFile(
blobId: string,
name: string,
mimeType: string
): Promise<Required<ContextFile>> {
let fileId = nanoid();
const existsBlob = this.config.files.find(f => f.blobId === blobId);
if (existsBlob) {
// use exists file id if the blob exists
// we assume that the file content pointed to by the same blobId is consistent.
if (existsBlob.status === ContextEmbedStatus.finished) {
return this.fulfillFile(existsBlob);
}
fileId = existsBlob.id;
} else {
await this.saveFileRecord(fileId, file => ({
...file,
blobId,
chunkSize: 0,
name,
mimeType,
error: null,
createdAt: Date.now(),
}));
}
return this.fulfillFile(this.getFile(fileId) as ContextFile);
}
getFile(fileId: string): ContextFile | undefined {
return this.config.files.find(f => f.id === fileId);
}
async getFileContent(
fileId: string,
chunk?: number
): Promise<string | undefined> {
const file = this.getFile(fileId);
if (!file) return undefined;
return this.models.copilotContext.getFileContent(
this.contextId,
fileId,
chunk
);
}
async removeFile(fileId: string): Promise<boolean> {
await this.models.copilotContext.deleteFileEmbedding(
this.contextId,
fileId
);
this.config.files = this.config.files.filter(f => f.id !== fileId);
await this.save();
return true;
}
/**
* Match the input text with the file chunks
* @param content input text to match
* @param topK number of similar chunks to return, default 5
* @param signal abort signal
* @param threshold relevance threshold for the similarity score, higher threshold means more similar chunks, default 0.7, good enough based on prior experiments
* @returns list of similar chunks
*/
async matchFiles(
content: string,
topK: number = 5,
signal?: AbortSignal,
scopedThreshold: number = 0.85,
threshold: number = 0.5,
routeContext?: EmbeddingRouteContext
): Promise<FileChunkSimilarity[]> {
if (!this.client) return [];
const options = this.embeddingOptions(signal, routeContext);
const embedding = await this.client.getEmbedding(content, options);
if (!embedding) return [];
const [context, workspace] = await Promise.all([
this.models.copilotContext.matchFileEmbedding(
embedding,
this.id,
topK * 2,
scopedThreshold
),
this.models.copilotWorkspace.matchFileEmbedding(
this.workspaceId,
embedding,
topK * 2,
threshold
),
]);
const files = new Map(this.files.map(f => [f.id, f]));
return this.client.reRank(
content,
[
...context
.filter(f => files.has(f.fileId))
.map(c => {
const { blobId, name, mimeType } = files.get(
c.fileId
) as Required<ContextFile>;
return { ...c, blobId, name, mimeType };
}),
...workspace,
],
topK,
options
);
}
/**
* Match the input text with the workspace chunks
* @param content input text to match
* @param topK number of similar chunks to return, default 5
* @param signal abort signal
* @param threshold relevance threshold for the similarity score, higher threshold means more similar chunks, default 0.7, good enough based on prior experiments
* @returns list of similar chunks
*/
async matchWorkspaceDocs(
content: string,
topK: number = 5,
signal?: AbortSignal,
scopedThreshold: number = 0.85,
threshold: number = 0.5,
routeContext?: EmbeddingRouteContext
) {
if (!this.client) return [];
const options = this.embeddingOptions(signal, routeContext);
const embedding = await this.client.getEmbedding(content, options);
if (!embedding) return [];
const docIds = this.docIds;
const [inContext, workspace] = await Promise.all([
this.models.copilotContext.matchWorkspaceEmbedding(
embedding,
this.workspaceId,
topK * 2,
scopedThreshold,
docIds
),
this.models.copilotContext.matchWorkspaceEmbedding(
embedding,
this.workspaceId,
topK * 2,
threshold
),
]);
const result = await this.client.reRank(
content,
[...inContext, ...workspace],
topK,
options
);
// sort result, doc recorded in context first
const docIdSet = new Set(docIds);
return result.toSorted(
(a, b) =>
(docIdSet.has(a.docId) ? -1 : 1) - (docIdSet.has(b.docId) ? -1 : 1) ||
(a.distance || Infinity) - (b.distance || Infinity)
);
}
async saveDocRecord(
docId: string,
cb: (
record: Pick<ContextDoc, 'id' | 'status'> &
Partial<Omit<ContextDoc, 'id' | 'status'>>
) => ContextDoc
) {
const docs = [this.config.docs, ...this.config.categories.map(c => c.docs)]
.flat()
.filter(d => d.id === docId);
for (const doc of docs) {
Object.assign(doc, cb({ ...doc }));
}
await this.save();
}
async saveFileRecord(
fileId: string,
cb: (
record: Pick<ContextFile, 'id' | 'status'> &
Partial<Omit<ContextFile, 'id' | 'status'>>
) => ContextFile
) {
const files = this.config.files;
const file = files.find(f => f.id === fileId);
if (file) {
Object.assign(file, cb({ ...file }));
} else {
const file = { id: fileId, status: ContextEmbedStatus.processing };
files.push(cb(file));
}
await this.save();
}
async save() {
await this.dispatcher?.(this.config);
}
async [Symbol.asyncDispose]() {
await this.save();
}
}
packages/backend/server/src/plugins/copilot/conversation/inbox.ts
+28 -1
View File
@@ -8,6 +8,7 @@ import {
sniffMime,
} from '../../../base';
import { PermissionAccess } from '../../../core/permission';
import { Models } from '../../../models';
import { processImage } from '../../../native';
import { CompatSubmissionStore } from '../compat/submission-store';
import type { PromptMessage } from '../providers/types';
@@ -30,6 +31,7 @@ export class ConversationInboxService {
constructor(
private readonly chatSession: ChatSessionService,
private readonly ac: PermissionAccess,
private readonly models: Models,
private readonly storage: CopilotStorage,
private readonly submissions: CompatSubmissionStore
) {}
@@ -48,6 +50,26 @@ export class ConversationInboxService {
options.blob ? [options.blob] : options.blobs || []
);
const focusSelectors = options.params?.focusSelectors;
const hasWorkspaceContext =
attachments.length > 0 ||
blobs.length > 0 ||
(Array.isArray(options.params?.scopeSelectors) &&
options.params.scopeSelectors.length > 0) ||
(Array.isArray(options.params?.preferredSourceIds) &&
options.params.preferredSourceIds.length > 0) ||
(focusSelectors === undefined
? session.config.focus.selectors.length > 0
: Array.isArray(focusSelectors) && focusSelectors.length > 0);
if (
hasWorkspaceContext &&
!(await this.models.workspace.get(session.config.workspaceId))
) {
throw new BadRequestException(
"Local workspaces don't support attachments or references."
);
}
if (blobs.length) {
await this.ac
.user(userId)
@@ -86,7 +108,12 @@ export class ConversationInboxService {
filename,
attachmentBuffer
);
attachments.push({ attachment, mimeType: attachmentMimeType });
attachments.push({
kind: 'url',
url: attachment,
mimeType: attachmentMimeType,
fileName: blob.filename,
});
}
return await this.submissions.create({
packages/backend/server/src/plugins/copilot/conversation/store.ts
+24
View File
@@ -12,6 +12,10 @@ import {
type Turn,
turnFromChatMessage,
} from '../core';
import {
type SessionFocus,
SessionFocusSchema,
} from '../runtime/contracts/shared';
import { type ChatMessage, ChatMessageSchema } from '../types';
type SessionRecord = NonNullable<
@@ -68,6 +72,11 @@ export class ConversationStore {
return parsed.data;
}
private toFocus(focus: unknown): SessionFocus {
const parsed = SessionFocusSchema.safeParse(focus);
return parsed.success ? parsed.data : { selectors: [] };
}
async create(
seed: ConversationSeed,
reuseLatestChat = false
@@ -83,6 +92,7 @@ export class ConversationStore {
conversation: Conversation;
turns: Turn[];
promptName: string;
focus: SessionFocus;
}
| undefined
> {
@@ -95,6 +105,7 @@ export class ConversationStore {
conversation: this.toConversation(session),
turns: this.toTurns(session),
promptName: session.promptName,
focus: this.toFocus(session.focus),
};
}
@@ -102,6 +113,7 @@ export class ConversationStore {
| {
conversation: Conversation;
promptName: string;
focus: SessionFocus;
}
| undefined
> {
@@ -121,6 +133,7 @@ export class ConversationStore {
updatedAt: session.updatedAt,
},
promptName: session.promptName,
focus: this.toFocus(session.focus),
};
}
@@ -142,6 +155,7 @@ export class ConversationStore {
turnFromChatMessage(message, session.id)
),
promptName: session.promptName,
focus: this.toFocus(session.focus),
}));
}
@@ -163,6 +177,7 @@ export class ConversationStore {
updatedAt: session.updatedAt,
} satisfies Conversation,
promptName: session.promptName,
focus: this.toFocus(session.focus),
}));
}
@@ -185,10 +200,19 @@ export class ConversationStore {
userId: string;
turn: Turn;
compatSubmissionId?: string;
focus?: SessionFocus;
artifacts?: Array<{
artifactId: string;
role: string;
displayName?: string;
metadata?: Record<string, unknown>;
}>;
}) {
const message = await this.models.copilotSession.appendMessage({
sessionId: input.sessionId,
userId: input.userId,
focus: input.focus,
artifacts: input.artifacts,
message: (() => {
const { id: _id, ...message } = chatMessageFromTurn(input.turn);
return { ...message, compatSubmissionId: input.compatSubmissionId };
packages/backend/server/src/plugins/copilot/core/adapters.ts
+16 -6
View File
@@ -1,4 +1,5 @@
import type { PromptMessage, StreamObject } from '../providers/types';
import { promptAttachmentMimeType } from '../providers/utils';
import {
streamObjectToToolEvent,
toolEventToStreamObject,
@@ -82,6 +83,7 @@ export const turnFromChatMessage = (
renderTrace: trace.renderTrace,
toolEvents: trace.toolEvents,
metadata: message.params ?? {},
scopeSnapshot: message.scopeSnapshot,
createdAt: message.createdAt,
});
};
@@ -95,14 +97,22 @@ export const chatMessageFromTurn = (turn: Turn): ChatMessage => {
content: turn.content,
attachments: turn.attachments.length ? turn.attachments : undefined,
params: turn.metadata,
scopeSnapshot: turn.scopeSnapshot,
streamObjects: renderTrace.length ? renderTrace : undefined,
createdAt: turn.createdAt,
};
};
export const promptMessageFromTurn = (turn: Turn): PromptMessage => ({
role: turn.role,
content: turn.content,
attachments: turn.attachments.length ? turn.attachments : undefined,
params: Object.keys(turn.metadata).length ? turn.metadata : undefined,
});
export const promptMessageFromTurn = (turn: Turn): PromptMessage => {
const attachments = turn.attachments.filter(attachment => {
const mimeType = promptAttachmentMimeType(attachment);
return !mimeType || mimeType.startsWith('image/');
});
return {
role: turn.role,
content: turn.content,
attachments: attachments.length ? attachments : undefined,
params: Object.keys(turn.metadata).length ? turn.metadata : undefined,
};
};
packages/backend/server/src/plugins/copilot/core/types.ts
+2
View File
@@ -6,6 +6,7 @@ import {
type ToolEvent,
ToolEventSchema,
} from '../runtime/contracts/runtime-event-contract';
import { TurnScopeSnapshotSchema } from '../runtime/contracts/shared';
const CanonicalDateSchema = z.coerce.date();
@@ -35,6 +36,7 @@ export const TurnSchema = z
renderTrace: z.array(StreamObjectSchema).default([]),
toolEvents: z.array(ToolEventSchema).default([]),
metadata: z.record(z.string(), z.any()).default({}),
scopeSnapshot: TurnScopeSnapshotSchema.nullable().optional(),
createdAt: CanonicalDateSchema,
})
.strict();
packages/backend/server/src/plugins/copilot/cron.ts
+1 -37
View File
@@ -1,19 +1,15 @@
import { Injectable, Logger } from '@nestjs/common';
import { Cron, CronExpression } from '@nestjs/schedule';
import { JOB_SIGNAL, JobQueue, OneDay, OnJob } from '../../base';
import { 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 {
'copilot.session.cleanupEmptySessions': {};
'copilot.session.generateMissingTitles': {};
'copilot.workspace.cleanupTrashedDocEmbeddings': {
nextSid?: number;
};
}
}
@@ -39,12 +35,6 @@ export class CopilotCronJobs {
{},
{ jobId: 'daily-copilot-generate-missing-titles' }
);
await this.jobs.add(
'copilot.workspace.cleanupTrashedDocEmbeddings',
{},
{ jobId: 'daily-copilot-cleanup-trashed-doc-embeddings' }
);
}
async triggerGenerateMissingTitles() {
@@ -82,30 +72,4 @@ export class CopilotCronJobs {
`Scheduled title generation for ${sessions.length} sessions`
);
}
@OnJob('copilot.workspace.cleanupTrashedDocEmbeddings')
async cleanupTrashedDocEmbeddings(
params: Jobs['copilot.workspace.cleanupTrashedDocEmbeddings']
) {
const nextSid = params.nextSid ?? 0;
// only consider workspaces that cleared their embeddings more than 24 hours ago
const oneDayAgo = new Date(Date.now() - OneDay);
const workspaces = await this.models.workspace.list(
{ sid: { gt: nextSid }, lastCheckEmbeddings: { lt: oneDayAgo } },
{ id: true, sid: true },
CLEANUP_EMBEDDING_JOB_BATCH_SIZE
);
if (!workspaces.length) {
return JOB_SIGNAL.Done;
}
for (const { id: workspaceId } of workspaces) {
await this.jobs.add(
'copilot.embedding.cleanupTrashedDocEmbeddings',
{ workspaceId },
{ jobId: `cleanup-trashed-doc-embeddings-${workspaceId}` }
);
}
params.nextSid = workspaces[workspaces.length - 1].sid;
return JOB_SIGNAL.Repeat;
}
}
packages/backend/server/src/plugins/copilot/delegated/realtime.ts
+144
View File
@@ -0,0 +1,144 @@
import { Injectable, OnModuleInit } from '@nestjs/common';
import { z } from 'zod';
import { EventBus } from '../../../base';
import { RealtimeRegistry, realtimeUserRoom } from '../../../core/realtime';
import { ChatSessionService } from '../session';
import { DelegatedEditorService } from './service';
const identity = {
requestId: z.string().uuid(),
runId: z.string().uuid(),
toolCallId: z.string().min(1).max(256),
sessionId: z.string().min(1),
workspaceId: z.string().min(1),
docId: z.string().min(1),
clientId: z.string().min(1).max(128),
editorStateId: z.string().min(1).max(128),
};
const responseSchema = z
.object({
...identity,
result: z.unknown().optional(),
error: z
.object({
code: z.string().min(1).max(64),
message: z.string().max(500),
retryable: z.boolean(),
})
.strict()
.optional(),
})
.strict()
.refine(
response =>
(response.result !== undefined) !== (response.error !== undefined),
{ message: 'Exactly one of result or error is required.' }
)
.refine(
response =>
Buffer.byteLength(JSON.stringify(response.result ?? null)) <= 512 * 1024,
{ message: 'Delegated tool result is too large.' }
);
@Injectable()
export class DelegatedEditorRealtimeProvider implements OnModuleInit {
constructor(
private readonly registry: RealtimeRegistry,
private readonly event: EventBus,
private readonly sessions: ChatSessionService,
private readonly delegated: DelegatedEditorService
) {}
onModuleInit() {
const leaseInput = z
.object({
clientId: z.string().min(1).max(128),
sessionId: z.string().min(1),
workspaceId: z.string().min(1),
docId: z.string().min(1),
editorStateId: z.string().min(1).max(128),
mode: z.enum(['page', 'edgeless']),
readonly: z.boolean(),
focused: z.boolean(),
capabilities: z
.array(
z.enum([
'frontend_get_editor_state',
'frontend_read_selection',
'frontend_read_nodes',
'frontend_snapshot_document',
])
)
.max(4),
})
.strict();
this.registry.registerRequest({
name: 'copilot.delegated.editor.upsert',
input: leaseInput,
handle: async (user, input, context) => {
const session = await this.sessions.get(input.sessionId);
if (
!user ||
!context?.connectionId ||
!session ||
session.config.userId !== user.id ||
session.config.workspaceId !== input.workspaceId ||
session.config.docId !== input.docId
) {
throw new Error('INVALID_DELEGATED_EDITOR_SESSION');
}
const lease = this.delegated.upsert(
user.id,
context.connectionId,
input
);
this.event.broadcast('copilot.delegated.editor.upserted', lease);
return { ok: true, expiresAt: lease.expiresAt };
},
});
this.registry.registerRequest({
name: 'copilot.delegated.editor.release',
input: z
.object({
clientId: z.string().min(1).max(128),
editorStateId: z.string().min(1).max(128),
})
.strict(),
handle: async (user, input) => {
if (user) {
this.delegated.release(user.id, input.clientId, input.editorStateId);
this.event.broadcast('copilot.delegated.editor.released', {
userId: user.id,
...input,
});
}
return { ok: true };
},
});
this.registry.registerRequest({
name: 'copilot.delegated.tool.respond',
input: responseSchema,
handle: async (user, response) => {
if (!user) return { accepted: false };
const accepted = this.delegated.receive(user.id, response);
this.event.broadcast('copilot.delegated.tool.responded', {
userId: user.id,
response,
});
return { accepted };
},
});
this.registry.registerTopic({
name: 'copilot.delegated.tool.requested',
input: z.object({ clientId: z.string().min(1).max(128) }).strict(),
authorize: async user => {
if (!user) throw new Error('AUTHENTICATION_REQUIRED');
},
room: (user, input) => {
if (!user) throw new Error('AUTHENTICATION_REQUIRED');
return realtimeUserRoom(user.id, `copilot:${input.clientId}`);
},
});
}
}
packages/backend/server/src/plugins/copilot/delegated/service.ts
+303
View File
@@ -0,0 +1,303 @@
import { randomUUID } from 'node:crypto';
import type {
DelegatedEditorLeaseInput,
DelegatedToolIdentity,
DelegatedToolName,
DelegatedToolResponse,
} from '@affine/realtime';
import { Injectable } from '@nestjs/common';
import { OnEvent } from '../../../base';
import { RealtimePublisher, realtimeUserRoom } from '../../../core/realtime';
import type { CopilotChatOptions } from '../providers/types';
type EditorLease = DelegatedEditorLeaseInput & {
userId: string;
connectionId: string;
expiresAt: number;
};
type PendingRequest = {
identity: DelegatedToolIdentity;
userId: string;
connectionId: string;
resolve: (response: DelegatedToolResponse) => void;
};
declare global {
interface Events {
'copilot.delegated.editor.upserted': EditorLease;
'copilot.delegated.editor.released': {
userId: string;
clientId: string;
editorStateId: string;
};
'copilot.delegated.tool.responded': {
userId: string;
response: DelegatedToolResponse;
};
}
}
const LEASE_TTL_MS = 30_000;
const TOOL_TIMEOUT_MS = 15_000;
@Injectable()
export class DelegatedEditorService {
private readonly leases = new Map<string, EditorLease>();
private readonly pending = new Map<string, PendingRequest>();
constructor(private readonly publisher: RealtimePublisher) {}
leaseKey(userId: string, clientId: string) {
return `${userId}:${clientId}`;
}
upsert(
userId: string,
connectionId: string,
input: DelegatedEditorLeaseInput
) {
const lease = {
...input,
userId,
connectionId,
expiresAt: Date.now() + LEASE_TTL_MS,
};
this.leases.set(this.leaseKey(userId, input.clientId), lease);
return lease;
}
release(userId: string, clientId: string, editorStateId: string) {
const key = this.leaseKey(userId, clientId);
const lease = this.leases.get(key);
if (lease?.editorStateId === editorStateId) {
this.leases.delete(key);
}
}
getLease(options: CopilotChatOptions, tool?: DelegatedToolName) {
if (!options?.user || !options.session || !options.workspace) return null;
const now = Date.now();
let selected: EditorLease | null = null;
for (const [key, lease] of this.leases) {
if (lease.expiresAt <= now) {
this.leases.delete(key);
continue;
}
if (
lease.userId === options.user &&
lease.sessionId === options.session &&
lease.workspaceId === options.workspace &&
lease.focused &&
(!tool || lease.capabilities.includes(tool)) &&
(!selected || lease.expiresAt > selected.expiresAt)
) {
selected = lease;
}
}
return selected;
}
async execute(
options: CopilotChatOptions,
tool: DelegatedToolName,
args: Record<string, unknown>,
signal?: AbortSignal,
execution?: { runId?: string; toolCallId?: string }
) {
const lease = this.getLease(options, tool);
if (!lease) {
return {
error: {
code: 'FRONTEND_UNAVAILABLE',
message: 'No focused editor is available for this session.',
retryable: true,
},
};
}
const identity = {
requestId: randomUUID(),
runId: execution?.runId ?? randomUUID(),
toolCallId: execution?.toolCallId ?? randomUUID(),
sessionId: lease.sessionId,
workspaceId: lease.workspaceId,
docId: lease.docId,
clientId: lease.clientId,
editorStateId: lease.editorStateId,
};
const deadlineAt = Date.now() + TOOL_TIMEOUT_MS;
const response = new Promise<DelegatedToolResponse>(resolve => {
this.pending.set(identity.requestId, {
identity,
userId: lease.userId,
connectionId: lease.connectionId,
resolve,
});
});
this.publisher.publish(
'copilot.delegated.tool.requested',
{ clientId: lease.clientId },
{ type: 'request', ...identity, tool, args, deadlineAt },
{ room: realtimeUserRoom(lease.userId, `copilot:${lease.clientId}`) }
);
let reason: 'aborted' | 'timeout' | undefined;
let timeout: ReturnType<typeof setTimeout> | undefined;
let abort: (() => void) | undefined;
const interrupted = new Promise<DelegatedToolResponse>(resolve => {
timeout = setTimeout(() => {
reason = 'timeout';
resolve({
...identity,
error: {
code: 'FRONTEND_TIMEOUT',
message: 'The focused editor did not respond before the deadline.',
retryable: true,
},
});
}, TOOL_TIMEOUT_MS);
timeout.unref?.();
abort = () => {
reason = 'aborted';
resolve({
...identity,
error: {
code: 'ABORTED',
message: 'The delegated read was cancelled.',
retryable: false,
},
});
};
if (signal?.aborted) {
abort();
} else {
signal?.addEventListener('abort', abort, { once: true });
}
});
const result = await Promise.race([response, interrupted]);
this.pending.delete(identity.requestId);
if (timeout) clearTimeout(timeout);
if (abort) signal?.removeEventListener('abort', abort);
if (reason) {
this.publisher.publish(
'copilot.delegated.tool.requested',
{ clientId: lease.clientId },
{ type: 'cancel', ...identity, reason },
{ room: realtimeUserRoom(lease.userId, `copilot:${lease.clientId}`) }
);
}
if (result.error) return { error: result.error };
if (
tool === 'frontend_get_editor_state' ||
!result.result ||
typeof result.result !== 'object' ||
Array.isArray(result.result)
) {
return result.result;
}
return {
...result.result,
source: {
type: 'document',
workspace_id: lease.workspaceId,
doc_id: lease.docId,
revision: lease.editorStateId,
},
};
}
receive(userId: string, response: DelegatedToolResponse) {
const request = this.pending.get(response.requestId);
if (
!request ||
request.userId !== userId ||
!this.sameIdentity(request.identity, response) ||
!this.validResult(request.identity, response)
) {
return false;
}
this.pending.delete(response.requestId);
request.resolve(response);
return true;
}
@OnEvent('copilot.delegated.editor.upserted', { suppressError: true })
onRemoteUpsert(lease: Events['copilot.delegated.editor.upserted']) {
this.leases.set(this.leaseKey(lease.userId, lease.clientId), lease);
}
@OnEvent('copilot.delegated.editor.released', { suppressError: true })
onRemoteRelease(event: Events['copilot.delegated.editor.released']) {
this.release(event.userId, event.clientId, event.editorStateId);
}
@OnEvent('copilot.delegated.tool.responded', { suppressError: true })
onRemoteResponse(event: Events['copilot.delegated.tool.responded']) {
this.receive(event.userId, event.response);
}
@OnEvent('realtime.connection.disconnected', { suppressError: true })
onDisconnect({ connectionId }: Events['realtime.connection.disconnected']) {
for (const [key, lease] of this.leases) {
if (lease.connectionId === connectionId) {
this.leases.delete(key);
}
}
for (const [requestId, request] of this.pending) {
if (request.connectionId !== connectionId) continue;
this.pending.delete(requestId);
request.resolve({
...request.identity,
error: {
code: 'FRONTEND_DISCONNECTED',
message: 'The focused editor disconnected during the read.',
retryable: true,
},
});
this.publisher.publish(
'copilot.delegated.tool.requested',
{ clientId: request.identity.clientId },
{ type: 'cancel', ...request.identity, reason: 'disconnect' },
{
room: realtimeUserRoom(
request.userId,
`copilot:${request.identity.clientId}`
),
}
);
}
}
private sameIdentity(
expected: DelegatedToolIdentity,
actual: DelegatedToolIdentity
) {
return (
expected.requestId === actual.requestId &&
expected.runId === actual.runId &&
expected.toolCallId === actual.toolCallId &&
expected.sessionId === actual.sessionId &&
expected.workspaceId === actual.workspaceId &&
expected.docId === actual.docId &&
expected.clientId === actual.clientId &&
expected.editorStateId === actual.editorStateId
);
}
private validResult(
identity: DelegatedToolIdentity,
response: DelegatedToolResponse
) {
if (response.error) return true;
return Boolean(
response.result &&
typeof response.result === 'object' &&
'editor_state_id' in response.result &&
response.result.editor_state_id === identity.editorStateId
);
}
}
packages/backend/server/src/plugins/copilot/embedding/client.ts
-237
View File
@@ -1,237 +0,0 @@
/* oxlint-disable import/no-cycle -- Embedding delegates to the shared capability runtime. */
import { createHash } from 'node:crypto';
import { forwardRef, Inject, Injectable, Logger } from '@nestjs/common';
import { CopilotFailedToGenerateEmbedding } from '../../../base/error/errors.gen';
import {
ChunkSimilarity,
Embedding,
EMBEDDING_DIMENSIONS,
} from '../../../models';
import { type CopilotRerankRequest } from '../providers/types';
import { CapabilityRuntime } from '../runtime/capability-runtime';
import {
type EmbeddingCallOptionsInput,
EmbeddingClient,
normalizeEmbeddingCallOptions,
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 runtime: EmbeddingRuntime) {
super();
}
override async configured(): Promise<boolean> {
const result = await this.runtime.embeddingConfigured('route-selected');
if (!result) {
this.logger.warn(
'Copilot embedding client is not configured properly, please check your configuration.'
);
}
return result;
}
async getEmbeddings(
input: string[],
options?: EmbeddingCallOptionsInput
): Promise<Embedding[]> {
const normalizedOptions = normalizeEmbeddingCallOptions(options);
const modelId = 'route-selected';
const embeddings = await this.runtime.embed(modelId, input, {
dimensions: EMBEDDING_DIMENSIONS,
signal: normalizedOptions.signal,
user: normalizedOptions.userId,
workspace: normalizedOptions.workspaceId,
byokLeaseId: normalizedOptions.byokLeaseId,
featureKind: normalizedOptions.featureKind ?? 'embedding',
});
if (embeddings.length !== input.length) {
throw new CopilotFailedToGenerateEmbedding({
provider: modelId,
message: `Expected ${input.length} embeddings, got ${embeddings.length}`,
});
}
return Array.from(embeddings.entries()).map(([index, embedding]) => ({
index,
embedding,
content: input[index],
}));
}
private getTargetId<T extends ChunkSimilarity>(embedding: T) {
return 'docId' in embedding && typeof embedding.docId === 'string'
? embedding.docId
: 'fileId' in embedding && typeof embedding.fileId === 'string'
? embedding.fileId
: '';
}
private async getEmbeddingRelevance<
Chunk extends ChunkSimilarity = ChunkSimilarity,
>(
query: string,
embeddings: Chunk[],
options?: EmbeddingCallOptionsInput
): Promise<ReRankResult> {
const normalizedOptions = normalizeEmbeddingCallOptions(options);
if (!embeddings.length) return [];
const rerankRequest: CopilotRerankRequest = {
query,
candidates: embeddings.map((embedding, index) => ({
id: String(index),
text: embedding.content,
})),
};
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) => {
const chunk = embeddings[i];
return {
chunk: chunk.chunk,
targetId: this.getTargetId(chunk),
score: Math.max(score, 1 - (chunk.distance || -Infinity)),
};
});
} catch (error) {
this.logger.error('Failed to parse rerank results', error);
// silent error, will fallback to default sorting in parent method
return [];
}
}
override async reRank<Chunk extends ChunkSimilarity = ChunkSimilarity>(
query: string,
embeddings: Chunk[],
topK: number,
options?: EmbeddingCallOptionsInput
): Promise<Chunk[]> {
const normalizedOptions = normalizeEmbeddingCallOptions(options);
// search in context and workspace may find same chunks, de-duplicate them
const { deduped: dedupedEmbeddings } = embeddings.reduce(
(acc, e) => {
const key = `${this.getTargetId(e)}:${e.chunk}`;
if (!acc.seen.has(key)) {
acc.seen.add(key);
acc.deduped.push(e);
}
return acc;
},
{ deduped: [] as Chunk[], seen: new Set<string>() }
);
const sortedEmbeddings = dedupedEmbeddings.toSorted(
(a, b) => (a.distance ?? Infinity) - (b.distance ?? Infinity)
);
const chunks = sortedEmbeddings.reduce(
(acc, e) => {
const targetId = this.getTargetId(e);
const key = `${targetId}:${e.chunk}`;
acc[key] = e;
return acc;
},
{} as Record<string, Chunk>
);
try {
// The rerank prompt is expected to handle the full deduped candidate list.
const ranks = await this.getEmbeddingRelevance(
query,
sortedEmbeddings,
normalizedOptions
);
if (sortedEmbeddings.length !== ranks.length) {
// llm return wrong result, fallback to default sorting
this.logger.warn(
`Batch size mismatch: expected ${sortedEmbeddings.length}, got ${ranks.length}`
);
return await super.reRank(
query,
dedupedEmbeddings,
topK,
normalizedOptions
);
}
const highConfidenceChunks = ranks
.flat()
.toSorted((a, b) => b.score - a.score)
.filter(r => r.score > 0.5)
.map(r => chunks[`${r.targetId}:${r.chunk}`])
.filter(Boolean);
this.logger.verbose(
`ReRank completed: ${highConfidenceChunks.length} high-confidence results found, total ${sortedEmbeddings.length} embeddings`,
highConfidenceChunks.length !== sortedEmbeddings.length
? JSON.stringify(ranks)
: undefined
);
return highConfidenceChunks.slice(0, topK);
} catch (error) {
this.logger.warn('ReRank failed, falling back to default sorting', error);
return await super.reRank(
query,
dedupedEmbeddings,
topK,
normalizedOptions
);
}
}
}
@Injectable()
export class CopilotEmbeddingClientService {
private client: EmbeddingClient | undefined;
constructor(
@Inject(forwardRef(() => CapabilityRuntime))
private readonly runtime: EmbeddingRuntime
) {}
async refresh() {
const client = new ProductionEmbeddingClient(this.runtime);
await client.configured();
this.client = client;
return this.client;
}
getClient() {
return this.client;
}
}
export class MockEmbeddingClient extends EmbeddingClient {
private embed(content: string) {
const seed = createHash('sha256').update(content).digest();
return Array.from({ length: EMBEDDING_DIMENSIONS }, (_, index) => {
const byte = seed[index % seed.length];
return byte / 255;
});
}
async getEmbeddings(input: string[]): Promise<Embedding[]> {
return input.map((content, i) => ({
index: i,
content,
embedding: this.embed(content),
}));
}
}
packages/backend/server/src/plugins/copilot/embedding/index.ts
+6 -4
View File
@@ -1,4 +1,6 @@
export { CopilotEmbeddingClientService, MockEmbeddingClient } from './client';
export { CopilotEmbeddingJob } from './job';
export type { Chunk, DocFragment } from './types';
export { EmbeddingClient } from './types';
export { NativeEmbeddingService } from './native';
export { CopilotRerankService } from './rerank';
export {
EMBEDDING_RERANK_RUNTIME,
type EmbeddingRerankRuntime,
} from './route-context';
packages/backend/server/src/plugins/copilot/embedding/job.ts
-674
View File
@@ -1,674 +0,0 @@
import { Injectable, Logger } from '@nestjs/common';
import {
BlobNotFound,
CallMetric,
CopilotContextFileNotSupported,
EventBus,
JobQueue,
mapAnyError,
OneDay,
OnEvent,
OnJob,
} from '../../../base';
import { DocReader } from '../../../core/doc';
import { WorkspaceBlobStorage } from '../../../core/storage';
import { readAllDocIdsFromWorkspaceSnapshot } from '../../../core/utils/blocksuite';
import { Models } from '../../../models';
import { CopilotStorage } from '../storage';
import { readStream } from '../utils';
import { CopilotEmbeddingClientService } from './client';
import type { Chunk, DocFragment, EmbeddingCallOptions } from './types';
import { EmbeddingClient } from './types';
@Injectable()
export class CopilotEmbeddingJob {
private readonly logger = new Logger(CopilotEmbeddingJob.name);
private readonly workspaceJobAbortController: Map<string, AbortController> =
new Map();
private supportEmbedding = false;
private client: EmbeddingClient | undefined;
constructor(
private readonly embeddingClients: CopilotEmbeddingClientService,
private readonly doc: DocReader,
private readonly event: EventBus,
private readonly models: Models,
private readonly queue: JobQueue,
private readonly storage: CopilotStorage,
private readonly workspaceStorage: WorkspaceBlobStorage
) {}
@OnEvent('config.init')
async onConfigInit() {
await this.setup();
}
@OnEvent('config.changed')
async onConfigChanged() {
await this.setup();
}
private async setup() {
this.supportEmbedding =
await this.models.copilotContext.checkEmbeddingAvailable();
if (this.supportEmbedding) {
this.client = await this.embeddingClients.refresh();
}
}
// public this client to allow overriding in tests
get embeddingClient() {
return this.client as EmbeddingClient;
}
@CallMetric('ai', 'addFileEmbeddingQueue')
async addFileEmbeddingQueue(
file: Jobs['copilot.embedding.files'],
options?: { priority?: number }
) {
if (!this.supportEmbedding) return;
await this.queue.add('copilot.embedding.files', file, {
priority: options?.priority,
});
}
@CallMetric('ai', 'addBlobEmbeddingQueue')
async addBlobEmbeddingQueue(blob: Jobs['copilot.embedding.blobs']) {
if (!this.supportEmbedding) return;
await this.queue.add('copilot.embedding.blobs', blob);
}
@OnEvent('workspace.doc.embedding')
async addDocEmbeddingQueue(
docs: Events['workspace.doc.embedding'],
options?: { contextId: string; priority: number }
) {
if (!this.supportEmbedding) return;
for (const { workspaceId, docId } of docs) {
const jobId = `workspace:embedding:${workspaceId}:${docId}`;
const job = await this.queue.get(jobId, 'copilot.embedding.docs');
// if the job exists and is older than 5 minute, remove it
if (job && job.timestamp + 5 * 60 * 1000 < Date.now()) {
this.logger.verbose(`Removing old embedding job ${jobId}`);
await this.queue.remove(jobId, 'copilot.embedding.docs');
}
await this.queue.add(
'copilot.embedding.docs',
{
contextId: options?.contextId,
workspaceId,
docId,
},
{
jobId: `workspace:embedding:${workspaceId}:${docId}`,
priority: options?.priority ?? 1,
timestamp: Date.now(),
}
);
}
}
@OnEvent('workspace.updated')
async onWorkspaceConfigUpdate({
id,
enableDocEmbedding,
}: Events['workspace.updated']) {
// trigger workspace embedding
this.event.emit('workspace.embedding', {
workspaceId: id,
enableDocEmbedding,
});
}
@OnEvent('workspace.embedding')
async addWorkspaceEmbeddingQueue({
workspaceId,
enableDocEmbedding,
}: Events['workspace.embedding']) {
if (!this.supportEmbedding || !this.embeddingClient) return;
if (enableDocEmbedding === undefined) {
enableDocEmbedding =
await this.models.workspace.allowEmbedding(workspaceId);
}
if (enableDocEmbedding) {
const toBeEmbedDocIds =
await this.models.copilotWorkspace.findDocsToEmbed(workspaceId);
if (!toBeEmbedDocIds.length) {
return;
}
// filter out trashed docs
const rootSnapshot = await this.models.doc.getSnapshot(
workspaceId,
workspaceId
);
if (!rootSnapshot) {
this.logger.warn(
`Root snapshot for workspace ${workspaceId} not found, skipping embedding.`
);
return;
}
const allDocIds = new Set(
readAllDocIdsFromWorkspaceSnapshot(rootSnapshot.blob)
);
this.logger.log(
`Trigger embedding for ${toBeEmbedDocIds.length} docs in workspace ${workspaceId}`
);
const finalToBeEmbedDocIds = toBeEmbedDocIds.filter(docId =>
allDocIds.has(docId)
);
for (const docId of finalToBeEmbedDocIds) {
await this.queue.add(
'copilot.embedding.docs',
{
workspaceId,
docId,
},
{
jobId: `workspace:embedding:${workspaceId}:${docId}`,
priority: 1,
}
);
}
} else {
const controller = this.workspaceJobAbortController.get(workspaceId);
if (controller) {
controller.abort();
this.workspaceJobAbortController.delete(workspaceId);
}
}
}
@OnJob('copilot.embedding.updateDoc')
async addDocEmbeddingQueueFromEvent(
doc: Jobs['copilot.embedding.updateDoc']
) {
if (!this.supportEmbedding || !this.embeddingClient) return;
await this.queue.add(
'copilot.embedding.docs',
{
workspaceId: doc.workspaceId,
docId: doc.docId,
},
{
jobId: `workspace:embedding:${doc.workspaceId}:${doc.docId}`,
priority: 2,
}
);
}
private async deleteDocEmbedding(doc: {
workspaceId: string;
docId: string;
}) {
await this.queue.remove(
`workspace:embedding:${doc.workspaceId}:${doc.docId}`,
'copilot.embedding.docs'
);
await this.models.copilotContext.purgeWorkspaceEmbedding(
doc.workspaceId,
doc.docId
);
}
@OnJob('copilot.embedding.reconcileDocumentCleanup')
async reconcileDocumentCleanup({
workspaceId,
docId,
cleanupVersion,
}: Jobs['copilot.embedding.reconcileDocumentCleanup']) {
const root = await this.doc.getDoc(workspaceId, workspaceId);
if (!root) {
throw new Error(`workspace root ${workspaceId} not found`);
}
const live = readAllDocIdsFromWorkspaceSnapshot(root.bin, true).includes(
docId
);
if (live) {
if (!(await this.doc.getDoc(workspaceId, docId))) {
throw new Error(`restored document ${workspaceId}/${docId} not found`);
}
await this.addDocEmbeddingQueueFromEvent({ workspaceId, docId });
} else {
await this.deleteDocEmbedding({ workspaceId, docId });
}
await this.queue.add('backendRuntime.ackDocumentCleanupEffect', {
workspaceId,
docId,
cleanupVersion,
effect: 'copilot',
});
}
private async readCopilotBlob(
userId: string,
workspaceId: string,
blobId: string,
fileName: string
) {
const { body } = await this.storage.get(userId, workspaceId, blobId);
if (!body) throw new BlobNotFound({ spaceId: workspaceId, blobId });
const buffer = await readStream(body);
return new File([buffer], fileName);
}
private async readWorkspaceBlob(
workspaceId: string,
blobId: string,
fileName: string
) {
const { body } = await this.workspaceStorage.get(workspaceId, blobId);
if (!body) throw new BlobNotFound({ spaceId: workspaceId, blobId });
const buffer = await readStream(body);
return new File([buffer], fileName);
}
private workspaceIndexingOptions(
workspaceId: string,
signal?: AbortSignal,
userId?: string
): EmbeddingCallOptions {
return {
workspaceId,
userId,
signal,
featureKind: 'workspace_indexing',
};
}
@OnJob('copilot.embedding.files')
async embedPendingFile({
userId,
workspaceId,
contextId,
blobId,
fileId,
fileName,
}: Jobs['copilot.embedding.files']) {
if (!this.supportEmbedding || !this.embeddingClient) return;
try {
const file = await this.readCopilotBlob(
userId,
workspaceId,
blobId,
fileName
);
// no need to check if embeddings is empty, will throw internally
const chunks = await this.embeddingClient.getFileChunks(file);
const total = chunks.reduce((acc, c) => acc + c.length, 0);
for (const chunk of chunks) {
const embeddings = await this.embeddingClient.generateEmbeddings(
chunk,
this.workspaceIndexingOptions(workspaceId, undefined, userId)
);
if (contextId) {
// for context files
await this.models.copilotContext.insertFileEmbedding(
contextId,
fileId,
embeddings
);
} else {
// for workspace files
await this.models.copilotWorkspace.insertFileEmbeddings(
workspaceId,
fileId,
embeddings
);
}
}
this.event.emit('workspace.file.embed.finished', {
contextId,
workspaceId,
fileId,
chunkSize: total,
});
} catch (error: any) {
this.event.emit('workspace.file.embed.failed', {
contextId,
workspaceId,
fileId,
error: mapAnyError(error).message,
});
// passthrough error to job queue
throw error;
}
}
@OnJob('copilot.embedding.blobs')
async embedPendingBlob({
workspaceId,
contextId,
blobId,
}: Jobs['copilot.embedding.blobs']) {
if (!this.supportEmbedding || !this.embeddingClient) return;
try {
const file = await this.readWorkspaceBlob(workspaceId, blobId, 'blob');
const chunks = await this.embeddingClient.getFileChunks(file);
const total = chunks.reduce((acc, c) => acc + c.length, 0);
for (const chunk of chunks) {
const embeddings = await this.embeddingClient.generateEmbeddings(
chunk,
this.workspaceIndexingOptions(workspaceId)
);
await this.models.copilotWorkspace.insertBlobEmbeddings(
workspaceId,
blobId,
embeddings
);
}
if (contextId) {
this.event.emit('workspace.blob.embed.finished', {
contextId,
blobId,
chunkSize: total,
});
}
} catch (error: any) {
if (contextId) {
this.event.emit('workspace.blob.embed.failed', {
contextId,
blobId,
error: mapAnyError(error).message,
});
}
throw error;
}
}
private async getDocFragment(
workspaceId: string,
docId: string
): Promise<DocFragment | null> {
const docContent = await this.doc.getFullDocContent(workspaceId, docId);
const authors = await this.models.doc.getAuthors(workspaceId, docId);
if (docContent && authors) {
const { title, summary } = docContent;
const { createdAt, updatedAt, createdByUser, updatedByUser } = authors;
return {
title: title || 'Untitled',
summary,
createdAt: createdAt.toDateString(),
updatedAt: updatedAt.toDateString(),
createdBy: createdByUser?.name,
updatedBy: updatedByUser?.name,
};
}
return null;
}
private formatDocChunks(chunks: Chunk[], fragment: DocFragment): Chunk[] {
return chunks.map(chunk => ({
index: chunk.index,
content: [
`Title: ${fragment.title}`,
`Created at: ${fragment.createdAt}`,
`Updated at: ${fragment.updatedAt}`,
fragment.createdBy ? `Created by: ${fragment.createdBy}` : undefined,
fragment.updatedBy ? `Updated by: ${fragment.updatedBy}` : undefined,
chunk.content,
]
.filter(Boolean)
.join('\n'),
}));
}
private getWorkspaceSignal(workspaceId: string) {
let controller = this.workspaceJobAbortController.get(workspaceId);
if (!controller) {
controller = new AbortController();
this.workspaceJobAbortController.set(workspaceId, controller);
}
return controller.signal;
}
private normalize(s: string) {
return s.replaceAll(/[\p{White_Space}]+/gu, '');
}
@OnJob('copilot.embedding.docs')
async embedPendingDocs({
contextId,
workspaceId,
docId,
}: Jobs['copilot.embedding.docs']) {
if (!this.supportEmbedding || !this.embeddingClient) return;
if (workspaceId === docId || docId.includes('$')) return;
const signal = this.getWorkspaceSignal(workspaceId);
try {
const hasNewDoc = await this.models.doc.exists(
workspaceId,
docId.split(':space:')[1] || ''
);
const needEmbedding =
await this.models.copilotWorkspace.checkDocNeedEmbedded(
workspaceId,
docId
);
this.logger.debug(
`Check if doc ${docId} in workspace ${workspaceId} needs embedding: ${needEmbedding}`
);
if (needEmbedding) {
if (signal.aborted) {
this.logger.debug(
`Doc ${docId} in workspace ${workspaceId} is aborted, skipping embedding.`
);
return;
}
// if doc id deprecated, skip embedding and fulfill empty embedding
const fragment = !hasNewDoc
? await this.getDocFragment(workspaceId, docId)
: undefined;
if (!hasNewDoc && fragment) {
// fast fall for empty doc, journal is easily to create a empty doc
if (fragment.summary.trim()) {
const existsContent =
await this.models.copilotContext.getWorkspaceContent(
workspaceId,
docId
);
if (
existsContent &&
this.normalize(existsContent) === this.normalize(fragment.summary)
) {
this.logger.debug(
`Doc ${docId} in workspace ${workspaceId} has no content change, skipping embedding.`
);
if (contextId) {
this.event.emit('workspace.doc.embed.finished', {
contextId,
docId,
});
}
return;
}
const embeddings = await this.embeddingClient.getFileEmbeddings(
new File(
[fragment.summary],
`${fragment.title || 'Untitled'}.md`
),
chunks => this.formatDocChunks(chunks, fragment),
this.workspaceIndexingOptions(workspaceId, signal)
);
for (const chunks of embeddings) {
await this.models.copilotContext.insertWorkspaceEmbedding(
workspaceId,
docId,
chunks
);
}
this.logger.debug(
`Doc ${docId} in workspace ${workspaceId} has summary, embedding done.`
);
} else {
// for empty doc, insert empty embedding
this.logger.debug(
`Doc ${docId} in workspace ${workspaceId} has no summary, fulfilling empty embedding.`
);
await this.models.copilotContext.fulfillEmptyEmbedding(
workspaceId,
docId
);
}
} else {
this.logger.debug(
`Doc ${docId} in workspace ${workspaceId} has no fragment, fulfilling empty embedding.`
);
await this.models.copilotContext.fulfillEmptyEmbedding(
workspaceId,
docId
);
}
}
if (contextId) {
this.event.emit('workspace.doc.embed.finished', {
contextId,
docId,
});
}
} catch (error: any) {
if (contextId) {
this.event.emit('workspace.doc.embed.failed', {
contextId,
docId,
});
}
if (
error instanceof CopilotContextFileNotSupported &&
error.message.includes('no content found')
) {
this.logger.debug(
`Doc ${docId} in workspace ${workspaceId} has no content, fulfilling empty embedding.`
);
// if the doc is empty, we still need to fulfill the embedding
await this.models.copilotContext.fulfillEmptyEmbedding(
workspaceId,
docId
);
return;
}
// log error and skip the job
this.logger.error(
`Error embedding doc ${docId} in workspace ${workspaceId}`,
error
);
}
}
@OnJob('copilot.embedding.cleanupTrashedDocEmbeddings')
async cleanupTrashedDocEmbeddings({
workspaceId,
}: Jobs['copilot.embedding.cleanupTrashedDocEmbeddings']) {
const workspace = await this.models.workspace.get(workspaceId);
if (!workspace) {
this.logger.warn(`workspace ${workspaceId} not found`);
return;
}
const oneMonthAgo = new Date(Date.now() - OneDay * 30);
const snapshot = await this.models.doc.getSnapshot(
workspaceId,
workspaceId
);
if (!snapshot) {
// maybe local workspace or empty workspace
this.logger.verbose(`workspace root snapshot ${workspaceId} not found`);
// mark last check time to avoid repeated checking
await this.models.workspace.update(
workspaceId,
{ lastCheckEmbeddings: new Date() },
false
);
return;
} else if (
// always check if never cleared
workspace.lastCheckEmbeddings > new Date(0) &&
snapshot.updatedAt < oneMonthAgo
) {
this.logger.verbose(
`workspace ${workspaceId} is too old, skipping embeddings cleanup`
);
await this.models.workspace.update(
workspaceId,
{ lastCheckEmbeddings: new Date() },
false
);
return;
}
const [docIdsInEmbedding, docIdsInSnapshots] = await Promise.all([
this.models.copilotContext.listWorkspaceDocEmbedding(workspaceId),
this.models.copilotWorkspace.listEmbeddableDocIds(workspaceId),
]);
if (!docIdsInEmbedding.length && !docIdsInSnapshots.length) {
this.logger.verbose(
`No doc embeddings and snapshots found in workspace ${workspaceId}, skipping cleanup`
);
await this.models.workspace.update(
workspaceId,
{ lastCheckEmbeddings: new Date() },
false
);
return;
}
const docIdsInWorkspace = readAllDocIdsFromWorkspaceSnapshot(snapshot.blob);
const docIdsInWorkspaceSet = new Set(docIdsInWorkspace);
const deletedDocIds = new Set(
[...docIdsInEmbedding, ...docIdsInSnapshots].filter(
docId => !docIdsInWorkspaceSet.has(docId)
)
);
for (const docId of deletedDocIds) {
const isPlaceholder = await this.models.copilotWorkspace.hasPlaceholder(
workspaceId,
docId
);
if (isPlaceholder) continue;
await this.models.copilotContext.deleteWorkspaceEmbedding(
workspaceId,
docId
);
}
await this.models.workspace.update(
workspaceId,
{ lastCheckEmbeddings: new Date() },
false
);
}
@OnEvent('workspace.updated')
async onWorkspaceUpdated({ id }: Events['workspace.updated']) {
if (!this.supportEmbedding) return;
await this.queue.add('copilot.embedding.cleanupTrashedDocEmbeddings', {
workspaceId: id,
});
}
}
packages/backend/server/src/plugins/copilot/embedding/native.ts
+201
View File
@@ -0,0 +1,201 @@
import { Injectable, OnApplicationBootstrap } from '@nestjs/common';
import { nanoid } from 'nanoid';
import { metrics } from '../../../base';
import { BackendRuntimeProvider } from '../../../core/backend-runtime';
import type { DocChunkSimilarity } from '../../../models';
import type {
RuntimeEmbeddingCandidate,
RuntimeRetrievalScope,
} from '../../../native';
import { CopilotRerankService } from './rerank';
import type { EmbeddingRouteContext } from './route-context';
@Injectable()
export class NativeEmbeddingService implements OnApplicationBootstrap {
private supportEmbedding = false;
constructor(
private readonly runtime: BackendRuntimeProvider,
private readonly rerank: CopilotRerankService
) {}
async onApplicationBootstrap() {
this.supportEmbedding = (await this.health()).enabled;
}
get canEmbedding() {
return this.supportEmbedding;
}
async health() {
const health = await this.runtime.embeddingHealth();
metrics.ai.counter('embedding_capability_check').add(1, {
state: health.state,
enabled: health.enabled,
reason: health.reason ?? 'none',
schema: String(health.schemaVersion ?? 0),
worker: health.workerRunning ? 'running' : 'stopped',
});
return health;
}
async progress(workspaceId: string) {
return await this.runtime.embeddingWorkspaceProgress(workspaceId);
}
async readSourceContent(
workspaceId: string,
sourceKind: 'document' | 'artifact',
sourceKey: string,
retrieval: RuntimeRetrievalScope,
maxChars?: number,
cursor?: string
) {
return await this.runtime.readEmbeddingSourceContent({
workspaceId,
sourceKind,
sourceKey,
retrieval,
maxChars,
cursor,
});
}
async match(
workspaceId: string,
query: string,
sourceKind: 'document' | 'artifact',
retrieval: RuntimeRetrievalScope,
limit: number,
signal?: AbortSignal
): Promise<RuntimeEmbeddingCandidate[]> {
const startedAt = performance.now();
signal?.throwIfAborted();
const requestId = nanoid();
const abort = () => {
void this.runtime
.cancelEmbeddingCandidateRequest(requestId)
.catch(() => {});
};
signal?.addEventListener('abort', abort, { once: true });
try {
const candidates = await this.runtime.matchEmbeddingCandidates({
requestId,
workspaceId,
query,
sourceKind,
retrieval,
limit,
});
signal?.throwIfAborted();
metrics.ai
.histogram('embedding_candidate_latency_ms')
.record(performance.now() - startedAt, {
corpus: sourceKind,
mode: retrieval.mode,
outcome: 'success',
});
return candidates;
} catch (error) {
metrics.ai.counter('embedding_operation_failure').add(1, {
operation: 'match',
kind: sourceKind,
code: embeddingErrorCode(error),
});
throw error;
} finally {
signal?.removeEventListener('abort', abort);
}
}
async matchWorkspaceDocCandidates(
workspaceId: string,
content: string,
topK = 5,
docIds?: string[]
): Promise<DocChunkSimilarity[]> {
const retrieval: RuntimeRetrievalScope = {
mode: docIds ? 'required' : 'workspace',
requiredDocIds: docIds ?? [],
requiredArtifactIds: [],
preferredSourceIds: [],
};
return (
await this.match(workspaceId, content, 'document', retrieval, topK * 2)
)
.filter(candidate => candidate.docId)
.map(candidate => ({
docId: candidate.docId as string,
chunk: candidate.chunk,
content: candidate.content,
distance: candidate.distance,
unitId: candidate.unitId ?? '',
visibility: (candidate.visibility ?? 'page') as
| 'page'
| 'edgeless'
| 'both',
blockId: candidate.blockId ?? undefined,
elementId: candidate.elementId ?? undefined,
frameId: candidate.frameId ?? undefined,
}));
}
async rerankWorkspaceDocs(
workspaceId: string,
content: string,
candidates: DocChunkSimilarity[],
topK = 5,
routeContext?: EmbeddingRouteContext
) {
if (!candidates.length) return [];
return await this.rerank.rerank(
content,
candidates,
topK,
workspaceId,
routeContext
);
}
async recordQueueCounts() {
const counts = await this.runtime.embeddingQueueCounts();
for (const status of [
'pending',
'running',
'retryWait',
'ready',
'failed',
] as const) {
metrics.ai
.gauge('embedding_queue_status')
.record(Number(counts[status]), { status });
}
metrics.ai
.gauge('embedding_vector_rows')
.record(Number(counts.activeVectorRows), { state: 'active' });
metrics.ai
.gauge('embedding_vector_rows')
.record(Number(counts.inactiveVectorRows), { state: 'inactive' });
metrics.ai
.gauge('embedding_index_size_bytes')
.record(Number(counts.indexBytes));
metrics.ai
.gauge('embedding_index_retry')
.record(Number(counts.retryingIndexes), { measure: 'indexes' });
metrics.ai
.gauge('embedding_index_retry')
.record(Number(counts.maxIndexRetrySeconds), {
measure: 'max_delay_seconds',
});
}
}
function embeddingErrorCode(error: unknown) {
if (!(error instanceof Error)) return 'unknown';
if (error.message.includes('resource_exceeded')) return 'resource_exceeded';
if (error.message.includes('embedding_unavailable')) return 'unavailable';
if (error.message.includes('not_found')) return 'not_found';
if (error.message.includes('disabled')) return 'disabled';
return 'failed';
}
packages/backend/server/src/plugins/copilot/embedding/realtime.ts
+57
View File
@@ -0,0 +1,57 @@
import { Injectable, OnModuleInit } from '@nestjs/common';
import { z } from 'zod';
import { Config } from '../../../base/config';
import { PermissionAccess } from '../../../core/permission';
import {
RealtimeRegistry,
realtimeWorkspaceEmbeddingProgressRoom,
registerRealtimeLiveQuery,
} from '../../../core/realtime';
import { assertCopilotEnabled } from '../availability';
import { NativeEmbeddingService } from './native';
@Injectable()
export class CopilotEmbeddingRealtimeProvider implements OnModuleInit {
constructor(
private readonly ac: PermissionAccess,
private readonly embedding: NativeEmbeddingService,
private readonly registry: RealtimeRegistry,
private readonly config: Config
) {}
onModuleInit() {
const input = z.object({ workspaceId: z.string() });
registerRealtimeLiveQuery(this.registry, {
request: {
name: 'workspace.embedding.progress.get',
input,
handle: async (user, payload) => {
await this.assertCopilot(user.id, payload.workspaceId);
const health = await this.embedding.health();
return health.enabled
? await this.embedding.progress(payload.workspaceId)
: { total: 0, embedded: 0 };
},
},
topic: {
name: 'workspace.embedding.progress.changed',
input,
authorize: async (user, payload) => {
await this.assertCopilot(user.id, payload.workspaceId);
},
room: (_user, payload) =>
realtimeWorkspaceEmbeddingProgressRoom(payload.workspaceId),
},
});
}
private async assertCopilot(userId: string, workspaceId: string) {
assertCopilotEnabled(this.config);
await this.ac
.user(userId)
.workspace(workspaceId)
.allowLocal()
.assert('Workspace.Copilot');
}
}
packages/backend/server/src/plugins/copilot/embedding/rerank.ts
+69
View File
@@ -0,0 +1,69 @@
import { Inject, Injectable } from '@nestjs/common';
import { ModuleRef } from '@nestjs/core';
import type { ChunkSimilarity } from '../../../models';
import {
EMBEDDING_RERANK_RUNTIME,
type EmbeddingRerankRuntime,
type EmbeddingRouteContext,
} from './route-context';
@Injectable()
export class CopilotRerankService {
constructor(
@Inject(ModuleRef)
private readonly moduleRef: ModuleRef
) {}
async rerank<T extends ChunkSimilarity>(
query: string,
candidates: T[],
topK: number,
workspaceId: string,
routeContext: EmbeddingRouteContext = {},
signal?: AbortSignal
): Promise<T[]> {
if (signal?.aborted) throw new Error('SEARCH_ABORTED');
if (!candidates.length) return [];
try {
const runtime = this.moduleRef.get<EmbeddingRerankRuntime>(
EMBEDDING_RERANK_RUNTIME,
{ strict: false }
);
const scores = await runtime.rerank(
'route-selected',
{
query,
candidates: candidates.map((candidate, index) => ({
id: String(index),
text: candidate.content,
})),
},
{
workspace: workspaceId,
byokLeaseId: routeContext.byokLeaseId,
featureKind: 'rerank',
signal,
}
);
if (signal?.aborted) throw new Error('SEARCH_ABORTED');
if (scores.length !== candidates.length) {
return candidates
.toSorted(
(a, b) => (a.distance ?? Infinity) - (b.distance ?? Infinity)
)
.slice(0, topK);
}
return candidates
.map((candidate, index) => ({ candidate, score: scores[index] }))
.toSorted((a, b) => b.score - a.score)
.slice(0, topK)
.map(item => item.candidate);
} catch (error) {
if (signal?.aborted) throw new Error('SEARCH_ABORTED', { cause: error });
return candidates
.toSorted((a, b) => (a.distance ?? Infinity) - (b.distance ?? Infinity))
.slice(0, topK);
}
}
}
packages/backend/server/src/plugins/copilot/embedding/route-context.ts
+18
View File
@@ -0,0 +1,18 @@
export type EmbeddingRouteContext = {
byokLeaseId?: string;
};
export const EMBEDDING_RERANK_RUNTIME = Symbol('EMBEDDING_RERANK_RUNTIME');
export interface EmbeddingRerankRuntime {
rerank(
modelId: string,
request: { query: string; candidates: { id: string; text: string }[] },
options: {
workspace: string;
byokLeaseId?: string;
featureKind: 'rerank';
signal?: AbortSignal;
}
): Promise<number[]>;
}
packages/backend/server/src/plugins/copilot/embedding/types.ts
-252
View File
@@ -1,252 +0,0 @@
import { File } from 'node:buffer';
import { z } from 'zod';
import { CopilotContextFileNotSupported } from '../../../base';
import type { PageDocContent } from '../../../core/utils/blocksuite';
import { ChunkSimilarity, Embedding } from '../../../models';
import { parseDoc } from '../../../native';
import type { ByokFeatureKind } from '../byok/types';
declare global {
interface Events {
'workspace.embedding': {
workspaceId: string;
enableDocEmbedding?: boolean;
};
'workspace.blob.embed.finished': {
contextId: string;
blobId: string;
chunkSize: number;
};
'workspace.blob.embed.failed': {
contextId: string;
blobId: string;
error: string;
};
'workspace.doc.embedding': Array<{
workspaceId: string;
docId: string;
}>;
'workspace.doc.embed.failed': {
contextId: string;
docId: string;
};
'workspace.doc.embed.finished': {
contextId: string;
docId: string;
};
'workspace.file.embed.finished': {
contextId?: string;
workspaceId: string;
fileId: string;
chunkSize: number;
};
'workspace.file.embed.failed': {
contextId?: string;
workspaceId: string;
fileId: string;
error: string;
};
}
interface Jobs {
'copilot.embedding.docs': {
contextId?: string;
workspaceId: string;
docId: string;
};
'copilot.embedding.updateDoc': {
workspaceId: string;
docId: string;
};
'copilot.embedding.reconcileDocumentCleanup': {
workspaceId: string;
docId: string;
cleanupVersion: string;
};
'copilot.embedding.files': {
contextId?: string;
userId: string;
workspaceId: string;
blobId: string;
fileId: string;
fileName: string;
};
'copilot.embedding.blobs': {
contextId?: string;
workspaceId: string;
blobId: string;
};
'copilot.embedding.cleanupTrashedDocEmbeddings': {
workspaceId: string;
};
}
}
export type DocFragment = PageDocContent & {
createdAt: string;
createdBy?: string;
updatedAt: string;
updatedBy?: string;
};
export type Chunk = {
index: number;
content: string;
};
export type EmbeddingCallOptions = {
signal?: AbortSignal;
userId?: string;
workspaceId?: string;
byokLeaseId?: string;
featureKind?: Extract<
ByokFeatureKind,
'embedding' | 'workspace_indexing' | 'rerank'
>;
};
export type EmbeddingCallOptionsInput = AbortSignal | EmbeddingCallOptions;
export type EmbeddingRouteContext = Pick<
EmbeddingCallOptions,
'userId' | 'byokLeaseId'
>;
export function normalizeEmbeddingCallOptions(
options?: EmbeddingCallOptionsInput
): EmbeddingCallOptions {
if (!options) {
return {};
}
if ('aborted' in options && 'addEventListener' in options) {
return { signal: options };
}
return options;
}
export abstract class EmbeddingClient {
async configured() {
return true;
}
async getFileEmbeddings(
file: File,
chunkMapper: (chunk: Chunk[]) => Chunk[],
options?: EmbeddingCallOptionsInput
): Promise<Embedding[][]> {
const normalizedOptions = normalizeEmbeddingCallOptions(options);
const chunks = await this.getFileChunks(file, normalizedOptions.signal);
const chunkedEmbeddings = await Promise.all(
chunks.map(chunk =>
this.generateEmbeddings(chunkMapper(chunk), normalizedOptions)
)
);
return chunkedEmbeddings;
}
async getFileChunks(file: File, signal?: AbortSignal): Promise<Chunk[][]> {
const buffer = Buffer.from(await file.arrayBuffer());
let doc;
try {
doc = await parseDoc(file.name, buffer);
} catch (e: any) {
throw new CopilotContextFileNotSupported({
fileName: file.name,
message: e?.message || e?.toString?.() || 'format not supported',
});
}
if (doc && !signal?.aborted) {
if (!doc.chunks.length) {
throw new CopilotContextFileNotSupported({
fileName: file.name,
message: 'no content found',
});
}
const input = doc.chunks.toSorted((a, b) => a.index - b.index);
// chunk input into 128 every array
const chunks: Chunk[][] = [];
for (let i = 0; i < input.length; i += 128) {
chunks.push(input.slice(i, i + 128));
}
return chunks;
}
throw new CopilotContextFileNotSupported({
fileName: file.name,
message: 'failed to parse file',
});
}
async generateEmbeddings(
chunks: Chunk[],
options?: EmbeddingCallOptionsInput
): Promise<Embedding[]> {
const normalizedOptions = normalizeEmbeddingCallOptions(options);
const retry = 3;
let embeddings: Embedding[] = [];
let error = null;
for (let i = 0; i < retry; i++) {
try {
embeddings = await this.getEmbeddings(
chunks.map(c => c.content),
normalizedOptions
);
break;
} catch (e) {
error = e;
}
}
if (error) throw error;
// fix the index of the embeddings
return embeddings.map(e => ({ ...e, index: chunks[e.index].index }));
}
async reRank<Chunk extends ChunkSimilarity = ChunkSimilarity>(
_query: string,
embeddings: Chunk[],
topK: number,
_options?: EmbeddingCallOptionsInput
): Promise<Chunk[]> {
// sort by distance with ascending order
return embeddings
.toSorted((a, b) => (a.distance ?? Infinity) - (b.distance ?? Infinity))
.slice(0, topK);
}
async getEmbedding(query: string, options?: EmbeddingCallOptionsInput) {
const embedding = await this.getEmbeddings([query], options);
return embedding?.[0]?.embedding;
}
abstract getEmbeddings(
input: string[],
options?: EmbeddingCallOptionsInput
): Promise<Embedding[]>;
}
const ReRankItemSchema = z.object({
chunk: z.number().describe('The chunk index of the search result.'),
targetId: z.string().describe('The id of the target.'),
score: z
.number()
.min(0)
.max(10)
.describe(
'The relevance score of the results should be 0-10, with 0 being the least relevant and 10 being the most relevant.'
),
});
export type ReRankResult = z.infer<typeof ReRankItemSchema>[];
packages/backend/server/src/plugins/copilot/index.ts
+1 -8
View File
@@ -17,7 +17,6 @@ import { McpCredentialService } from './mcp/credential';
import { McpCredentialResolver } from './mcp/resolver';
import {
COPILOT_API_PROVIDERS,
COPILOT_CONTEXT_REALTIME_PROVIDERS,
COPILOT_FEATURE_PROVIDERS,
COPILOT_KERNEL_PROVIDERS,
COPILOT_TRANSCRIPT_REALTIME_PROVIDERS,
@@ -49,17 +48,11 @@ export class CopilotAvailabilityModule {}
export class CopilotKernelModule {}
@Module({
imports: [PermissionModule, CopilotAvailabilityModule],
imports: [PermissionModule, CopilotAvailabilityModule, CopilotKernelModule],
providers: [...COPILOT_TRANSCRIPT_REALTIME_PROVIDERS],
})
export class CopilotRealtimeModule {}
@Module({
imports: [PermissionModule, CopilotAvailabilityModule],
providers: [...COPILOT_CONTEXT_REALTIME_PROVIDERS],
})
export class CopilotEmbeddingRealtimeModule {}
@Module({
imports: [...COPILOT_SHARED_IMPORTS, CopilotKernelModule],
providers: [...COPILOT_FEATURE_PROVIDERS],
packages/backend/server/src/plugins/copilot/mcp/provider.ts
+39 -91
View File
@@ -1,13 +1,10 @@
import { Injectable } from '@nestjs/common';
import { McpAccessMode } from '@prisma/client';
import { pick } from 'lodash-es';
import z from 'zod/v3';
import { DocReader, DocWriter } from '../../../core/doc';
import { PermissionAccess } from '../../../core/permission';
import { clearEmbeddingChunk } from '../../../models';
import { IndexerService } from '../../indexer';
import { CopilotContextService } from '../context/service';
import { DocumentRetrievalService } from '../retrieval/document';
type McpTextContent = {
type: 'text';
@@ -103,8 +100,7 @@ export class WorkspaceMcpProvider {
private readonly ac: PermissionAccess,
private readonly reader: DocReader,
private readonly writer: DocWriter,
private readonly context: CopilotContextService,
private readonly indexer: IndexerService
private readonly retrieval: DocumentRetrievalService
) {}
async for(
@@ -154,105 +150,57 @@ export class WorkspaceMcpProvider {
},
});
const semanticSearch = defineTool({
name: 'semantic_search',
title: 'Semantic Search',
const docSearch = defineTool({
name: 'doc_search',
title: 'Document Search',
description:
'Retrieve conceptually related passages by performing vector-based semantic similarity search across embedded documents; use this tool only when exact keyword search fails or the user explicitly needs meaning-level matches (e.g., paraphrases, synonyms, broader concepts, recent documents).',
parser: z.object({ query: z.string() }),
'Search persisted workspace documents and return bounded passages with Page or canvas locators. Retrieval strategy is selected by the server and never includes files, blobs, attachments, or the web.',
parser: z.object({
query: z.string().trim().min(1).max(2000),
doc_ids: z.array(z.string().min(1).max(128)).max(50).optional(),
limit: z.number().int().min(1).max(20).optional(),
}),
inputSchema: {
type: 'object',
properties: {
query: { type: 'string' },
doc_ids: {
type: 'array',
items: { type: 'string' },
maxItems: 50,
},
limit: { type: 'integer', minimum: 1, maximum: 20 },
},
required: ['query'],
additionalProperties: false,
},
execute: async ({ query }, options) => {
const trimmed = query.trim();
if (!trimmed) {
return toolError('Query is required for semantic search.');
}
const chunks = await this.context.matchWorkspaceDocs(
workspaceId,
trimmed,
5,
execute: async ({ query, doc_ids, limit }, options) => {
const result = await this.retrieval.search(
{ user: userId, workspace: workspaceId },
query,
doc_ids,
limit ?? 10,
options.signal
);
const abortedAfterMatch = abortIfNeeded(options.signal);
if (abortedAfterMatch) return abortedAfterMatch;
const docs = await this.ac
.user(userId)
.workspace(workspaceId)
.docs(
chunks.filter(chunk => 'docId' in chunk),
'Doc.Read'
);
const abortedAfterDocs = abortIfNeeded(options.signal);
if (abortedAfterDocs) return abortedAfterDocs;
if (!docs || docs.length === 0) {
return toolText('No matching documents found.');
}
return {
content: docs.map(doc => ({
type: 'text',
text: clearEmbeddingChunk(doc).content,
})),
};
return toolText(
JSON.stringify({
retrieval_mode: result.retrievalMode,
degraded_reason: result.degradedReason,
hits: result.hits.map(hit => ({
doc_id: hit.docId,
title: hit.title,
excerpt: hit.excerpt,
visibility: hit.visibility,
block_id: hit.blockId,
element_id: hit.elementId,
frame_id: hit.frameId,
})),
})
);
},
});
const keywordSearch = defineTool({
name: 'keyword_search',
title: 'Keyword Search',
description:
'Fuzzy search all workspace documents for the exact keyword or phrase supplied and return passages ranked by textual match. Use this tool by default whenever a straightforward term-based or keyword-base lookup is sufficient.',
parser: z.object({ query: z.string() }),
inputSchema: {
type: 'object',
properties: {
query: { type: 'string' },
},
required: ['query'],
additionalProperties: false,
},
execute: async ({ query }, options) => {
const trimmed = query.trim();
if (!trimmed) return toolError('Query is required for keyword search.');
let docs = await this.indexer.searchDocsByKeyword(workspaceId, trimmed);
const abortedAfterSearch = abortIfNeeded(options.signal);
if (abortedAfterSearch) return abortedAfterSearch;
docs = await this.ac
.user(userId)
.workspace(workspaceId)
.docs(docs, 'Doc.Read');
const abortedAfterDocs = abortIfNeeded(options.signal);
if (abortedAfterDocs) return abortedAfterDocs;
if (!docs || docs.length === 0) {
return toolText('No matching documents found.');
}
return {
content: docs.map(doc => ({
type: 'text',
text: JSON.stringify(pick(doc, 'docId', 'title', 'createdAt')),
})),
};
},
});
const tools = [readDocument, semanticSearch, keywordSearch];
const tools = [readDocument, docSearch];
if (
accessMode === McpAccessMode.READ_WRITE &&
packages/backend/server/src/plugins/copilot/module-providers.ts
+21 -22
View File
@@ -4,23 +4,26 @@ import { CompatHistoryProjector } from './compat/history-projector';
import { HistoryPromptPreloadProjector } from './compat/history-prompt-preload-projector';
import { HistoryVisibilityPolicy } from './compat/history-visibility-policy';
import { CompatSubmissionStore } from './compat/submission-store';
import {
CopilotContextResolver,
CopilotContextRootResolver,
CopilotContextService,
CopilotEmbeddingRealtimeProvider,
} from './context';
import { ConversationInboxService } from './conversation/inbox';
import { ConversationPolicy } from './conversation/policy';
import { ConversationStore } from './conversation/store';
import { CopilotCronJobs } from './cron';
import { DelegatedEditorRealtimeProvider } from './delegated/realtime';
import { DelegatedEditorService } from './delegated/service';
import {
CopilotEmbeddingClientService,
CopilotEmbeddingJob,
CopilotRerankService,
EMBEDDING_RERANK_RUNTIME,
NativeEmbeddingService,
} from './embedding';
import { CopilotEmbeddingRealtimeProvider } from './embedding/realtime';
import { WorkspaceMcpProvider } from './mcp/provider';
import { PromptService } from './prompt';
import { CopilotResolver, UserCopilotResolver } from './resolver';
import { ArtifactRetrievalService } from './retrieval/artifact';
import {
DOCUMENT_VECTOR_SEARCH,
DocumentRetrievalService,
} from './retrieval/document';
import { ActionRuntimeBridge } from './runtime/action-runtime-bridge';
import { CapabilityRuntime } from './runtime/capability-runtime';
import { CopilotRuntimeEventConsumer } from './runtime/copilot-runtime-event-consumer';
@@ -61,14 +64,19 @@ export const COPILOT_RUNTIME_PROVIDERS = [
HistoryPromptPreloadProjector,
CompatSubmissionStore,
HistoryVisibilityPolicy,
CopilotContextService,
CopilotEmbeddingClientService,
NativeEmbeddingService,
CopilotRerankService,
PromptService,
{ provide: DOCUMENT_VECTOR_SEARCH, useExisting: NativeEmbeddingService },
DocumentRetrievalService,
ArtifactRetrievalService,
DelegatedEditorService,
ActionRuntimeBridge,
CopilotRuntimeEventConsumer,
PromptRuntime,
ConversationHost,
CapabilityRuntime,
{ provide: EMBEDDING_RERANK_RUNTIME, useExisting: CapabilityRuntime },
ToolRuntime,
AttachmentMaterializer,
AttachmentAdmissionHost,
@@ -79,18 +87,11 @@ export const COPILOT_RUNTIME_PROVIDERS = [
TurnPersistence,
];
export const COPILOT_CONTEXT_REALTIME_PROVIDERS = [
CopilotEmbeddingRealtimeProvider,
];
export const COPILOT_CONTEXT_PROVIDERS = [
CopilotContextResolver,
...COPILOT_CONTEXT_REALTIME_PROVIDERS,
];
export const COPILOT_TRANSCRIPT_REALTIME_PROVIDERS = [
CopilotTranscriptionReader,
CopilotTranscriptRealtimeProvider,
CopilotEmbeddingRealtimeProvider,
DelegatedEditorRealtimeProvider,
];
export const COPILOT_TRANSCRIPT_PROVIDERS = [
@@ -108,11 +109,10 @@ export const COPILOT_WORKSPACE_PROVIDERS = [
export const COPILOT_RESOLVER_PROVIDERS = [
CopilotResolver,
UserCopilotResolver,
CopilotContextRootResolver,
WorkspaceByokResolver,
];
export const COPILOT_JOB_PROVIDERS = [CopilotEmbeddingJob, CopilotCronJobs];
export const COPILOT_JOB_PROVIDERS = [CopilotCronJobs];
export const COPILOT_MCP_PROVIDERS = [WorkspaceMcpProvider];
@@ -123,7 +123,6 @@ export const COPILOT_KERNEL_PROVIDERS = [
export const COPILOT_FEATURE_PROVIDERS = [
TurnOrchestrator,
...COPILOT_CONTEXT_PROVIDERS,
...COPILOT_TRANSCRIPT_PROVIDERS,
...COPILOT_WORKSPACE_PROVIDERS,
...COPILOT_JOB_PROVIDERS,
packages/backend/server/src/plugins/copilot/providers/types.ts
+10 -4
View File
@@ -11,6 +11,7 @@ import {
type StreamObject,
StreamObjectSchema,
} from '../runtime/contracts/runtime-event-contract';
import { RetrievalScopeSchema } from '../runtime/contracts/shared';
// Owner map:
// - provider/profile/config schemas in this file are backend host ingress.
@@ -74,17 +75,21 @@ export const VertexSchema: JSONSchema = {
export const PromptToolsSchema = z
.enum([
'blobRead',
'artifactRead',
'artifactSearch',
'codeArtifact',
'conversationSummary',
// work with indexer
'docRead',
'docCanvasRead',
'docSearch',
'docCreate',
'docUpdate',
'docUpdateMeta',
'docKeywordSearch',
// work with embeddings
'docSemanticSearch',
'frontendGetEditorState',
'frontendReadSelection',
'frontendReadNodes',
'frontendSnapshotDocument',
// work with exa/model internal tools
'webSearch',
// artifact tools
@@ -280,6 +285,7 @@ const CopilotProviderOptionsSchema = z.object({
'transcript',
])
.optional(),
retrievalScope: RetrievalScopeSchema.optional(),
});
export const CopilotChatOptionsSchema = CopilotProviderOptionsSchema.merge(
packages/backend/server/src/plugins/copilot/providers/utils.ts
+7 -35
View File
@@ -174,8 +174,8 @@ export class TextStreamParser {
result += `\nCrawling the web "${chunk.input.url}"\n`;
break;
}
case 'doc_keyword_search': {
result += `\nSearching the keyword "${chunk.input.query}"\n`;
case 'doc_search': {
result += `\nSearching workspace documents for "${chunk.input.query}"\n`;
break;
}
case 'doc_read': {
@@ -196,27 +196,11 @@ export class TextStreamParser {
);
result = this.addPrefix(result);
switch (chunk.toolName) {
case 'doc_semantic_search': {
const output = chunk.output;
if (Array.isArray(output)) {
result += `\nFound ${output.length} document${output.length !== 1 ? 's' : ''} related to “${chunk.input.query}”.\n`;
} else if (typeof output === 'string') {
result += `\n${output}\n`;
} else {
const message = asRecord(output)?.message;
this.logger.warn(
`Unexpected result type for doc_semantic_search: ${
typeof message === 'string' ? message : 'Unknown error'
}`
);
}
break;
}
case 'doc_keyword_search': {
const output = chunk.output;
if (Array.isArray(output)) {
result += `\nFound ${output.length} document${output.length !== 1 ? 's' : ''} related to “${chunk.input.query}”.\n`;
result += `\n${this.getKeywordSearchLinks(output)}\n`;
case 'doc_search': {
const output = asRecord(chunk.output);
const hits = output?.hits;
if (Array.isArray(hits)) {
result += `\nFound ${hits.length} document${hits.length !== 1 ? 's' : ''} related to “${chunk.input.query}”.\n`;
}
break;
}
@@ -287,18 +271,6 @@ export class TextStreamParser {
}, '');
return links;
}
private getKeywordSearchLinks(
list: {
docId: string;
title: string;
}[]
): string {
const links = list.reduce((acc, result) => {
return acc + `\n\n[${result.title}](${result.docId})\n\n`;
}, '');
return links;
}
}
export class StreamObjectParser {
packages/backend/server/src/plugins/copilot/resolver.ts
+3
View File
@@ -236,6 +236,9 @@ class ChatMessageType implements Partial<ChatMessage> {
@Field(() => GraphQLJSON, { nullable: true })
params!: Record<string, string> | undefined;
@Field(() => GraphQLJSON, { nullable: true })
scopeSnapshot!: ChatMessage['scopeSnapshot'];
@Field(() => Date)
createdAt!: Date;
}
packages/backend/server/src/plugins/copilot/retrieval/artifact.ts
+187
View File
@@ -0,0 +1,187 @@
import { Injectable } from '@nestjs/common';
import { PrismaClient } from '@prisma/client';
import { AccessDenied } from '../../../base';
import { PermissionAccess } from '../../../core/permission';
import type { RuntimeRetrievalScope } from '../../../native';
import { NativeEmbeddingService } from '../embedding/native';
@Injectable()
export class ArtifactRetrievalService {
constructor(
private readonly access: PermissionAccess,
private readonly embedding: NativeEmbeddingService,
private readonly db: PrismaClient
) {}
private async authorize(userId: string, workspaceId: string) {
return await this.access
.user(userId)
.workspace(workspaceId)
.allowLocal()
.can('Workspace.Read');
}
async search(options: {
userId: string;
workspaceId: string;
query: string;
retrieval: RuntimeRetrievalScope;
limit: number;
messageId?: string;
signal?: AbortSignal;
}) {
if (!(await this.authorize(options.userId, options.workspaceId))) {
throw new AccessDenied();
}
let degraded = false;
let matched: Awaited<ReturnType<NativeEmbeddingService['match']>> = [];
try {
matched = await this.embedding.match(
options.workspaceId,
options.query,
'artifact',
options.retrieval,
options.limit,
options.signal
);
} catch (error) {
if (options.signal?.aborted) throw error;
degraded = true;
}
const matchedIds = new Set(matched.map(hit => hit.artifactId));
const missingRequired =
options.retrieval.mode === 'required'
? options.retrieval.requiredArtifactIds
.filter(id => !matchedIds.has(id))
.slice(0, Math.max(0, options.limit - matched.length))
: [];
const directAttempts = await Promise.allSettled(
missingRequired.map(async artifactId => {
const source = await this.embedding.readSourceContent(
options.workspaceId,
'artifact',
artifactId,
options.retrieval,
20_000
);
return {
sourceKind: 'artifact',
sourceKey: artifactId,
artifactId,
content: source.content,
distance: 0,
chunk: 0,
};
})
);
const direct = directAttempts.flatMap(result =>
result.status === 'fulfilled' ? [result.value] : []
);
degraded ||= direct.length !== directAttempts.length;
const hits = [...matched, ...direct].map(hit => ({
...hit,
artifactId: hit.artifactId ?? hit.sourceKey,
}));
const metadata = await this.loadMetadata(
options.workspaceId,
hits.map(hit => hit.artifactId),
options.retrieval,
options.messageId
);
return {
hits: hits.map(hit => ({ ...hit, ...metadata.get(hit.artifactId) })),
degraded,
} as const;
}
async read(options: {
userId: string;
workspaceId: string;
artifactId: string;
retrieval: RuntimeRetrievalScope;
messageId?: string;
maxChars?: number;
cursor?: string;
}) {
if (!(await this.authorize(options.userId, options.workspaceId))) {
throw new AccessDenied();
}
const [result, metadata] = await Promise.all([
this.embedding.readSourceContent(
options.workspaceId,
'artifact',
options.artifactId,
options.retrieval,
options.maxChars,
options.cursor
),
this.loadMetadata(
options.workspaceId,
[options.artifactId],
options.retrieval,
options.messageId
),
]);
return {
...result,
name: metadata.get(options.artifactId)?.name ?? result.name,
mimeType: metadata.get(options.artifactId)?.mimeType ?? result.mimeType,
};
}
private async loadMetadata(
workspaceId: string,
artifactIds: string[],
retrieval: RuntimeRetrievalScope,
messageId?: string
) {
const ids = [...new Set(artifactIds)];
if (!ids.length) {
return new Map<string, { name?: string; mimeType: string }>();
}
const [artifacts, occurrences] = await Promise.all([
this.db.workspaceArtifact.findMany({
where: {
workspaceId,
id: { in: ids },
...(retrieval.mode === 'workspace' ? { libraryOwned: true } : {}),
},
select: {
id: true,
displayName: true,
canonicalMediaType: true,
},
}),
retrieval.mode === 'required' && messageId
? this.db.aiMessageArtifact.findMany({
where: {
workspaceId,
messageId,
artifactId: { in: ids },
role: 'attachment',
},
select: { artifactId: true, displayName: true },
})
: [],
]);
const occurrenceNames = new Map(
occurrences.map(occurrence => [
occurrence.artifactId,
occurrence.displayName ?? undefined,
])
);
return new Map(
artifacts.map(artifact => [
artifact.id,
{
name:
occurrenceNames.get(artifact.id) ??
artifact.displayName ??
undefined,
mimeType: artifact.canonicalMediaType,
},
])
);
}
}
packages/backend/server/src/plugins/copilot/retrieval/document.ts
+264
View File
@@ -0,0 +1,264 @@
import { Inject, Injectable } from '@nestjs/common';
import { Config, SearchProviderNotFound } from '../../../base';
import { PermissionAccess } from '../../../core/permission';
import type { DocVisibility } from '../../../core/utils/blocksuite';
import { type DocChunkSimilarity, Models } from '../../../models';
import { IndexerService } from '../../indexer/service';
import type { SearchDoc } from '../../indexer/types';
import type { EmbeddingRouteContext } from '../embedding/route-context';
type DocumentSearchContext =
| {
user?: string;
workspace?: string;
byokLeaseId?: string;
}
| undefined;
type DocumentVectorSearch = {
readonly canEmbedding: boolean;
matchWorkspaceDocCandidates(
workspaceId: string,
content: string,
topK?: number,
docIds?: string[]
): Promise<DocChunkSimilarity[]>;
rerankWorkspaceDocs(
workspaceId: string,
content: string,
candidates: DocChunkSimilarity[],
topK?: number,
routeContext?: EmbeddingRouteContext
): Promise<DocChunkSimilarity[]>;
};
export const DOCUMENT_VECTOR_SEARCH = Symbol('DOCUMENT_VECTOR_SEARCH');
export type DocumentSearchHit = {
docId: string;
title: string;
excerpt: string;
visibility: DocVisibility;
blockId?: string;
elementId?: string;
frameId?: string;
updatedAt?: Date;
score: number;
unitId: string;
};
type Candidate = DocumentSearchHit & { channels: Set<'lexical' | 'vector'> };
type ProjectedSearchDoc = SearchDoc &
Required<
Pick<
SearchDoc,
'unitId' | 'projectionVersion' | 'sourceHash' | 'visibility'
>
>;
function hasProjectionMetadata(hit: SearchDoc): hit is ProjectedSearchDoc {
return Boolean(
hit.unitId && hit.projectionVersion && hit.sourceHash && hit.visibility
);
}
function hasVectorProjectionMetadata(hit: DocChunkSimilarity) {
return Boolean(hit.unitId && hit.visibility);
}
@Injectable()
export class DocumentRetrievalService {
constructor(
private readonly config: Config,
private readonly ac: PermissionAccess,
private readonly indexer: IndexerService,
@Inject(DOCUMENT_VECTOR_SEARCH)
private readonly context: DocumentVectorSearch,
private readonly models: Models
) {}
async search(
options: DocumentSearchContext,
query: string,
docIds: string[] | undefined,
requestedLimit: number,
signal?: AbortSignal
) {
if (!options?.user || !options.workspace) {
throw new Error('INVALID_SEARCH_CONTEXT');
}
const userId = options.user;
const workspaceId = options.workspace;
const limit = Math.min(requestedLimit, 20);
const routeContext = {
userId,
byokLeaseId: options.byokLeaseId,
};
const [lexicalAttempt, vectorAttempt] = await Promise.allSettled([
this.config.indexer.enabled
? this.indexer
.searchDocsByKeyword(workspaceId, query, {
limit: Math.max(limit * 3, 20),
docIds,
})
.catch(error => {
if (error instanceof SearchProviderNotFound) return null;
throw error;
})
: null,
this.context.canEmbedding
? this.context.matchWorkspaceDocCandidates(
workspaceId,
query,
Math.max(limit * 3, 20),
docIds
)
: null,
]);
if (signal?.aborted) throw new Error('SEARCH_ABORTED');
const lexicalResult =
lexicalAttempt.status === 'fulfilled' ? lexicalAttempt.value : null;
const vectorResult =
vectorAttempt.status === 'fulfilled' ? vectorAttempt.value : null;
const lexical = lexicalResult
? await this.readable(
userId,
workspaceId,
lexicalResult.filter(hasProjectionMetadata)
)
: [];
const vectorScoped = (vectorResult ?? []).filter(
candidate =>
hasVectorProjectionMetadata(candidate) &&
(!docIds || docIds.includes(candidate.docId))
);
const readableVector = vectorScoped.length
? await this.readable(userId, workspaceId, vectorScoped)
: [];
let vector = null;
if (vectorResult !== null) {
try {
vector = await this.context.rerankWorkspaceDocs(
workspaceId,
query,
readableVector,
Math.max(limit * 3, 20),
routeContext
);
} catch {
vector = null;
}
}
if (signal?.aborted) throw new Error('SEARCH_ABORTED');
const metas = await this.models.doc.findMetas(
(vector ?? []).map(candidate => ({
workspaceId,
docId: candidate.docId,
})),
{ select: { title: true } }
);
const metaByDoc = new Map(
metas
.filter((meta): meta is NonNullable<typeof meta> => meta !== null)
.map(meta => [meta.docId, meta])
);
const candidates = new Map<string, Candidate>();
const merge = (
hit: DocumentSearchHit,
channel: 'lexical' | 'vector',
rank: number
) => {
const key = `${hit.docId}:${hit.unitId}`;
const score = 1 / (60 + rank);
const existing = candidates.get(key);
if (existing) {
existing.score += score;
existing.channels.add(channel);
} else {
candidates.set(key, { ...hit, score, channels: new Set([channel]) });
}
};
lexical.forEach((hit, index) =>
merge(this.fromLexical(hit), 'lexical', index + 1)
);
vector?.forEach((hit, index) => {
const meta = metaByDoc.get(hit.docId);
merge(
{
docId: hit.docId,
title: meta?.title ?? '',
excerpt: hit.content,
visibility: hit.visibility as DocVisibility,
blockId: hit.blockId,
elementId: hit.elementId,
frameId: hit.frameId,
score: 0,
unitId: hit.unitId,
},
'vector',
index + 1
);
});
if (lexicalResult === null && vector === null) {
throw new Error('SEARCH_UNAVAILABLE');
}
const perDoc = new Map<string, number>();
const hits = [...candidates.values()]
.sort(
(left, right) =>
right.score - left.score || left.unitId.localeCompare(right.unitId)
)
.filter(hit => {
const count = perDoc.get(hit.docId) ?? 0;
if (count >= 3) return false;
perDoc.set(hit.docId, count + 1);
return true;
})
.slice(0, limit)
.map(({ channels: _, ...hit }) => hit);
const hasLexical = lexicalResult !== null;
const retrievalMode =
hasLexical && vector ? 'hybrid' : hasLexical ? 'lexical' : 'vector';
return {
retrievalMode,
degradedReason:
retrievalMode === 'hybrid'
? undefined
: lexicalResult
? 'VECTOR_UNAVAILABLE'
: 'LEXICAL_UNAVAILABLE',
hits,
} as const;
}
private async readable<T extends { docId: string }>(
userId: string,
workspaceId: string,
candidates: T[]
) {
return (
(await this.ac
.user(userId)
.workspace(workspaceId)
.docs(candidates, 'Doc.Read')) ?? []
);
}
private fromLexical(hit: ProjectedSearchDoc): DocumentSearchHit {
return {
docId: hit.docId,
title: hit.title,
excerpt: hit.highlight || '',
visibility: hit.visibility as DocVisibility,
blockId: hit.blockId,
elementId: hit.elementId,
frameId: hit.frameId,
updatedAt: hit.updatedAt,
score: 0,
unitId: hit.unitId,
};
}
}
packages/backend/server/src/plugins/copilot/runtime/capability-runtime.ts
+5
View File
@@ -1,4 +1,6 @@
/* oxlint-disable import/no-cycle -- Tool callbacks can invoke nested Copilot prompts. */
import { randomUUID } from 'node:crypto';
import { Injectable } from '@nestjs/common';
import { Config } from '../../../base/config';
@@ -198,6 +200,7 @@ export class CapabilityRuntime {
options: RuntimeOptions
) {
const { request, toolSet } = await this.prepareChat(messages, options);
const runId = randomUUID();
const rawStream = this.backend.streamCopilot<
LlmToolLoopStreamEvent | CopilotRuntimeEvent
>(
@@ -218,6 +221,8 @@ export class CapabilityRuntime {
await executeToolCall(toolSet, toolRequest, {
signal: options.signal,
messages,
runId,
toolCallId: toolRequest.callId,
})
);
},
packages/backend/server/src/plugins/copilot/runtime/contracts/shared.ts
+49
View File
@@ -33,6 +33,55 @@ export const JsonObjectSchema = z.record(JsonValueSchema);
export const NonEmptyStringSchema = z.string().trim().min(1);
export const ScopeSelectorSchema = z
.object({
kind: z.enum(['document', 'tag', 'collection', 'favorite', 'artifact']),
id: NonEmptyStringSchema,
name: z.string().optional(),
source: z.enum(['draft', 'focus', 'message']),
})
.strict();
export const ScopeSelectorsSchema = ScopeSelectorSchema.array().max(100);
export const ClientScopeSelectorSchema = ScopeSelectorSchema.omit({
kind: true,
source: true,
}).extend({
kind: z.enum(['document', 'tag', 'collection', 'favorite']),
});
export const RetrievalScopeSchema = z
.object({
mode: z.enum(['workspace', 'required']),
requiredDocIds: z.array(z.string()),
requiredArtifactIds: z.array(z.string()),
preferredSourceIds: z.array(z.string()),
})
.strict();
export const TurnScopeSnapshotSchema = z
.object({
version: z.number().int().positive(),
resolvedAt: z.string(),
selectors: ScopeSelectorsSchema,
requiredDocIds: z.array(z.string()),
requiredArtifactIds: z.array(z.string()),
preferredSourceIds: z.array(z.string()),
retrieval: RetrievalScopeSchema,
})
.strict();
export const SessionFocusSchema = z
.object({
selectors: ScopeSelectorsSchema,
})
.strict();
export type ScopeSelector = z.infer<typeof ScopeSelectorSchema>;
export type TurnScopeSnapshot = z.infer<typeof TurnScopeSnapshotSchema>;
export type SessionFocus = z.infer<typeof SessionFocusSchema>;
export const ToolDefinitionBaseSchema = z
.object({
name: NonEmptyStringSchema,
packages/backend/server/src/plugins/copilot/runtime/copilot-runtime-event-consumer.ts
+15 -7
View File
@@ -53,7 +53,9 @@ export class CopilotRuntimeEventConsumer {
) {
for (const event of events) {
try {
if (event.type === 'usage') {
if (event.type === 'route_selected') {
await this.recordSelection(event, context);
} else if (event.type === 'usage') {
await this.recordUsage(event, context);
} else if (event.type === 'route_failed') {
await this.recordFailure(event, context);
@@ -68,6 +70,18 @@ export class CopilotRuntimeEventConsumer {
}
}
private async recordSelection(
event: Extract<CopilotRuntimeEvent, { type: 'route_selected' }>,
context: CopilotRuntimeEventContext
) {
if (context.workspaceId && event.route.source === 'server') {
await this.models.copilotWorkspaceByokConfig.touchUsed(
context.workspaceId,
event.route.profileId
);
}
}
private async recordUsage(
event: Extract<CopilotRuntimeEvent, { type: 'usage' }>,
context: CopilotRuntimeEventContext
@@ -100,12 +114,6 @@ export class CopilotRuntimeEventConsumer {
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(
packages/backend/server/src/plugins/copilot/runtime/hosts/conversation-host.ts
+128 -3
View File
@@ -2,19 +2,30 @@ import { Injectable } from '@nestjs/common';
import {
CopilotMessageNotFound,
CopilotSelectedSourcesLimitExceeded,
CopilotSessionNotFound,
Mutex,
} from '../../../../base';
import { BackendRuntimeProvider } from '../../../../core/backend-runtime';
import { CompatSubmissionStore } from '../../compat/submission-store';
import { ConversationPolicy } from '../../conversation/policy';
import {
canonicalizeTurnTrace,
promptMessageFromTurn,
type Turn,
turnFromChatMessage,
} from '../../core';
import type { PromptParams } from '../../providers/types';
import { ChatSession, ChatSessionService } from '../../session';
import { ChatQuerySchema } from '../../types';
import {
ClientScopeSelectorSchema,
type ScopeSelector,
ScopeSelectorSchema,
type SessionFocus,
TurnScopeSnapshotSchema,
} from '../contracts/shared';
import { AttachmentAdmissionHost } from './attachment-admission';
export type PreparedConversationTurn = {
messageId?: string;
@@ -35,9 +46,114 @@ export class ConversationHost {
private readonly sessions: ChatSessionService,
private readonly submissions: CompatSubmissionStore,
private readonly mutex: Mutex,
private readonly policy: ConversationPolicy
private readonly policy: ConversationPolicy,
private readonly runtime: BackendRuntimeProvider,
private readonly attachmentAdmission: AttachmentAdmissionHost
) {}
private selectors(
value: unknown,
source: ScopeSelector['source']
): ScopeSelector[] {
if (value === undefined) return [];
return ClientScopeSelectorSchema.array()
.max(100)
.parse(value)
.map(selector => ({ ...selector, source }));
}
private mergeSelectors(...groups: ScopeSelector[][]): ScopeSelector[] {
const merged = new Map<string, ScopeSelector>();
for (const selector of groups.flat()) {
merged.set(`${selector.kind}:${selector.id}`, selector);
}
return [...merged.values()];
}
private async prepareMessageState(
session: ChatSession,
params: Record<string, any>,
attachments: NonNullable<
Parameters<AttachmentAdmissionHost['admitPromptAttachments']>[0]
>
) {
const {
scopeSelectors: rawSelectors,
focusSelectors: rawFocus,
preferredSourceIds: rawPreferred,
...metadata
} = params;
const focus: SessionFocus =
rawFocus === undefined
? session.config.focus
: { selectors: this.selectors(rawFocus, 'focus') };
const admitted = await this.attachmentAdmission.admitPromptAttachments(
attachments,
{
userId: session.config.userId,
workspaceId: session.config.workspaceId,
sessionId: session.config.sessionId,
}
);
const artifacts = await Promise.all(
admitted.map(async source => {
const artifact = await this.runtime.putWorkspaceArtifact(
{
workspaceId: session.config.workspaceId,
mimeType: source.mimeType,
fileName: source.fileName,
libraryOwned: false,
},
Buffer.from(source.data, 'base64')
);
return {
artifactId: artifact.id,
role: 'attachment',
displayName: source.fileName,
metadata: { mimeType: artifact.canonicalMediaType },
};
})
);
const artifactSelectors = artifacts.map(
({ artifactId, displayName }): ScopeSelector => ({
kind: 'artifact',
id: artifactId,
name: displayName,
source: 'message',
})
);
const selectors = this.mergeSelectors(
focus.selectors,
this.selectors(rawSelectors, 'draft'),
artifactSelectors
);
const preferredSourceIds =
rawPreferred === undefined
? []
: ScopeSelectorSchema.shape.id.array().max(100).parse(rawPreferred);
let compiledScope: Awaited<
ReturnType<BackendRuntimeProvider['compileTurnScope']>
>;
try {
compiledScope = await this.runtime.compileTurnScope({
workspaceId: session.config.workspaceId,
userId: session.config.userId,
selectors,
preferredSourceIds,
});
} catch (error) {
if (
error instanceof Error &&
error.message.includes('scope_required_document_limit_exceeded')
) {
throw new CopilotSelectedSourcesLimitExceeded();
}
throw error;
}
const scopeSnapshot = TurnScopeSnapshotSchema.parse(compiledScope);
return { artifacts, focus, metadata, scopeSnapshot };
}
private async loadAcceptedTurn(
session: ChatSession,
sessionId: string,
@@ -180,16 +296,25 @@ export class ConversationHost {
session.revertLatestMessage(true);
}
const prepared = await this.prepareMessageState(
session,
submission.params ?? {},
submission.attachments ?? []
);
const turn = await this.sessions.appendTurn({
sessionId,
userId: session.config.userId,
compatSubmissionId: messageId,
focus: prepared.focus,
artifacts: prepared.artifacts,
turn: {
conversationId: sessionId,
role: 'user',
content: submission.content ?? '',
attachments: submission.attachments ?? [],
metadata: submission.params ?? {},
metadata: prepared.metadata,
scopeSnapshot: prepared.scopeSnapshot,
renderTrace: [],
toolEvents: [],
createdAt: submission.createdAt,
@@ -245,7 +370,7 @@ export class ConversationHost {
return {
...latestTurn.metadata,
content: latestTurn.content,
attachments: latestTurn.attachments,
attachments: promptMessageFromTurn(latestTurn).attachments ?? [],
};
}
packages/backend/server/src/plugins/copilot/runtime/tool-runtime.ts
+94 -45
View File
@@ -1,38 +1,43 @@
/* oxlint-disable import/no-cycle -- Tools can invoke nested prompts and semantic search. */
import { Injectable } from '@nestjs/common';
import { forwardRef, Inject, Injectable } from '@nestjs/common';
import { Config } from '../../../base';
import { DocReader, DocWriter } from '../../../core/doc';
import { PermissionAccess } from '../../../core/permission';
import { Models } from '../../../models';
import { IndexerService } from '../../indexer';
import { CopilotContextService } from '../context/service';
import { DelegatedEditorService } from '../delegated/service';
import {
type CopilotChatOptions,
type CopilotChatTools,
} from '../providers/types';
import { ArtifactRetrievalService } from '../retrieval/artifact';
import { DocumentRetrievalService } from '../retrieval/document';
import {
buildBlobContentGetter,
buildDocCanvasGetter,
buildDocContentGetter,
buildDocCreateHandler,
buildDocKeywordSearchGetter,
buildDocSearchGetter,
buildDocumentSearch,
buildDocUpdateHandler,
buildDocUpdateMetaHandler,
type CopilotTool,
type CopilotToolSet,
createBlobReadTool,
createArtifactReadTool,
createArtifactSearchTool,
createCodeArtifactTool,
createConversationSummaryTool,
createDocCanvasReadTool,
createDocComposeTool,
createDocCreateTool,
createDocKeywordSearchTool,
createDocReadTool,
createDocSemanticSearchTool,
createDocSearchTool,
createDocUpdateMetaTool,
createDocUpdateTool,
createExaCrawlTool,
createExaSearchTool,
createFrontendEditorStateTool,
createFrontendNodesTool,
createFrontendSelectionTool,
createFrontendSnapshotTool,
createSectionEditTool,
} from '../tools';
import { PromptRuntime } from './prompt-runtime';
@@ -47,12 +52,14 @@ export class ToolRuntime {
constructor(
private readonly config: Config,
private readonly ac: PermissionAccess,
private readonly context: CopilotContextService,
private readonly docReader: DocReader,
private readonly docWriter: DocWriter,
private readonly models: Models,
private readonly promptRuntime: PromptRuntime,
private readonly indexerService: IndexerService
@Inject(forwardRef(() => PromptRuntime))
private readonly promptRuntime: Pick<PromptRuntime, 'runText'>,
private readonly retrieval: DocumentRetrievalService,
private readonly artifactRetrieval: ArtifactRetrievalService,
private readonly delegated: DelegatedEditorService
) {}
async getTools(
@@ -80,6 +87,14 @@ export class ToolRuntime {
},
});
const documentScope =
options.retrievalScope?.mode === 'required'
? {
mode: 'selected' as const,
allowedDocIds: options.retrievalScope.requiredDocIds,
}
: undefined;
for (const tool of options.tools) {
const toolDef = resolveProviderSpecificTool?.(tool, model);
if (toolDef) {
@@ -97,13 +112,17 @@ export class ToolRuntime {
}
switch (tool) {
case 'blobRead': {
const docContext = options.session
? await this.context.getBySessionId(options.session)
: null;
const getBlobContent = buildBlobContentGetter(this.ac, docContext);
tools.blob_read = createBlobReadTool(
getBlobContent.bind(null, options)
case 'artifactRead': {
tools.artifact_read = createArtifactReadTool(
this.artifactRetrieval,
options
);
break;
}
case 'artifactSearch': {
tools.artifact_search = createArtifactSearchTool(
this.artifactRetrieval,
options
);
break;
}
@@ -118,40 +137,70 @@ export class ToolRuntime {
);
break;
}
case 'docSemanticSearch': {
const searchDocs = buildDocSearchGetter(
this.ac,
this.context,
options.session,
this.models
);
tools.doc_semantic_search = createDocSemanticSearchTool(
searchDocs.bind(null, options)
);
break;
}
case 'docKeywordSearch': {
if (this.config.indexer.enabled) {
const searchDocs = buildDocKeywordSearchGetter(
this.ac,
this.indexerService,
this.models
);
tools.doc_keyword_search = createDocKeywordSearchTool(
searchDocs.bind(null, options)
);
}
break;
}
case 'docRead': {
const getDoc = buildDocContentGetter(
this.ac,
this.docReader,
this.models
this.models,
documentScope
);
tools.doc_read = createDocReadTool(getDoc.bind(null, options));
break;
}
case 'docCanvasRead': {
const readCanvas = buildDocCanvasGetter(
this.ac,
this.docReader,
this.models,
documentScope
);
tools.doc_canvas_read = createDocCanvasReadTool(
readCanvas.bind(null, options)
);
break;
}
case 'docSearch': {
tools.doc_search = createDocSearchTool(
buildDocumentSearch(this.retrieval, options, documentScope)
);
break;
}
case 'frontendGetEditorState': {
if (this.delegated.getLease(options, 'frontend_get_editor_state')) {
tools.frontend_get_editor_state = createFrontendEditorStateTool(
this.delegated,
options
);
}
break;
}
case 'frontendReadSelection': {
if (this.delegated.getLease(options, 'frontend_read_selection')) {
tools.frontend_read_selection = createFrontendSelectionTool(
this.delegated,
options
);
}
break;
}
case 'frontendReadNodes': {
if (this.delegated.getLease(options, 'frontend_read_nodes')) {
tools.frontend_read_nodes = createFrontendNodesTool(
this.delegated,
options
);
}
break;
}
case 'frontendSnapshotDocument': {
if (this.delegated.getLease(options, 'frontend_snapshot_document')) {
tools.frontend_snapshot_document = createFrontendSnapshotTool(
this.delegated,
options
);
}
break;
}
case 'docCreate': {
const createDoc = buildDocCreateHandler(this.ac, this.docWriter);
tools.doc_create = createDocCreateTool(createDoc.bind(null, options));
packages/backend/server/src/plugins/copilot/runtime/tool/footnotes.ts
+146
View File
@@ -0,0 +1,146 @@
import type { DocSource } from '../../tools/types';
import type { EnrichedToolResultEvent } from './native-runtime-adapter';
export type AttachmentFootnote = {
artifactId: string;
fileName: string;
fileType: string;
};
function pickAttachmentFootnote(value: unknown): AttachmentFootnote | null {
if (!value || typeof value !== 'object') return null;
const record = value as Record<string, unknown>;
if (record.source && typeof record.source === 'object') {
const source = pickAttachmentFootnote(record.source);
if (source) return source;
}
const artifactId =
typeof record.artifactId === 'string'
? record.artifactId
: typeof record.artifact_id === 'string'
? record.artifact_id
: undefined;
const fileName =
typeof record.fileName === 'string'
? record.fileName
: typeof record.name === 'string'
? record.name
: 'Attachment';
const fileType =
typeof record.fileType === 'string'
? record.fileType
: typeof record.mimeType === 'string'
? record.mimeType
: typeof record.mime_type === 'string'
? record.mime_type
: 'application/octet-stream';
return artifactId ? { artifactId, fileName, fileType } : null;
}
export function collectAttachmentFootnotes(
event: EnrichedToolResultEvent
): AttachmentFootnote[] {
if (!['artifact_read', 'artifact_search'].includes(event.name)) return [];
if (!event.output || typeof event.output !== 'object') return [];
const output = event.output as Record<string, unknown>;
if (event.name === 'artifact_search' && Array.isArray(output.hits)) {
return output.hits
.map(pickAttachmentFootnote)
.filter((item): item is AttachmentFootnote => item !== null);
}
const item = pickAttachmentFootnote(output);
return item ? [item] : [];
}
export function formatAttachmentFootnotes(
attachments: AttachmentFootnote[],
options: { includeReferences?: boolean } = {}
) {
const references =
options.includeReferences === false
? ''
: attachments.map((_, index) => `[^attachment-${index + 1}]`).join('');
const definitions = attachments
.map(
(attachment, index) =>
`[^attachment-${index + 1}]: ${JSON.stringify({
type: 'attachment',
artifactId: attachment.artifactId,
fileName: attachment.fileName,
fileType: attachment.fileType,
})}`
)
.join('\n');
return references
? `\n\n${references}\n\n${definitions}`
: `\n\n${definitions}`;
}
function pickDocumentFootnote(value: unknown): DocSource | null {
if (!value || typeof value !== 'object') return null;
const source = value as Record<string, unknown>;
if (source.type !== 'document') return null;
const workspaceId = source.workspace_id ?? source.workspaceId;
const docId = source.doc_id ?? source.docId;
if (typeof workspaceId !== 'string' || typeof docId !== 'string') return null;
const optional = (snake: string, camel: string) => {
const candidate = source[snake] ?? source[camel];
return typeof candidate === 'string' ? candidate : undefined;
};
return {
type: 'document',
workspace_id: workspaceId,
doc_id: docId,
title: typeof source.title === 'string' ? source.title : '',
revision: optional('revision', 'revision'),
visibility: optional('visibility', 'visibility') as
| DocSource['visibility']
| undefined,
block_id: optional('block_id', 'blockId'),
element_id: optional('element_id', 'elementId'),
frame_id: optional('frame_id', 'frameId'),
};
}
export function collectDocumentFootnotes(event: EnrichedToolResultEvent) {
if (
![
'doc_read',
'doc_canvas_read',
'doc_search',
'frontend_read_selection',
'frontend_read_nodes',
'frontend_snapshot_document',
].includes(event.name)
)
return [];
if (!event.output || typeof event.output !== 'object') return [];
const output = event.output as Record<string, unknown>;
const direct = pickDocumentFootnote(output.source);
if (direct) return [direct];
return Array.isArray(output.hits)
? output.hits
.map(hit =>
pickDocumentFootnote((hit as Record<string, unknown>)?.source)
)
.filter((source): source is DocSource => source !== null)
: [];
}
export function formatDocumentFootnotes(documents: DocSource[]) {
const unique = [
...new Map(documents.map(document => [document.doc_id, document])).values(),
];
const references = unique.map((_, index) => `[^doc-${index + 1}]`).join('');
const definitions = unique
.map(
(document, index) =>
`[^doc-${index + 1}]: ${JSON.stringify({
type: 'doc',
docId: document.doc_id,
...(document.title ? { title: document.title } : {}),
})}`
)
.join('\n');
return `\n\n${references}\n\n${definitions}`;
}
packages/backend/server/src/plugins/copilot/runtime/tool/native-adapter.ts
+50 -88
View File
@@ -7,19 +7,21 @@ import {
CitationFootnoteFormatter,
TextStreamParser,
} from '../../providers/utils';
import type { DocSource } from '../../tools/types';
import { projectRuntimeEventToStreamObject } from '../contracts/runtime-event-contract';
import {
type AttachmentFootnote,
collectAttachmentFootnotes,
collectDocumentFootnotes,
formatAttachmentFootnotes,
formatDocumentFootnotes,
} from './footnotes';
import {
type EnrichedToolCallEvent,
type EnrichedToolResultEvent,
NativeRuntimeAdapter,
} from './native-runtime-adapter';
type AttachmentFootnote = {
blobId: string;
fileName: string;
fileType: string;
};
export type NativeProviderAdapterOptions = {
maxSteps?: number;
nodeTextMiddleware?: NodeTextMiddleware[];
@@ -34,79 +36,6 @@ type NativeStreamDispatch = ConstructorParameters<
typeof NativeRuntimeAdapter
>[0];
function pickAttachmentFootnote(value: unknown): AttachmentFootnote | null {
if (!value || typeof value !== 'object') {
return null;
}
const record = value as Record<string, unknown>;
const blobId =
typeof record.blobId === 'string'
? record.blobId
: typeof record.blob_id === 'string'
? record.blob_id
: undefined;
const fileName =
typeof record.fileName === 'string'
? record.fileName
: typeof record.name === 'string'
? record.name
: undefined;
const fileType =
typeof record.fileType === 'string'
? record.fileType
: typeof record.mimeType === 'string'
? record.mimeType
: 'application/octet-stream';
if (!blobId || !fileName) {
return null;
}
return { blobId, fileName, fileType };
}
function collectAttachmentFootnotes(
event: EnrichedToolResultEvent
): AttachmentFootnote[] {
if (event.name === 'blob_read') {
const item = pickAttachmentFootnote(event.output);
return item ? [item] : [];
}
if (event.name === 'doc_semantic_search' && Array.isArray(event.output)) {
return event.output
.map(item => pickAttachmentFootnote(item))
.filter((item): item is AttachmentFootnote => item !== null);
}
return [];
}
function formatAttachmentFootnotes(
attachments: AttachmentFootnote[],
options: { includeReferences?: boolean } = {}
) {
const references =
options.includeReferences === false
? ''
: attachments.map((_, index) => `[^${index + 1}]`).join('');
const definitions = attachments
.map((attachment, index) => {
return `[^${index + 1}]: ${JSON.stringify({
type: 'attachment',
blobId: attachment.blobId,
fileName: attachment.fileName,
fileType: attachment.fileType,
})}`;
})
.join('\n');
return references
? `\n\n${references}\n\n${definitions}`
: `\n\n${definitions}`;
}
export class NativeProviderAdapter {
readonly logger = new Logger(NativeProviderAdapter.name);
readonly #runtime: NativeRuntimeAdapter;
@@ -180,6 +109,9 @@ export class NativeProviderAdapter {
const citationFormatter = this.#enableCitationFootnote
? new CitationFootnoteFormatter()
: null;
const attachmentFootnotes = new Map<string, AttachmentFootnote>();
const documentFootnotes = new Map<string, DocSource>();
let hasAttachmentFootnoteReference = false;
let streamPartId = 0;
const usageState: {
model?: string;
@@ -210,6 +142,9 @@ export class NativeProviderAdapter {
}
case 'text_delta': {
const textEvent = event as unknown as { text: string };
if (textEvent.text.includes('[^attachment-')) {
hasAttachmentFootnoteReference = true;
}
if (textParser) {
yield textParser.parse({
type: 'text-delta',
@@ -247,8 +182,14 @@ export class NativeProviderAdapter {
break;
}
case 'tool_result': {
if (!textParser) break;
const normalized = event as EnrichedToolResultEvent;
collectAttachmentFootnotes(normalized).forEach(attachment => {
attachmentFootnotes.set(attachment.artifactId, attachment);
});
collectDocumentFootnotes(normalized).forEach(document => {
documentFootnotes.set(JSON.stringify(document), document);
});
if (!textParser) break;
yield textParser.parse({
type: 'tool-result',
toolCallId: normalized.call_id,
@@ -280,7 +221,17 @@ export class NativeProviderAdapter {
usageState.usage = doneEvent.usage ?? usageState.usage;
const footnotes = textParser?.end() ?? '';
const citations = citationFormatter?.end() ?? '';
const tails = [citations, footnotes].filter(Boolean).join('\n');
const attachments = attachmentFootnotes.size
? formatAttachmentFootnotes([...attachmentFootnotes.values()], {
includeReferences: !hasAttachmentFootnoteReference,
})
: '';
const documents = documentFootnotes.size
? formatDocumentFootnotes([...documentFootnotes.values()])
: '';
const tails = [citations, attachments, documents, footnotes]
.filter(Boolean)
.join('\n');
if (tails) {
yield `\n${tails}`;
}
@@ -310,7 +261,8 @@ export class NativeProviderAdapter {
? new CitationFootnoteFormatter()
: null;
const fallbackAttachmentFootnotes = new Map<string, AttachmentFootnote>();
let hasFootnoteReference = false;
const fallbackDocumentFootnotes = new Map<string, DocSource>();
let hasAttachmentFootnoteReference = false;
const usageState: {
model?: string;
usage?: Extract<LlmToolLoopStreamEvent, { type: 'usage' }>['usage'];
@@ -340,8 +292,8 @@ export class NativeProviderAdapter {
}
case 'text_delta': {
const textEvent = event as unknown as { text: string };
if (textEvent.text.includes('[^')) {
hasFootnoteReference = true;
if (textEvent.text.includes('[^attachment-')) {
hasAttachmentFootnoteReference = true;
}
yield { type: 'text-delta', textDelta: textEvent.text };
break;
@@ -363,7 +315,10 @@ export class NativeProviderAdapter {
const normalized = event as EnrichedToolResultEvent;
const attachments = collectAttachmentFootnotes(normalized);
attachments.forEach(attachment => {
fallbackAttachmentFootnotes.set(attachment.blobId, attachment);
fallbackAttachmentFootnotes.set(attachment.artifactId, attachment);
});
collectDocumentFootnotes(normalized).forEach(document => {
fallbackDocumentFootnotes.set(JSON.stringify(document), document);
});
const streamObject = projectRuntimeEventToStreamObject(
event as LlmToolLoopStreamEvent
@@ -394,18 +349,25 @@ export class NativeProviderAdapter {
usageState.usage = doneEvent.usage ?? usageState.usage;
const citations = citationFormatter?.end() ?? '';
if (citations) {
hasFootnoteReference = true;
yield { type: 'text-delta', textDelta: `\n${citations}` };
}
if (!citations && fallbackAttachmentFootnotes.size > 0) {
if (fallbackAttachmentFootnotes.size > 0) {
yield {
type: 'text-delta',
textDelta: formatAttachmentFootnotes(
Array.from(fallbackAttachmentFootnotes.values()),
{ includeReferences: !hasFootnoteReference }
{ includeReferences: !hasAttachmentFootnoteReference }
),
};
}
if (fallbackDocumentFootnotes.size > 0) {
yield {
type: 'text-delta',
textDelta: formatDocumentFootnotes([
...fallbackDocumentFootnotes.values(),
]),
};
}
break;
}
case 'provider_selected':

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