diff --git a/docs/configuration/deployments.mdx b/docs/configuration/deployments.mdx index f1fb2cbe..22d4804f 100644 --- a/docs/configuration/deployments.mdx +++ b/docs/configuration/deployments.mdx @@ -50,6 +50,24 @@ vp run deploy:worker:staging Remote migrations and deploy commands affect shared environments. Run them only when you are intentionally doing release work. +## Workspace Search Vector Indexes + +Workspace search uses 1,024-dimension BGE-M3 embeddings with cosine distance. Vectorize indexes and their metadata indexes are persistent infrastructure, so create them before deploying the Worker: + +```bash +wrangler vectorize create thinkex-workspace-search-staging --dimensions 1024 --metric cosine +wrangler vectorize create-metadata-index thinkex-workspace-search-staging --propertyName itemId --type string +wrangler vectorize create-metadata-index thinkex-workspace-search-staging --propertyName parentId --type string +wrangler vectorize create-metadata-index thinkex-workspace-search-staging --propertyName type --type string + +wrangler vectorize create thinkex-workspace-search --dimensions 1024 --metric cosine +wrangler vectorize create-metadata-index thinkex-workspace-search --propertyName itemId --type string +wrangler vectorize create-metadata-index thinkex-workspace-search --propertyName parentId --type string +wrangler vectorize create-metadata-index thinkex-workspace-search --propertyName type --type string +``` + +Create metadata indexes before inserting vectors. Vectors written first are not retroactively indexed for metadata filtering. + ## Secrets Runtime and deploy secrets are synced into GitHub Actions secrets from Infisical. Production and staging workflows require Cloudflare credentials, PostHog variables, and the runtime secrets declared by the Worker configuration. diff --git a/src/features/workspaces/ai/ai-tool-registry.ts b/src/features/workspaces/ai/ai-tool-registry.ts index d466fadb..b108d340 100644 --- a/src/features/workspaces/ai/ai-tool-registry.ts +++ b/src/features/workspaces/ai/ai-tool-registry.ts @@ -53,6 +53,7 @@ export const AI_TOOL_REGISTRY = defineAiToolRegistry({ }), workspace_list_items: readTool({ icon: "file", title: "List workspace", visibility: "hidden" }), workspace_read_items: readTool({ icon: "file", title: "Read workspace" }), + workspace_search: readTool({ icon: "search", title: "Search workspace" }), workspace_rename_item: writeTool({ icon: "edit", title: "Rename item" }), workspace_move_items: writeTool({ icon: "edit", title: "Move items" }), workspace_create_items: writeTool({ icon: "edit", title: "Create items" }), diff --git a/src/features/workspaces/ai/workspace-citations.ts b/src/features/workspaces/ai/workspace-citations.ts index 52c13088..2ad7e818 100644 --- a/src/features/workspaces/ai/workspace-citations.ts +++ b/src/features/workspaces/ai/workspace-citations.ts @@ -1,6 +1,7 @@ import { isToolUIPart, type UIMessage } from "ai"; import { z } from "zod"; +import { getWorkspaceToolResultAdapter } from "#/features/workspaces/ai/workspace-tool-result-adapters"; import { getWorkspaceLocationKey, parseWorkspaceReference, @@ -9,7 +10,6 @@ import { type WorkspaceLocation, workspaceReferenceRecordSchema, } from "#/features/workspaces/locations/workspace-location"; -import { workspaceReadItemsOutputSchema } from "#/features/workspaces/content/workspace-content-contract"; export const WORKSPACE_CITATIONS_DATA_PART_TYPE = "data-workspace-citations"; const MAX_WORKSPACE_CITATIONS_PER_MESSAGE = 50; @@ -107,14 +107,9 @@ export function collectWorkspaceReferenceRecords( const toolName = part.type === "dynamic-tool" ? part.toolName : part.type.split("-").slice(1).join("-"); - if (toolName !== "workspace_read_items") { - continue; - } - - const parsed = workspaceReadItemsOutputSchema.safeParse(part.output); - if (parsed.success) { - records.push(...parsed.data.references); - } + records.push( + ...(getWorkspaceToolResultAdapter(toolName)?.collectReferences(part.output) ?? []), + ); } } diff --git a/src/features/workspaces/ai/workspace-tool-result-adapters.ts b/src/features/workspaces/ai/workspace-tool-result-adapters.ts new file mode 100644 index 00000000..4e2333b3 --- /dev/null +++ b/src/features/workspaces/ai/workspace-tool-result-adapters.ts @@ -0,0 +1,47 @@ +import type { JSONValue } from "ai"; +import type { z } from "zod"; + +import { workspaceReadItemsOutputSchema } from "#/features/workspaces/content/workspace-content-contract"; +import { createWorkspaceReadItemsModelOutput } from "#/features/workspaces/content/workspace-read-references"; +import type { WorkspaceReferenceRecord } from "#/features/workspaces/locations/workspace-location"; +import { workspaceSearchOutputSchema } from "#/features/workspaces/search/workspace-search-contract"; +import { createWorkspaceSearchModelOutput } from "#/features/workspaces/search/workspace-search-references"; + +function defineWorkspaceToolResultAdapter(input: { + collectReferences: (output: z.output) => readonly WorkspaceReferenceRecord[]; + outputSchema: TSchema; + projectOutput: (output: z.output) => unknown; +}) { + return { + collectReferences: (output: unknown) => { + const parsed = input.outputSchema.safeParse(output); + return parsed.success ? input.collectReferences(parsed.data) : []; + }, + projectOutput: (output: unknown) => { + return input.projectOutput(input.outputSchema.parse(output)) as JSONValue; + }, + }; +} + +const workspaceReadItemsResultAdapter = defineWorkspaceToolResultAdapter({ + collectReferences: (output) => output.references, + outputSchema: workspaceReadItemsOutputSchema, + projectOutput: createWorkspaceReadItemsModelOutput, +}); + +const workspaceSearchResultAdapter = defineWorkspaceToolResultAdapter({ + collectReferences: (output) => output.references, + outputSchema: workspaceSearchOutputSchema, + projectOutput: createWorkspaceSearchModelOutput, +}); + +const workspaceToolResultAdapters = { + workspace_read_items: workspaceReadItemsResultAdapter, + workspace_search: workspaceSearchResultAdapter, +} as const; + +export function getWorkspaceToolResultAdapter(name: string) { + return Object.hasOwn(workspaceToolResultAdapters, name) + ? workspaceToolResultAdapters[name as keyof typeof workspaceToolResultAdapters] + : null; +} diff --git a/src/features/workspaces/ai/workspace-tools.ts b/src/features/workspaces/ai/workspace-tools.ts index fc85fe8a..48aae743 100644 --- a/src/features/workspaces/ai/workspace-tools.ts +++ b/src/features/workspaces/ai/workspace-tools.ts @@ -2,13 +2,11 @@ import type { ToolSet } from "ai"; import type { AIThreadContext } from "#/features/workspaces/ai/ai-thread-metadata"; import { defineAIThreadTool } from "#/features/workspaces/ai/ai-thread-tool"; -import { workspaceReadItemsOutputSchema } from "#/features/workspaces/content/workspace-content-contract"; -import { createWorkspaceReadItemsModelOutput } from "#/features/workspaces/content/workspace-read-references"; +import { getWorkspaceToolResultAdapter } from "#/features/workspaces/ai/workspace-tool-result-adapters"; import type { WorkspaceReferenceRecord } from "#/features/workspaces/locations/workspace-location"; import { workspaceToolDefinitions, getWorkspaceToolScopes, - type WorkspaceToolDefinition, } from "#/features/workspaces/operations/workspace-tool-definitions"; import { createWorkspaceAccessContext, @@ -17,14 +15,14 @@ import { } from "#/features/workspaces/operations/workspace-access-context"; type WorkspaceThreadToolConfig = { - definition: WorkspaceToolDefinition; + definition: (typeof workspaceToolDefinitions)[number]; getThreadContext: () => Promise; onWorkspaceReferences?: (records: readonly WorkspaceReferenceRecord[]) => void; }; function createWorkspaceThreadTool(input: WorkspaceThreadToolConfig) { const { definition } = input; - const isWorkspaceRead = definition.name === "workspace_read_items"; + const resultAdapter = getWorkspaceToolResultAdapter(definition.name); return defineAIThreadTool({ description: definition.description, @@ -32,20 +30,18 @@ function createWorkspaceThreadTool(input: WorkspaceThreadToolConfig) { inputExamples: definition.inputExamples, outputSchema: definition.outputSchema, strict: true, - ...(isWorkspaceRead + ...(resultAdapter ? { toModelOutput: ({ output }) => ({ type: "json" as const, - value: createWorkspaceReadItemsModelOutput( - workspaceReadItemsOutputSchema.parse(output), - ), + value: resultAdapter.projectOutput(output), }), } : {}), execute: async (args, context) => { const thread = await requireThreadContext(input.getThreadContext); - const output = await definition.execute( + const output = await definition.executeUnknown( args, createThreadWorkspaceAccessContext( thread, @@ -54,11 +50,9 @@ function createWorkspaceThreadTool(input: WorkspaceThreadToolConfig) { ), ); - if (isWorkspaceRead && input.onWorkspaceReferences) { - const parsed = workspaceReadItemsOutputSchema.safeParse(output); - if (parsed.success) { - input.onWorkspaceReferences(parsed.data.references); - } + const references = resultAdapter?.collectReferences(output) ?? []; + if (references.length > 0 && input.onWorkspaceReferences) { + input.onWorkspaceReferences(references); } return output; diff --git a/src/features/workspaces/documents/document-item-content.ts b/src/features/workspaces/documents/document-item-content.ts index 334bd4ef..2cf3fc88 100644 --- a/src/features/workspaces/documents/document-item-content.ts +++ b/src/features/workspaces/documents/document-item-content.ts @@ -47,17 +47,3 @@ export function persistDocumentItemContentUpdate(input: { WHERE id = ${input.itemId} AND deleted_at IS NULL `; } - -export function touchWorkspaceItemUpdatedAt(input: { - itemId: string; - sql: WorkspaceKernelSql; - updatedAt?: number; -}) { - const updatedAt = input.updatedAt ?? Date.now(); - - input.sql` - UPDATE kernel_items - SET updated_at = ${updatedAt} - WHERE id = ${input.itemId} AND deleted_at IS NULL - `; -} diff --git a/src/features/workspaces/extraction/request-workspace-file-extraction.ts b/src/features/workspaces/extraction/request-workspace-file-extraction.ts index 11b2d8ea..447fd93d 100644 --- a/src/features/workspaces/extraction/request-workspace-file-extraction.ts +++ b/src/features/workspaces/extraction/request-workspace-file-extraction.ts @@ -1,7 +1,7 @@ import { env } from "cloudflare:workers"; -import { sha256Base64UrlText } from "#/lib/binary"; import type { WorkspaceFileExtractionWorkflowParams } from "#/features/workspaces/extraction/types"; +import { getWorkspaceFileExtractionWorkflowId } from "#/features/workspaces/extraction/workspace-file-extraction-workflow-id"; import { getWorkspaceKernel } from "#/features/workspaces/kernel/workspace-kernel-access"; import type { WorkspaceFileAssetKind } from "#/features/workspaces/model/workspace-file"; import { recordOperationalFailure } from "#/integrations/observability/operational-events"; @@ -16,7 +16,10 @@ export async function requestWorkspaceFileExtraction(input: { let workflowId: string | null = null; try { - workflowId = await getWorkspaceFileExtractionWorkflowId(input); + workflowId = await getWorkspaceFileExtractionWorkflowId({ + ...input, + runKey: "initial", + }); const params = { workspaceId: input.workspaceId, itemId: input.itemId, @@ -71,15 +74,3 @@ export async function requestWorkspaceFileExtraction(input: { } } } - -async function getWorkspaceFileExtractionWorkflowId(input: { - workspaceId: string; - itemId: string; - assetKind: WorkspaceFileAssetKind; -}) { - const digest = await sha256Base64UrlText( - `${input.workspaceId}:${input.itemId}:${input.assetKind}-extraction:v2`, - ); - - return `${input.assetKind}-${digest.slice(0, 48)}`; -} diff --git a/src/features/workspaces/extraction/workspace-file-extraction-reconciler.ts b/src/features/workspaces/extraction/workspace-file-extraction-reconciler.ts new file mode 100644 index 00000000..35178188 --- /dev/null +++ b/src/features/workspaces/extraction/workspace-file-extraction-reconciler.ts @@ -0,0 +1,85 @@ +import type { WorkspaceFileExtractionWorkflowParams } from "#/features/workspaces/extraction/types"; +import { getWorkspaceFileExtractionWorkflowId } from "#/features/workspaces/extraction/workspace-file-extraction-workflow-id"; +import { workspaceExtractionStallThresholdMs } from "#/features/workspaces/extraction/workspace-projection-readiness"; +import type { WorkspaceKernelSql } from "#/features/workspaces/kernel/workspace-kernel-schema"; +import { workspaceFileAssetKindSchema } from "#/features/workspaces/model/workspace-file"; + +const extractionHealingVersion = "extraction-healing-v1"; +const failedExtractionCooldownMs = 15 * 60_000; +const workflowBatchSize = 100; + +export async function reconcileWorkspaceFileExtractions(input: { + sql: WorkspaceKernelSql; + workflow: Workflow; + workspaceId: string; +}) { + const now = Date.now(); + const candidates = input.sql<{ + asset_kind: string | number | boolean | null; + id: string; + object_key: string; + projection_updated_at: number; + }>` + SELECT + json_extract(i.metadata_json, '$.assetKind') AS asset_kind, + i.id, + i.object_key, + COALESCE(p.updated_at, i.created_at) AS projection_updated_at + FROM kernel_items i + LEFT JOIN kernel_item_projections p + ON p.item_id = i.id AND p.format = 'pages' + WHERE i.deleted_at IS NULL + AND i.type = 'file' + AND i.object_key IS NOT NULL + AND ( + (p.item_id IS NULL AND i.created_at <= ${now - workspaceExtractionStallThresholdMs}) + OR ( + p.status = 'failed' + AND p.updated_at <= ${now - failedExtractionCooldownMs} + ) + OR ( + p.status = 'processing' + AND p.updated_at <= ${now - workspaceExtractionStallThresholdMs} + ) + OR ( + p.status = 'ready' + AND (p.object_key IS NULL OR p.source_hash IS NULL) + AND p.updated_at <= ${now - workspaceExtractionStallThresholdMs} + ) + ) + ORDER BY i.created_at ASC + `; + for (let offset = 0; offset < candidates.length; offset += workflowBatchSize) { + const workflows = ( + await Promise.all( + candidates.slice(offset, offset + workflowBatchSize).map(async (candidate) => { + const assetKind = workspaceFileAssetKindSchema.safeParse(candidate.asset_kind); + if (!assetKind.success) { + return null; + } + + const runKey = `${extractionHealingVersion}:${candidate.object_key}:${candidate.projection_updated_at}`; + const params = { + actorUserId: null, + assetKind: assetKind.data, + itemId: candidate.id, + requestId: extractionHealingVersion, + workspaceId: input.workspaceId, + } satisfies WorkspaceFileExtractionWorkflowParams; + return { + id: await getWorkspaceFileExtractionWorkflowId({ + assetKind: params.assetKind, + itemId: candidate.id, + runKey, + workspaceId: input.workspaceId, + }), + params, + }; + }), + ) + ).filter((workflow) => workflow !== null); + if (workflows.length > 0) { + await input.workflow.createBatch(workflows); + } + } +} diff --git a/src/features/workspaces/extraction/workspace-file-extraction-workflow-id.ts b/src/features/workspaces/extraction/workspace-file-extraction-workflow-id.ts new file mode 100644 index 00000000..af87089a --- /dev/null +++ b/src/features/workspaces/extraction/workspace-file-extraction-workflow-id.ts @@ -0,0 +1,17 @@ +import type { WorkspaceFileAssetKind } from "#/features/workspaces/model/workspace-file"; +import { sha256Base64UrlText } from "#/lib/binary"; + +export async function getWorkspaceFileExtractionWorkflowId(input: { + workspaceId: string; + itemId: string; + assetKind: WorkspaceFileAssetKind; + runKey: string; +}) { + const identity = + input.runKey === "initial" + ? `${input.workspaceId}:${input.itemId}:${input.assetKind}-extraction:v2` + : `${input.workspaceId}:${input.itemId}:${input.assetKind}:${input.runKey}-extraction:v3`; + const digest = await sha256Base64UrlText(identity); + + return `${input.assetKind}-${digest.slice(0, 48)}`; +} diff --git a/src/features/workspaces/extraction/workspace-page-projection.ts b/src/features/workspaces/extraction/workspace-page-projection.ts index 19a55124..b8e301a4 100644 --- a/src/features/workspaces/extraction/workspace-page-projection.ts +++ b/src/features/workspaces/extraction/workspace-page-projection.ts @@ -175,10 +175,12 @@ export async function readWorkspacePageProjection(input: { // Consume each R2 body before opening the next one; never retain a batch of live responses. for (const pageNumber of selectedPageNumbers) { - const object = await input.bucket.get(getWorkspacePageObjectKey(prefix, pageNumber)); - if (!object) { - throw new Error(`Extracted page ${pageNumber} was not found.`); - } + const object = await getWorkspacePageProjectionObject({ + bucket: input.bucket, + pageMetadataByNumber, + pageNumber, + prefix, + }); totalBytes += object.size; if (totalBytes > maxPageReadBytes) { @@ -186,12 +188,6 @@ export async function readWorkspacePageProjection(input: { throw new WorkspacePageSelectionError("page_selection_too_large"); } - const manifestPage = pageMetadataByNumber?.get(pageNumber); - if (manifestPage && manifestPage.markdownBytes !== object.size) { - await object.body.cancel(); - throw new Error(`Extracted page ${pageNumber} does not match its manifest.`); - } - pages.push({ markdown: await object.text(), pageNumber, @@ -211,6 +207,56 @@ export async function readWorkspacePageProjection(input: { }; } +export async function* iterateWorkspacePageProjection(input: { + bucket: R2Bucket; + expectedSourceHash: string; + manifestObjectKey: string; +}): AsyncGenerator<{ markdown: string; pageNumber: number }> { + const manifest = await readWorkspacePageProjectionManifest(input.bucket, input.manifestObjectKey); + if (manifest.sourceHash !== input.expectedSourceHash) { + throw new Error("Workspace page projection source does not match its published revision."); + } + + const pageMetadataByNumber = manifest.pages + ? new Map(manifest.pages.map((page) => [page.pageNumber, page] as const)) + : null; + const prefix = getManifestPrefix(input.manifestObjectKey); + + for (let pageNumber = 1; pageNumber <= manifest.pageCount; pageNumber += 1) { + const object = await getWorkspacePageProjectionObject({ + bucket: input.bucket, + pageMetadataByNumber, + pageNumber, + prefix, + }); + + yield { + markdown: await object.text(), + pageNumber, + }; + } +} + +async function getWorkspacePageProjectionObject(input: { + bucket: R2Bucket; + pageMetadataByNumber: ReadonlyMap | null; + pageNumber: number; + prefix: string; +}) { + const object = await input.bucket.get(getWorkspacePageObjectKey(input.prefix, input.pageNumber)); + if (!object) { + throw new Error(`Extracted page ${input.pageNumber} was not found.`); + } + + const manifestPage = input.pageMetadataByNumber?.get(input.pageNumber); + if (manifestPage && manifestPage.markdownBytes !== object.size) { + await object.body.cancel(); + throw new Error(`Extracted page ${input.pageNumber} does not match its manifest.`); + } + + return object; +} + function requireManifestPage( pagesByNumber: ReadonlyMap, pageNumber: number, diff --git a/src/features/workspaces/extraction/workspace-projection-readiness.ts b/src/features/workspaces/extraction/workspace-projection-readiness.ts index 0a50de42..2a4bdc7c 100644 --- a/src/features/workspaces/extraction/workspace-projection-readiness.ts +++ b/src/features/workspaces/extraction/workspace-projection-readiness.ts @@ -8,7 +8,7 @@ import type { ReadWorkspaceKernelFileProjectionResult } from "#/features/workspa * a healthy run can legitimately stay `processing` for a little over half an hour. * The threshold sits above that worst case so slow retries are never mislabelled. */ -const extractionStallThresholdMs = 45 * 60_000; +export const workspaceExtractionStallThresholdMs = 45 * 60_000; const minimumRetryAfterSeconds = 15; const maximumRetryAfterSeconds = 120; @@ -58,7 +58,7 @@ export function resolveWorkspaceProjectionReadiness( if (projection.status === "processing") { const elapsedMs = Math.max(0, now - Date.parse(projection.updatedAt)); - if (elapsedMs > extractionStallThresholdMs) { + if (elapsedMs > workspaceExtractionStallThresholdMs) { return { state: "stalled", elapsedSeconds: Math.round(elapsedMs / 1000) }; } diff --git a/src/features/workspaces/kernel/workspace-kernel-access.ts b/src/features/workspaces/kernel/workspace-kernel-access.ts index 2644b61d..04976123 100644 --- a/src/features/workspaces/kernel/workspace-kernel-access.ts +++ b/src/features/workspaces/kernel/workspace-kernel-access.ts @@ -36,6 +36,12 @@ import { import type { ListWorkspaceKernelItemsResult } from "#/features/workspaces/kernel/workspace-kernel-list"; import type { WorkspaceFileAssetKind } from "#/features/workspaces/model/workspace-file"; import type { WorkspaceCommandResult } from "#/features/workspaces/realtime/messages"; +import type { + WorkspaceSearchFailure, + WorkspaceSearchInput, + WorkspaceSearchResult, + WorkspaceSearchStatus, +} from "#/features/workspaces/search/workspace-search-contract"; import { assertCanMutateWorkspace, assertCanReadWorkspace, @@ -114,6 +120,11 @@ export interface WorkspaceKernelClient { actorUserId?: string | null; clientMutationId?: string | null; }): Promise>; + searchWorkspace(input: WorkspaceSearchInput): Promise<{ + failed: WorkspaceSearchFailure[]; + results: WorkspaceSearchResult[]; + status: WorkspaceSearchStatus; + }>; purgeForDeletion(): Promise; } diff --git a/src/features/workspaces/kernel/workspace-kernel-events.ts b/src/features/workspaces/kernel/workspace-kernel-events.ts index 1b1edcd6..bb6f5b7d 100644 --- a/src/features/workspaces/kernel/workspace-kernel-events.ts +++ b/src/features/workspaces/kernel/workspace-kernel-events.ts @@ -9,17 +9,20 @@ export class WorkspaceKernelEventBus { private readonly workspaceId: () => string; private readonly getNextRevision: () => number; private readonly broadcast: (message: WorkspaceRealtimeServerMessage) => void; + private readonly onCommit?: (event: WorkspaceRealtimeEvent) => void; constructor(input: { sql: WorkspaceKernelSql; workspaceId: () => string; getNextRevision: () => number; broadcast: (message: WorkspaceRealtimeServerMessage) => void; + onCommit?: (event: WorkspaceRealtimeEvent) => void; }) { this.sql = input.sql; this.workspaceId = input.workspaceId; this.getNextRevision = input.getNextRevision; this.broadcast = input.broadcast; + this.onCommit = input.onCommit; } commit(input: Omit) { @@ -52,6 +55,7 @@ export class WorkspaceKernelEventBus { ${createdAt} ) `; + this.onCommit?.(event); this.broadcast({ type: "workspace.event", workspaceId: this.workspaceId(), diff --git a/src/features/workspaces/kernel/workspace-kernel-item-commands.ts b/src/features/workspaces/kernel/workspace-kernel-item-commands.ts index 2459803a..e8247ca5 100644 --- a/src/features/workspaces/kernel/workspace-kernel-item-commands.ts +++ b/src/features/workspaces/kernel/workspace-kernel-item-commands.ts @@ -5,7 +5,6 @@ import { workspaceItemTypeSchema } from "#/features/workspaces/contracts"; import { buildWorkspaceItemCreateBootstrap, persistDocumentItemContentUpdate, - touchWorkspaceItemUpdatedAt, } from "#/features/workspaces/documents/document-item-content"; import type { WorkspaceKernelEventBus } from "#/features/workspaces/kernel/workspace-kernel-events"; import { @@ -354,23 +353,16 @@ export class WorkspaceKernelItemCommands { getWorkspaceKernelContentMimeType(type), ); - const now = Date.now(); + const currentItem = this.store.assertActiveItem(input.itemId); + const updatedAt = Math.max(Date.now(), currentItem.updated_at + 1); - if (type === "document") { - persistDocumentItemContentUpdate({ - content: input.content, - itemId: input.itemId, - metadataJson: item.metadata_json, - sql: this.sql, - updatedAt: now, - }); - } else { - touchWorkspaceItemUpdatedAt({ - itemId: input.itemId, - sql: this.sql, - updatedAt: now, - }); - } + persistDocumentItemContentUpdate({ + content: input.content, + itemId: input.itemId, + metadataJson: currentItem.metadata_json, + sql: this.sql, + updatedAt, + }); return this.commitItemEvent({ type: "workspace.item.content.updated", diff --git a/src/features/workspaces/kernel/workspace-kernel.ts b/src/features/workspaces/kernel/workspace-kernel.ts index 474a3717..25835820 100644 --- a/src/features/workspaces/kernel/workspace-kernel.ts +++ b/src/features/workspaces/kernel/workspace-kernel.ts @@ -3,6 +3,7 @@ import { Agent, type Connection, type ConnectionContext } from "agents"; import type { WorkspaceItemSummary } from "#/features/workspaces/contracts"; import { getDocumentSessionFromEnv } from "#/features/workspaces/document-session-access"; +import { reconcileWorkspaceFileExtractions } from "#/features/workspaces/extraction/workspace-file-extraction-reconciler"; import type { ResourcePurgeResult } from "#/features/workspaces/resource-purge-result"; import { WorkspaceKernelEventBus } from "#/features/workspaces/kernel/workspace-kernel-events"; import { WorkspaceKernelFileCommands } from "#/features/workspaces/kernel/workspace-kernel-file-commands"; @@ -63,12 +64,17 @@ import { recordOperationalOutcome, } from "#/integrations/observability/operational-events"; import { deleteR2Prefix } from "#/lib/r2"; +import { WorkspaceSearchProjection } from "#/features/workspaces/search/workspace-search-projection"; +import type { WorkspaceSearchInput } from "#/features/workspaces/search/workspace-search-contract"; const workspaceKernelInlineThresholdBytes = 1_500_000; +const workspaceExtractionHealingThrottleMs = 60_000; +const workspacePurgeMaximumAttempts = 5; export { setWorkspaceKernelUserHeaders }; export class WorkspaceKernel extends Agent { + private lastExtractionHealingRequestAt = 0; private readonly kernelSql: WorkspaceKernelSql = (strings, ...values) => this.sql(strings, ...values); private readonly workspace = new ShellWorkspace({ @@ -82,11 +88,35 @@ export class WorkspaceKernel extends Agent { sql: this.kernelSql, workspaceId: () => this.name, }); + private readonly search = new WorkspaceSearchProjection({ + ai: this.env.AI, + bucket: this.env.WORKSPACE_KERNEL_FILES, + getItems: () => this.store.getPageItems(), + requestRun: () => this.ctx.waitUntil(this.scheduleWorkspaceSearchIndexing()), + sql: this.kernelSql, + vectorize: this.env.WORKSPACE_SEARCH, + workspace: this.workspace, + workspaceId: () => this.name, + }); private readonly events = new WorkspaceKernelEventBus({ sql: this.kernelSql, workspaceId: () => this.name, getNextRevision: () => this.store.getNextRevision(), broadcast: (message) => this.broadcastRealtimeMessage(message), + onCommit: (event) => { + try { + this.search.observe(event); + } catch (error) { + recordOperationalFailure({ + error, + event: "workspace_search_projection", + fields: { + event_type: event.type, + workspace_id: this.name, + }, + }); + } + }, }); private readonly relations = new WorkspaceKernelRelations(this.kernelSql); private readonly itemCommands = new WorkspaceKernelItemCommands({ @@ -105,8 +135,13 @@ export class WorkspaceKernel extends Agent { workspaceId: () => this.name, }); - onStart() { + async onStart() { initializeWorkspaceKernelStorage(this.kernelSql); + this.search.initialize(); + if (this.search.hasPending()) { + await this.scheduleWorkspaceSearchIndexing(); + } + this.requestWorkspaceFileExtractionHealing(); } onConnect(connection: Connection, context: ConnectionContext) { @@ -120,6 +155,7 @@ export class WorkspaceKernel extends Agent { connection.setState({ user, }); + this.requestWorkspaceFileExtractionHealing(); this.broadcastPresenceSnapshot(); } @@ -283,6 +319,22 @@ export class WorkspaceKernel extends Agent { ); } + async searchWorkspace(input: WorkspaceSearchInput) { + this.requestWorkspaceFileExtractionHealing(); + if (this.search.hasPending()) { + this.ctx.waitUntil(this.scheduleWorkspaceSearchIndexing()); + } + return await this.search.search(input); + } + + async processWorkspaceSearchIndex() { + if (await this.search.processBatch()) { + // The current one-shot schedule is removed after this callback returns, + // so its successor must not deduplicate onto the executing row. + await this.scheduleWorkspaceSearchIndexing(false); + } + } + private async runMutation( operation: string, input: { actorUserId?: string | null; clientMutationId?: string | null }, @@ -314,11 +366,22 @@ export class WorkspaceKernel extends Agent { } } - async purgeForDeletion(): Promise { + async purgeForDeletion(input: { attempt?: number } = {}): Promise { const workspaceId = this.name; const documentItemIds = this.store.getAllDocumentItemIds(); let failed = 0; + try { + await this.search.purgeVectors(); + } catch (error) { + failed += 1; + recordOperationalFailure({ + error, + event: "workspace_search_purge", + fields: { workspace_id: workspaceId }, + }); + } + for (const itemId of documentItemIds) { try { await getDocumentSessionFromEnv(this.env, { @@ -338,7 +401,7 @@ export class WorkspaceKernel extends Agent { } } - await Promise.all([ + const r2PurgeResults = await Promise.allSettled([ deleteR2Prefix( this.env.WORKSPACE_KERNEL_FILES, getChatAttachmentWorkspacePrefix(workspaceId), @@ -348,9 +411,68 @@ export class WorkspaceKernel extends Agent { deleteR2Prefix(this.env.WORKSPACE_KERNEL_FILES, `workspace_file_objects/${workspaceId}/`), deleteR2Prefix(this.env.WORKSPACE_KERNEL_FILES, `workspace_file_uploads/${workspaceId}/`), ]); + const r2Failures = r2PurgeResults.filter((result) => result.status === "rejected"); + if (r2Failures.length > 0) { + failed += 1; + for (const result of r2Failures) { + recordOperationalFailure({ + error: result.reason, + event: "workspace_r2_purge", + fields: { workspace_id: workspaceId }, + }); + } + } - await this.ctx.storage.deleteAll(); - return { attempted: documentItemIds.length + 1, failed }; + // Keep the local inventory when a remote purge fails so cleanup can be retried. + if (failed === 0) { + await this.destroy(); + } else { + const attempt = input.attempt ?? 1; + if (attempt < workspacePurgeMaximumAttempts) { + await this.schedule( + attempt * 5, + "purgeForDeletion", + { attempt: attempt + 1 }, + { + // A retry scheduled from the executing one-shot needs its own row. + idempotent: false, + }, + ); + } + } + return { attempted: documentItemIds.length + 2, failed }; + } + + private async scheduleWorkspaceSearchIndexing(idempotent = true) { + await this.schedule(1, "processWorkspaceSearchIndex", undefined, { + idempotent, + retry: { + baseDelayMs: 250, + maxAttempts: 5, + maxDelayMs: 3_000, + }, + }); + } + + private requestWorkspaceFileExtractionHealing() { + const now = Date.now(); + if (now - this.lastExtractionHealingRequestAt < workspaceExtractionHealingThrottleMs) { + return; + } + this.lastExtractionHealingRequestAt = now; + this.ctx.waitUntil( + reconcileWorkspaceFileExtractions({ + sql: this.kernelSql, + workflow: this.env.WORKSPACE_FILE_EXTRACTION_WORKFLOW, + workspaceId: this.name, + }).catch((error) => { + recordOperationalFailure({ + error, + event: "workspace_file_extraction_healing", + fields: { workspace_id: this.name }, + }); + }), + ); } private broadcastPresenceSnapshot() { diff --git a/src/features/workspaces/operations/search-workspace.ts b/src/features/workspaces/operations/search-workspace.ts new file mode 100644 index 00000000..dab6dcae --- /dev/null +++ b/src/features/workspaces/operations/search-workspace.ts @@ -0,0 +1,23 @@ +import type { WorkspaceAccessContext } from "#/features/workspaces/operations/workspace-access-context"; +import { getAuthorizedWorkspaceKernel } from "#/features/workspaces/operations/workspace-operation-context"; +import type { + WorkspaceSearchInput, + WorkspaceSearchOutput, +} from "#/features/workspaces/search/workspace-search-contract"; +import { createWorkspaceSearchReferences } from "#/features/workspaces/search/workspace-search-references"; + +export async function searchWorkspaceOperation( + accessContext: WorkspaceAccessContext, + input: WorkspaceSearchInput, +): Promise { + const kernel = await getAuthorizedWorkspaceKernel({ + access: "read", + context: accessContext, + }); + const output = await kernel.searchWorkspace(input); + + return { + ...output, + references: createWorkspaceSearchReferences(output.results), + }; +} diff --git a/src/features/workspaces/operations/workspace-tool-definitions.ts b/src/features/workspaces/operations/workspace-tool-definitions.ts index 9927caf4..c7fbab64 100644 --- a/src/features/workspaces/operations/workspace-tool-definitions.ts +++ b/src/features/workspaces/operations/workspace-tool-definitions.ts @@ -8,6 +8,7 @@ import { listWorkspaceItemsOperation } from "#/features/workspaces/operations/li import { moveWorkspaceItemsOperation } from "#/features/workspaces/operations/move-items"; import { readWorkspaceItemsOperation } from "#/features/workspaces/operations/read-items"; import { renameWorkspaceItemOperation } from "#/features/workspaces/operations/rename-item"; +import { searchWorkspaceOperation } from "#/features/workspaces/operations/search-workspace"; import { workspaceCreateItemsInputExamples, workspaceCreateItemsInputSchema, @@ -31,6 +32,9 @@ import { workspaceReadItemsInputExamples, workspaceReadItemsInputSchema, workspaceReadItemsOutputSchema, + workspaceSearchInputExamples, + workspaceSearchInputSchema, + workspaceSearchOutputSchema, workspaceRenameItemInputExamples, workspaceRenameItemInputSchema, workspaceRenameItemOutputSchema, @@ -62,7 +66,7 @@ export function getWorkspaceToolScopes( return access === "read" ? ["workspace:read"] : workspaceAccessScopes; } -export type WorkspaceToolDefinition< +type WorkspaceToolDefinition< TName extends string = string, TInputSchema extends z.ZodTypeAny = z.ZodTypeAny, TOutputSchema extends z.ZodTypeAny = z.ZodTypeAny, @@ -146,7 +150,7 @@ export const workspaceToolDefinitions = [ name: "workspace_read_items", access: "read", description: - "Read ThinkEx documents and extracted files by absolute path. Documents return bounded line chunks; files support explicit physical-page selections. Continue either kind with the returned nextCursor.", + "Read ThinkEx documents and extracted files by absolute path. Documents return bounded line chunks; files support explicit physical-page selections. Continue either kind with the returned nextCursor. Uploaded files extract in the background, so a read can come back pending or report that extraction failed; each result carries the guidance for handling it.", inputSchema: workspaceReadItemsInputSchema, inputExamples: workspaceReadItemsInputExamples, outputSchema: workspaceReadItemsOutputSchema, @@ -156,6 +160,24 @@ export const workspaceToolDefinitions = [ return await readWorkspaceItemsOperation(context, { requests }); }, }), + defineWorkspaceTool({ + name: "workspace_search", + access: "read", + description: + "Search current ThinkEx workspace documents and extracted files by meaning and exact text. Optionally scope to an absolute item or folder path and filter content types.", + inputSchema: workspaceSearchInputSchema, + inputExamples: workspaceSearchInputExamples, + outputSchema: workspaceSearchOutputSchema, + summarizeResult: (result) => + summarizeWorkspaceCollectionResult({ + failed: result.failed, + items: result.results, + }), + effects: { destructive: false, idempotent: true }, + execute: async (args, context) => { + return await searchWorkspaceOperation(context, args); + }, + }), defineWorkspaceTool({ name: "workspace_rename_item", access: "write", diff --git a/src/features/workspaces/operations/workspace-tool-schemas.ts b/src/features/workspaces/operations/workspace-tool-schemas.ts index 27991c3c..8118d43b 100644 --- a/src/features/workspaces/operations/workspace-tool-schemas.ts +++ b/src/features/workspaces/operations/workspace-tool-schemas.ts @@ -16,8 +16,17 @@ import { } from "#/features/workspaces/contracts"; import { documentMarkdownEditSchema } from "#/features/workspaces/documents/document-markdown-edits"; import { workspaceFileAssetKindSchema } from "#/features/workspaces/model/workspace-file"; +import { + workspaceSearchInputSchema, + workspaceSearchOutputSchema, +} from "#/features/workspaces/search/workspace-search-contract"; -export { workspaceReadItemsInputSchema, workspaceReadItemsOutputSchema }; +export { + workspaceReadItemsInputSchema, + workspaceReadItemsOutputSchema, + workspaceSearchInputSchema, + workspaceSearchOutputSchema, +}; export const workspaceDocumentMarkdownMathInstruction = "For document Markdown math, use `$...$` for inline math and `$$...$$` on separate lines for block math. Escape literal currency dollar signs as `\\$`."; @@ -216,6 +225,20 @@ export const workspaceReadItemsInputExamples = createInputExamples< }, ); +export const workspaceSearchInputExamples = createInputExamples< + z.input +>( + { + query: "What does the market report say about adoption?", + path: "/Research", + types: ["document", "file"], + limit: 10, + }, + { + query: "photosynthesis experiment results", + }, +); + export const workspaceRenameItemInputExamples = createInputExamples< z.input >({ diff --git a/src/features/workspaces/search/workspace-search-batches.ts b/src/features/workspaces/search/workspace-search-batches.ts new file mode 100644 index 00000000..484828a8 --- /dev/null +++ b/src/features/workspaces/search/workspace-search-batches.ts @@ -0,0 +1,7 @@ +export function batchWorkspaceSearchValues(values: readonly T[], size: number): T[][] { + const batches: T[][] = []; + for (let index = 0; index < values.length; index += size) { + batches.push(values.slice(index, index + size)); + } + return batches; +} diff --git a/src/features/workspaces/search/workspace-search-chunks.ts b/src/features/workspaces/search/workspace-search-chunks.ts new file mode 100644 index 00000000..f4f0b262 --- /dev/null +++ b/src/features/workspaces/search/workspace-search-chunks.ts @@ -0,0 +1,96 @@ +const targetChunkCharacters = 1_800; +const minimumChunkCharacters = 900; +const overlapCharacters = 220; + +interface WorkspaceSearchTextChunk { + content: string; + endLine: number; + startLine: number; +} + +export function* iterateWorkspaceSearchTextChunks( + text: string, +): Generator { + const normalized = text.replace(/\r\n?/g, "\n"); + if (normalized.length === 0) { + yield { content: "", endLine: 0, startLine: 0 }; + return; + } + + let yielded = false; + let start = 0; + let startLine = 1; + let startLineCursor = 0; + + while (start < normalized.length) { + const hardEnd = Math.min(normalized.length, start + targetChunkCharacters); + const end = findChunkEnd(normalized, start, hardEnd); + const rawContent = normalized.slice(start, end); + const leadingWhitespace = rawContent.length - rawContent.trimStart().length; + const trailingWhitespace = rawContent.length - rawContent.trimEnd().length; + const contentStart = start + leadingWhitespace; + const contentEnd = Math.max(contentStart, end - trailingWhitespace); + const content = normalized.slice(contentStart, contentEnd); + + if (content) { + for ( + let newline = normalized.indexOf("\n", startLineCursor); + newline !== -1 && newline < contentStart; + newline = normalized.indexOf("\n", startLineCursor) + ) { + startLine += 1; + startLineCursor = newline + 1; + } + yielded = true; + yield { + content, + endLine: startLine + countLineBreaks(normalized, contentStart, contentEnd), + startLine, + }; + } + + if (end >= normalized.length) { + break; + } + + const overlapStart = Math.max(start + 1, end - overlapCharacters); + const nextParagraph = normalized.indexOf("\n\n", overlapStart); + const nextLine = normalized.indexOf("\n", overlapStart); + const boundary = + nextParagraph !== -1 && nextParagraph < end + ? nextParagraph + 2 + : nextLine !== -1 && nextLine < end + ? nextLine + 1 + : overlapStart; + start = Math.min(boundary, end); + } + + if (!yielded) { + yield { content: "", endLine: 0, startLine: 0 }; + } +} + +function findChunkEnd(text: string, start: number, hardEnd: number) { + if (hardEnd === text.length) { + return hardEnd; + } + + const minimumEnd = start + minimumChunkCharacters; + for (const separator of ["\n\n", "\n", ". "]) { + const boundary = text.lastIndexOf(separator, hardEnd - separator.length); + if (boundary >= minimumEnd) { + return boundary + separator.length; + } + } + + return hardEnd; +} + +function countLineBreaks(text: string, start: number, end: number) { + let count = 0; + for (let index = text.indexOf("\n", start); index !== -1 && index < end;) { + count += 1; + index = text.indexOf("\n", index + 1); + } + return count; +} diff --git a/src/features/workspaces/search/workspace-search-content.ts b/src/features/workspaces/search/workspace-search-content.ts new file mode 100644 index 00000000..33f5867a --- /dev/null +++ b/src/features/workspaces/search/workspace-search-content.ts @@ -0,0 +1,76 @@ +import { serializeTiptapDocumentToMarkdown } from "#/features/workspaces/documents/document-markdown"; +import { parseTiptapDocumentJson } from "#/features/workspaces/documents/tiptap-document"; +import { iterateWorkspacePageProjection } from "#/features/workspaces/extraction/workspace-page-projection"; +import { iterateWorkspaceSearchTextChunks } from "#/features/workspaces/search/workspace-search-chunks"; + +export interface WorkspaceSearchFileSystem { + readFile(path: string): Promise; +} + +interface WorkspaceSearchIndexSourceBase { + itemId: string; + name: string; + sourceVersion: string; +} + +export type WorkspaceSearchIndexSource = WorkspaceSearchIndexSourceBase & + ( + | { shellPath: string; type: "document" } + | { objectKey: string; sourceHash: string; type: "file" } + ); + +interface PreparedWorkspaceSearchChunk { + content: string; + endLine: number | null; + index: number; + pageNumber: number | null; + startLine: number | null; +} + +export async function* iteratePreparedWorkspaceSearchChunks(input: { + bucket: R2Bucket; + source: WorkspaceSearchIndexSource; + workspace: WorkspaceSearchFileSystem; +}): AsyncGenerator { + let index = 0; + + if (input.source.type === "document") { + const checkpoint = await input.workspace.readFile(input.source.shellPath); + if (checkpoint === null) { + throw new Error("Workspace document checkpoint was not found."); + } + const markdown = serializeTiptapDocumentToMarkdown(parseTiptapDocumentJson(checkpoint)); + for (const chunk of iterateWorkspaceSearchTextChunks(markdown)) { + yield { + content: chunk.content, + endLine: chunk.endLine, + index, + pageNumber: null, + startLine: chunk.startLine, + }; + index += 1; + } + return; + } + + for await (const page of iterateWorkspacePageProjection({ + bucket: input.bucket, + expectedSourceHash: input.source.sourceHash, + manifestObjectKey: input.source.objectKey, + })) { + for (const chunk of iterateWorkspaceSearchTextChunks(page.markdown)) { + yield { + content: chunk.content, + endLine: null, + index, + pageNumber: page.pageNumber, + startLine: null, + }; + index += 1; + } + } +} + +export function createWorkspaceSearchEmbeddingText(input: { content: string; title: string }) { + return [`Title: ${input.title}`, "", input.content].join("\n"); +} diff --git a/src/features/workspaces/search/workspace-search-contract.ts b/src/features/workspaces/search/workspace-search-contract.ts new file mode 100644 index 00000000..3cf93cea --- /dev/null +++ b/src/features/workspaces/search/workspace-search-contract.ts @@ -0,0 +1,81 @@ +import { z } from "zod"; + +import { workspaceReferenceRecordSchema } from "#/features/workspaces/locations/workspace-location"; +import { workspaceFileAssetKindSchema } from "#/features/workspaces/model/workspace-file"; + +const workspaceSearchItemTypeSchema = z.enum(["document", "file"]); + +export const workspaceSearchInputSchema = z.object({ + query: z + .string() + .trim() + .min(2) + .max(1_000) + .describe("Text or natural-language question to search for in the workspace."), + path: z + .string() + .min(1) + .optional() + .describe( + "Optional absolute workspace path. A folder searches recursively; an item searches only that item. Defaults to /.", + ), + types: z + .array(workspaceSearchItemTypeSchema) + .min(1) + .max(2) + .overwrite((types) => Array.from(new Set(types))) + .optional() + .describe("Optional content types to include. Defaults to documents and files."), + limit: z + .number() + .int() + .min(1) + .max(25) + .optional() + .describe("Maximum results to return. Defaults to 10 and is capped at 25."), +}); + +const workspaceSearchLocationSchema = z.discriminatedUnion("kind", [ + z.object({ + endLine: z.number().int().nonnegative(), + kind: z.literal("lines"), + startLine: z.number().int().nonnegative(), + }), + z.object({ + kind: z.literal("page"), + pageNumber: z.number().int().positive(), + }), +]); + +const workspaceSearchResultSchema = z.object({ + assetKind: workspaceFileAssetKindSchema.optional(), + excerpt: z.string(), + itemId: z.string().min(1), + location: workspaceSearchLocationSchema, + path: z.string().min(1), + title: z.string().min(1), + type: workspaceSearchItemTypeSchema, +}); + +const workspaceSearchFailureSchema = z.object({ + code: z.enum(["path_not_absolute", "path_not_found"]), + path: z.string(), +}); + +const workspaceSearchStatusSchema = z.enum(["ready", "indexing", "partial"]); + +export const workspaceSearchOutputSchema = z.object({ + failed: z.array(workspaceSearchFailureSchema), + references: z.array(workspaceReferenceRecordSchema), + results: z.array(workspaceSearchResultSchema), + status: workspaceSearchStatusSchema.describe( + "ready when coverage is current, indexing while the search projection catches up, or partial when some files lack usable extracted text or semantic indexing failed.", + ), +}); + +export type WorkspaceSearchInput = z.output; +export type WorkspaceSearchItemType = z.output; +export type WorkspaceSearchResult = z.output; +export type WorkspaceSearchFailure = z.output; +export type WorkspaceSearchStatus = z.output; +export type WorkspaceSearchOutput = z.output; diff --git a/src/features/workspaces/search/workspace-search-embeddings.ts b/src/features/workspaces/search/workspace-search-embeddings.ts new file mode 100644 index 00000000..5790a609 --- /dev/null +++ b/src/features/workspaces/search/workspace-search-embeddings.ts @@ -0,0 +1,35 @@ +import { batchWorkspaceSearchValues } from "#/features/workspaces/search/workspace-search-batches"; + +const workspaceSearchEmbeddingModel = "@cf/baai/bge-m3"; +const embeddingBatchSize = 16; + +export async function embedWorkspaceSearchTexts(ai: Ai, texts: string[]) { + const embeddings: number[][] = []; + for (const batch of batchWorkspaceSearchValues(texts, embeddingBatchSize)) { + const output: unknown = await ai.run(workspaceSearchEmbeddingModel, { + text: batch, + truncate_inputs: true, + }); + embeddings.push(...readEmbeddingData(output)); + } + return embeddings; +} + +function readEmbeddingData(output: unknown): number[][] { + if (!isRecord(output) || !Array.isArray(output.data)) { + throw new Error("Workspace search embedding response is missing vector data."); + } + + const data = output.data.filter( + (vector): vector is number[] => + Array.isArray(vector) && vector.every((value) => typeof value === "number"), + ); + if (data.length !== output.data.length) { + throw new Error("Workspace search embedding response contains invalid vector data."); + } + return data; +} + +function isRecord(value: unknown): value is Record { + return typeof value === "object" && value !== null && !Array.isArray(value); +} diff --git a/src/features/workspaces/search/workspace-search-indexer.ts b/src/features/workspaces/search/workspace-search-indexer.ts new file mode 100644 index 00000000..2b107ab5 --- /dev/null +++ b/src/features/workspaces/search/workspace-search-indexer.ts @@ -0,0 +1,548 @@ +import type { WorkspaceKernelSql } from "#/features/workspaces/kernel/workspace-kernel-schema"; +import { batchWorkspaceSearchValues } from "#/features/workspaces/search/workspace-search-batches"; +import { + createWorkspaceSearchEmbeddingText, + iteratePreparedWorkspaceSearchChunks, + type WorkspaceSearchFileSystem, + type WorkspaceSearchIndexSource, +} from "#/features/workspaces/search/workspace-search-content"; +import { embedWorkspaceSearchTexts } from "#/features/workspaces/search/workspace-search-embeddings"; +import { createWorkspaceSearchVectorMetadata } from "#/features/workspaces/search/workspace-search-scope"; +import { + buildWorkspaceSearchSourceVersion, + workspaceSearchIndexVersion, +} from "#/features/workspaces/search/workspace-search-version"; +import { recordOperationalFailure } from "#/integrations/observability/operational-events"; +import { sha256Base64UrlText } from "#/lib/binary"; + +const searchIndexBatchSize = 2; +const searchChunkProcessingBatchSize = 32; +const maximumIndexAttempts = 5; +const maximumVectorDeleteAttempts = 5; +const vectorDeleteBatchSize = 100; + +interface SearchSourceRow { + id: string; + name: string; + parent_id: string | null; + projection_object_key: string | null; + projection_source_hash: string | null; + projection_updated_at: number | null; + shell_path: string; + type: string; + updated_at: number; +} + +type ScopedWorkspaceSearchIndexSource = WorkspaceSearchIndexSource & { + parentId: string | null; +}; + +interface SearchIndexChunk { + chunkId: string; + content: string; + endLine: number | null; + index: number; + pageNumber: number | null; + startLine: number | null; +} + +export class WorkspaceSearchIndexer { + private readonly ai: Ai; + private readonly bucket: R2Bucket; + private readonly sql: WorkspaceKernelSql; + private readonly vectorize: VectorizeIndex; + private readonly workspace: WorkspaceSearchFileSystem; + private readonly workspaceId: () => string; + + constructor(input: { + ai: Ai; + bucket: R2Bucket; + sql: WorkspaceKernelSql; + vectorize: VectorizeIndex; + workspace: WorkspaceSearchFileSystem; + workspaceId: () => string; + }) { + this.ai = input.ai; + this.bucket = input.bucket; + this.sql = input.sql; + this.vectorize = input.vectorize; + this.workspace = input.workspace; + this.workspaceId = input.workspaceId; + } + + seedPendingItems() { + const now = Date.now(); + this.sql` + INSERT INTO kernel_search_pending (item_id, requested_at) + SELECT i.id, ${now} + FROM kernel_items i + LEFT JOIN kernel_item_projections p + ON p.item_id = i.id + AND p.format = 'pages' + AND p.status = 'ready' + LEFT JOIN kernel_search_items s ON s.item_id = i.id + WHERE i.deleted_at IS NULL + AND ( + i.type = 'document' + OR (i.type = 'file' AND p.source_hash IS NOT NULL) + ) + AND ( + s.item_id IS NULL + OR s.source_version != CASE + WHEN i.type = 'document' + THEN ${workspaceSearchIndexVersion} || ':document:' || i.updated_at + ELSE ${workspaceSearchIndexVersion} || ':file:' || i.updated_at || ':' || p.updated_at || ':' || p.source_hash + END + ) + ON CONFLICT(item_id) DO NOTHING + `; + this.sql` + INSERT INTO kernel_search_pending (item_id, requested_at) + SELECT s.item_id, ${now} + FROM kernel_search_items s + LEFT JOIN kernel_items i ON i.id = s.item_id AND i.deleted_at IS NULL + WHERE i.id IS NULL + ON CONFLICT(item_id) DO NOTHING + `; + } + + markPending(itemId: string) { + this.sql` + INSERT INTO kernel_search_pending (item_id, requested_at) + VALUES (${itemId}, ${Date.now()}) + ON CONFLICT(item_id) DO UPDATE SET + requested_at = excluded.requested_at, + attempts = 0 + `; + } + + hasPending() { + return Boolean( + this.sql<{ item_id: string }>` + SELECT item_id FROM kernel_search_pending + UNION ALL + SELECT vector_id AS item_id + FROM kernel_search_vector_deletes + WHERE attempts < ${maximumVectorDeleteAttempts} + LIMIT 1 + `[0], + ); + } + + async processBatch() { + await this.flushVectorDeletes(); + const pending = this.sql<{ item_id: string }>` + SELECT item_id + FROM kernel_search_pending + ORDER BY requested_at ASC + LIMIT ${searchIndexBatchSize} + `; + + const failures: unknown[] = []; + for (const row of pending) { + try { + await this.indexItem(row.item_id); + } catch (error) { + this.recordIndexFailure(row.item_id, error); + if (!this.markIndexAttemptFailed(row.item_id)) { + failures.push(error); + } + } + } + + await this.flushVectorDeletes(); + if (failures.length > 0) { + throw new AggregateError(failures, "Workspace search indexing failed."); + } + return this.hasPending(); + } + + async purgeVectors() { + const ids = new Set( + this.sql<{ vector_id: string }>` + SELECT chunk_id AS vector_id FROM kernel_search_chunks + UNION + SELECT vector_id FROM kernel_search_vector_deletes + `.map((row) => row.vector_id), + ); + + const failures: unknown[] = []; + for (const batch of batchWorkspaceSearchValues(Array.from(ids), vectorDeleteBatchSize)) { + try { + await this.vectorize.deleteByIds(batch); + this.deletePurgedVectorState(batch); + } catch (error) { + failures.push(error); + } + } + if (failures.length > 0) { + throw new AggregateError(failures, "Workspace search vector purge failed."); + } + } + + private async indexItem(itemId: string) { + const source = this.getIndexSource(itemId); + if (!source) { + this.removeIndexedItem(itemId); + return; + } + + const revisionKey = await sha256Base64UrlText( + `${this.workspaceId()}:${source.itemId}:${source.sourceVersion}`, + ); + if (!this.isCurrentSource(source)) { + return; + } + this.beginIndex(source); + + let chunks: SearchIndexChunk[] = []; + let indexError: unknown = null; + for await (const prepared of iteratePreparedWorkspaceSearchChunks({ + bucket: this.bucket, + source, + workspace: this.workspace, + })) { + const chunk = { + ...prepared, + chunkId: `s${revisionKey}-${prepared.index}`, + }; + this.insertChunk(source, chunk); + chunks.push(chunk); + + if (chunks.length === searchChunkProcessingBatchSize) { + if (indexError === null) { + try { + if (!(await this.indexChunkBatch(source, chunks))) { + return; + } + } catch (error) { + indexError = error; + } + } + chunks = []; + } + } + + if (chunks.length > 0 && indexError === null) { + try { + if (!(await this.indexChunkBatch(source, chunks))) { + return; + } + } catch (error) { + indexError = error; + } + } + if (!this.isCurrentSource(source)) { + return; + } + if (indexError !== null) { + throw indexError; + } + this.markVectorIndexReady(source); + } + + private getIndexSource(itemId: string): ScopedWorkspaceSearchIndexSource | null { + const row = this.sql` + SELECT + i.id, + i.type, + i.name, + i.parent_id, + i.shell_path, + i.updated_at, + p.object_key AS projection_object_key, + p.source_hash AS projection_source_hash, + p.updated_at AS projection_updated_at + FROM kernel_items i + LEFT JOIN kernel_item_projections p + ON p.item_id = i.id + AND p.format = 'pages' + AND p.status = 'ready' + WHERE i.id = ${itemId} + AND i.deleted_at IS NULL + AND i.type IN ('document', 'file') + LIMIT 1 + `[0]; + if (!row) { + return null; + } + + const source = { + itemId: row.id, + name: row.name, + parentId: row.parent_id, + }; + + if (row.type === "document") { + return { + ...source, + shellPath: row.shell_path, + sourceVersion: buildWorkspaceSearchSourceVersion({ + type: "document", + updatedAt: row.updated_at, + }), + type: "document", + }; + } + if ( + !row.projection_object_key || + !row.projection_source_hash || + row.projection_updated_at === null + ) { + return null; + } + + return { + ...source, + objectKey: row.projection_object_key, + sourceHash: row.projection_source_hash, + sourceVersion: buildWorkspaceSearchSourceVersion({ + projectionUpdatedAt: row.projection_updated_at, + sourceHash: row.projection_source_hash, + type: "file", + updatedAt: row.updated_at, + }), + type: "file", + }; + } + + private isCurrentSource(source: ScopedWorkspaceSearchIndexSource) { + const current = this.getIndexSource(source.itemId); + return ( + current?.sourceVersion === source.sourceVersion && + current.name === source.name && + current.parentId === source.parentId + ); + } + + private beginIndex(source: ScopedWorkspaceSearchIndexSource) { + this.discardItemChunks(source.itemId); + this.sql` + INSERT INTO kernel_search_items ( + item_id, + source_version, + vector_status + ) + VALUES ( + ${source.itemId}, + ${source.sourceVersion}, + 'pending' + ) + ON CONFLICT(item_id) DO UPDATE SET + source_version = excluded.source_version, + vector_status = 'pending' + `; + } + + private async indexChunkBatch( + source: ScopedWorkspaceSearchIndexSource, + chunks: readonly SearchIndexChunk[], + ) { + if (!this.isCurrentSource(source)) { + return false; + } + + const embeddings = await embedWorkspaceSearchTexts( + this.ai, + chunks.map((chunk) => + createWorkspaceSearchEmbeddingText({ + content: chunk.content, + title: source.name, + }), + ), + ); + if (embeddings.length !== chunks.length) { + throw new Error("Workspace search embedding response did not match the indexed chunks."); + } + if (!this.isCurrentSource(source)) { + return false; + } + + const vectors = chunks.map((chunk, index): VectorizeVector => { + const values = embeddings[index]; + if (!values) { + throw new Error("Workspace search embedding was missing for an indexed chunk."); + } + return { + id: chunk.chunkId, + metadata: createWorkspaceSearchVectorMetadata({ + itemId: source.itemId, + parentId: source.parentId, + type: source.type, + }), + namespace: this.workspaceId(), + values, + }; + }); + await this.vectorize.upsert(vectors); + if (!this.isCurrentSource(source)) { + for (const vector of vectors) { + this.queueVectorDelete(vector.id); + } + return false; + } + + this.clearVectorDeletes(vectors.map((vector) => vector.id)); + return true; + } + + private insertChunk(source: ScopedWorkspaceSearchIndexSource, chunk: SearchIndexChunk) { + this.sql` + INSERT INTO kernel_search_chunks ( + chunk_id, + item_id, + page_number, + start_line, + end_line + ) + VALUES ( + ${chunk.chunkId}, + ${source.itemId}, + ${chunk.pageNumber}, + ${chunk.startLine}, + ${chunk.endLine} + ) + `; + this.sql` + INSERT INTO kernel_search_fts (rowid, title, content) + SELECT rowid, ${source.name}, ${chunk.content} + FROM kernel_search_chunks + WHERE chunk_id = ${chunk.chunkId} + `; + } + + private markVectorIndexReady(source: ScopedWorkspaceSearchIndexSource) { + this.sql` + UPDATE kernel_search_items + SET vector_status = 'ready' + WHERE item_id = ${source.itemId} AND source_version = ${source.sourceVersion} + `; + this.sql`DELETE FROM kernel_search_pending WHERE item_id = ${source.itemId}`; + } + + private removeIndexedItem(itemId: string) { + this.discardItemChunks(itemId); + this.sql`DELETE FROM kernel_search_items WHERE item_id = ${itemId}`; + this.sql`DELETE FROM kernel_search_pending WHERE item_id = ${itemId}`; + } + + private discardItemChunks(itemId: string) { + this.sql` + INSERT OR IGNORE INTO kernel_search_vector_deletes (vector_id, requested_at) + SELECT chunk_id, ${Date.now()} + FROM kernel_search_chunks + WHERE item_id = ${itemId} + `; + this.deleteLocalChunks(itemId); + } + + private deleteLocalChunks(itemId: string) { + this.sql` + DELETE FROM kernel_search_fts + WHERE rowid IN ( + SELECT rowid FROM kernel_search_chunks WHERE item_id = ${itemId} + ) + `; + this.sql`DELETE FROM kernel_search_chunks WHERE item_id = ${itemId}`; + } + + private deletePurgedVectorState(vectorIds: readonly string[]) { + const vectorIdsJson = JSON.stringify(vectorIds); + this.sql` + DELETE FROM kernel_search_fts + WHERE rowid IN ( + SELECT rowid + FROM kernel_search_chunks + WHERE chunk_id IN (SELECT value FROM json_each(${vectorIdsJson})) + ) + `; + this.sql` + DELETE FROM kernel_search_chunks + WHERE chunk_id IN (SELECT value FROM json_each(${vectorIdsJson})) + `; + this.clearVectorDeletes(vectorIds); + } + + private queueVectorDelete(vectorId: string) { + this.sql` + INSERT INTO kernel_search_vector_deletes (vector_id, requested_at) + VALUES (${vectorId}, ${Date.now()}) + ON CONFLICT(vector_id) DO NOTHING + `; + } + + private clearVectorDeletes(vectorIds: readonly string[]) { + if (vectorIds.length === 0) { + return; + } + this.sql` + DELETE FROM kernel_search_vector_deletes + WHERE vector_id IN ( + SELECT value FROM json_each(${JSON.stringify(vectorIds)}) + ) + `; + } + + private async flushVectorDeletes() { + const ids = this.sql<{ vector_id: string }>` + SELECT vector_id + FROM kernel_search_vector_deletes + WHERE attempts < ${maximumVectorDeleteAttempts} + ORDER BY requested_at ASC + LIMIT ${vectorDeleteBatchSize} + `.map((row) => row.vector_id); + if (ids.length === 0) { + return; + } + + try { + await this.vectorize.deleteByIds(ids); + this.sql` + DELETE FROM kernel_search_vector_deletes + WHERE vector_id IN (SELECT value FROM json_each(${JSON.stringify(ids)})) + `; + } catch (error) { + this.sql` + UPDATE kernel_search_vector_deletes + SET attempts = attempts + 1 + WHERE vector_id IN (SELECT value FROM json_each(${JSON.stringify(ids)})) + `; + recordOperationalFailure({ + error, + event: "workspace_search_vector_cleanup", + fields: { workspace_id: this.workspaceId() }, + }); + } + } + + private recordIndexFailure(itemId: string, error: unknown) { + recordOperationalFailure({ + error, + event: "workspace_search_indexing", + fields: { + item_id: itemId, + workspace_id: this.workspaceId(), + }, + }); + } + + private markIndexAttemptFailed(itemId: string) { + const attempt = this.sql<{ attempts: number }>` + UPDATE kernel_search_pending + SET attempts = attempts + 1 + WHERE item_id = ${itemId} + RETURNING attempts + `[0]; + if (!attempt || attempt.attempts < maximumIndexAttempts) { + return false; + } + + this.sql` + UPDATE kernel_search_items + SET vector_status = 'failed' + WHERE item_id = ${itemId} + `; + this.sql`DELETE FROM kernel_search_pending WHERE item_id = ${itemId}`; + return true; + } +} diff --git a/src/features/workspaces/search/workspace-search-projection.ts b/src/features/workspaces/search/workspace-search-projection.ts new file mode 100644 index 00000000..d49c3716 --- /dev/null +++ b/src/features/workspaces/search/workspace-search-projection.ts @@ -0,0 +1,103 @@ +import type { WorkspaceItemSummary } from "#/features/workspaces/contracts"; +import type { WorkspaceKernelSql } from "#/features/workspaces/kernel/workspace-kernel-schema"; +import type { WorkspaceRealtimeEvent } from "#/features/workspaces/realtime/messages"; +import type { WorkspaceSearchFileSystem } from "#/features/workspaces/search/workspace-search-content"; +import type { WorkspaceSearchInput } from "#/features/workspaces/search/workspace-search-contract"; +import { WorkspaceSearchIndexer } from "#/features/workspaces/search/workspace-search-indexer"; +import { WorkspaceSearchQuery } from "#/features/workspaces/search/workspace-search-query"; +import { initializeWorkspaceSearchStorage } from "#/features/workspaces/search/workspace-search-schema"; + +export class WorkspaceSearchProjection { + private readonly indexer: WorkspaceSearchIndexer; + private readonly query: WorkspaceSearchQuery; + private readonly requestRun: () => void; + private readonly sql: WorkspaceKernelSql; + + constructor(input: { + ai: Ai; + bucket: R2Bucket; + getItems: () => WorkspaceItemSummary[]; + requestRun: () => void; + sql: WorkspaceKernelSql; + vectorize: VectorizeIndex; + workspace: WorkspaceSearchFileSystem; + workspaceId: () => string; + }) { + this.requestRun = input.requestRun; + this.sql = input.sql; + this.indexer = new WorkspaceSearchIndexer(input); + this.query = new WorkspaceSearchQuery(input); + } + + initialize() { + initializeWorkspaceSearchStorage(this.sql); + this.indexer.seedPendingItems(); + } + + observe(event: WorkspaceRealtimeEvent) { + switch (event.type) { + case "workspace.item.created": + if (event.payload.item.type === "document") { + this.indexer.markPending(event.payload.item.id); + } + break; + case "workspace.item.content.updated": + this.indexer.markPending(event.payload.item.id); + break; + case "workspace.item.renamed": + if (event.payload.item.type === "document" || event.payload.item.type === "file") { + this.indexer.markPending(event.payload.item.id); + } + break; + case "workspace.item.moved": + // Folder scope follows the live tree, so only moved content needs new metadata. + if (event.payload.item.type === "document" || event.payload.item.type === "file") { + this.indexer.markPending(event.payload.item.id); + } + break; + case "workspace.items.moved": + for (const item of event.payload.items) { + if (item.type === "document" || item.type === "file") { + this.indexer.markPending(item.id); + } + } + break; + case "workspace.item.projection.updated": + for (const fact of event.payload.itemFacts) { + this.indexer.markPending(fact.itemId); + } + break; + case "workspace.item.deleted": + for (const itemId of event.payload.deletedItemIds) { + this.indexer.markPending(itemId); + } + break; + case "workspace.item.color.updated": + case "workspace.relations.updated": + return; + default: + event satisfies never; + return; + } + + if (this.indexer.hasPending()) { + this.requestRun(); + } + } + + hasPending() { + return this.indexer.hasPending(); + } + + async processBatch() { + return await this.indexer.processBatch(); + } + + async search(input: WorkspaceSearchInput) { + return await this.query.search(input); + } + + async purgeVectors() { + await this.indexer.purgeVectors(); + } +} diff --git a/src/features/workspaces/search/workspace-search-query.ts b/src/features/workspaces/search/workspace-search-query.ts new file mode 100644 index 00000000..301e7a3f --- /dev/null +++ b/src/features/workspaces/search/workspace-search-query.ts @@ -0,0 +1,429 @@ +import type { WorkspaceItemSummary } from "#/features/workspaces/contracts"; +import { buildWorkspaceKernelItemPathIndex } from "#/features/workspaces/kernel/workspace-kernel-paths"; +import type { WorkspaceKernelSql } from "#/features/workspaces/kernel/workspace-kernel-schema"; +import { resolveWorkspaceFileTypeFromItem } from "#/features/workspaces/model/workspace-file"; +import type { + WorkspaceSearchFailure, + WorkspaceSearchInput, + WorkspaceSearchItemType, + WorkspaceSearchResult, + WorkspaceSearchStatus, +} from "#/features/workspaces/search/workspace-search-contract"; +import { batchWorkspaceSearchValues } from "#/features/workspaces/search/workspace-search-batches"; +import { embedWorkspaceSearchTexts } from "#/features/workspaces/search/workspace-search-embeddings"; +import { + fuseWorkspaceSearchRanks, + type WorkspaceSearchRankCandidate, +} from "#/features/workspaces/search/workspace-search-ranking"; +import { + resolveWorkspaceSearchScope, + type WorkspaceSearchLocalScope, + type WorkspaceSearchScope, + type WorkspaceSearchVectorFilter, +} from "#/features/workspaces/search/workspace-search-scope"; +import { recordOperationalFailure } from "#/integrations/observability/operational-events"; + +const semanticQueryConcurrency = 4; +const maximumSemanticQueries = 16; + +interface SearchChunkRow { + chunk_id: string; + content: string; + end_line: number | null; + item_id: string; + page_number: number | null; + start_line: number | null; +} + +interface SearchCandidate extends WorkspaceSearchRankCandidate { + content: string; + endLine: number | null; + pageNumber: number | null; + startLine: number | null; +} + +export class WorkspaceSearchQuery { + private readonly ai: Ai; + private readonly getItems: () => WorkspaceItemSummary[]; + private readonly sql: WorkspaceKernelSql; + private readonly vectorize: VectorizeIndex; + private readonly workspaceId: () => string; + + constructor(input: { + ai: Ai; + getItems: () => WorkspaceItemSummary[]; + sql: WorkspaceKernelSql; + vectorize: VectorizeIndex; + workspaceId: () => string; + }) { + this.ai = input.ai; + this.getItems = input.getItems; + this.sql = input.sql; + this.vectorize = input.vectorize; + this.workspaceId = input.workspaceId; + } + + async search(input: WorkspaceSearchInput): Promise<{ + failed: WorkspaceSearchFailure[]; + results: WorkspaceSearchResult[]; + status: WorkspaceSearchStatus; + }> { + const status = this.getStatus(); + const limit = input.limit ?? 10; + const types = input.types ?? ["document", "file"]; + const items = this.getItems(); + const resolvedScope = resolveWorkspaceSearchScope({ + items, + path: input.path ?? "/", + types, + }); + if (resolvedScope.status === "failed") { + return { failed: [resolvedScope.failure], results: [], status }; + } + const { scope } = resolvedScope; + const candidateLimit = Math.min(100, Math.max(50, limit * 6)); + const keyword = this.searchKeyword({ + candidateLimit, + query: input.query, + scope: scope.local, + types, + }); + const semantic = await this.searchSemanticWithFallback({ + candidateLimit, + query: input.query, + scope, + types, + }); + const itemsById = new Map(items.map((item) => [item.id, item])); + const paths = buildWorkspaceKernelItemPathIndex(items); + const ranked = fuseWorkspaceSearchRanks({ + keyword, + limit, + semantic: semantic.candidates, + }); + + return { + failed: [], + status: semantic.degraded ? "partial" : status, + results: ranked.flatMap((candidate) => { + const path = paths.get(candidate.itemId); + const item = itemsById.get(candidate.itemId); + if (!path || !item) { + return []; + } + const result = mapSearchResult(candidate, path, item, input.query); + return result ? [result] : []; + }), + }; + } + + private getStatus(): WorkspaceSearchStatus { + const unavailableFile = this.sql<{ item_id: string }>` + SELECT i.id AS item_id + FROM kernel_items i + LEFT JOIN kernel_item_projections p + ON p.item_id = i.id + AND p.format = 'pages' + AND p.status = 'ready' + AND p.object_key IS NOT NULL + AND p.source_hash IS NOT NULL + WHERE i.deleted_at IS NULL + AND i.type = 'file' + AND p.item_id IS NULL + LIMIT 1 + `[0]; + if (unavailableFile) { + return "partial"; + } + + const pending = this.sql<{ item_id: string }>` + SELECT item_id FROM kernel_search_pending LIMIT 1 + `[0]; + if (pending) { + return "indexing"; + } + + const failed = this.sql<{ item_id: string }>` + SELECT item_id + FROM kernel_search_items + WHERE vector_status = 'failed' + LIMIT 1 + `[0]; + return failed ? "partial" : "ready"; + } + + private async searchSemanticWithFallback(input: { + candidateLimit: number; + query: string; + scope: WorkspaceSearchScope; + types: WorkspaceSearchItemType[]; + }) { + try { + return await this.searchSemantic(input); + } catch (error) { + recordOperationalFailure({ + error, + event: "workspace_search_semantic", + fields: { workspace_id: this.workspaceId() }, + }); + return { candidates: [], degraded: true }; + } + } + + private searchKeyword(input: { + candidateLimit: number; + query: string; + scope: WorkspaceSearchLocalScope; + types: WorkspaceSearchItemType[]; + }): SearchCandidate[] { + const match = createFtsMatchExpression(input.query); + if (!match) { + return []; + } + + const localScope = serializeWorkspaceSearchLocalScope(input.scope); + const typesJson = JSON.stringify(input.types); + const rows = this.sql` + SELECT + c.chunk_id, + c.item_id, + kernel_search_fts.content, + c.page_number, + c.start_line, + c.end_line + FROM kernel_search_fts + JOIN kernel_search_chunks c ON c.rowid = kernel_search_fts.rowid + JOIN kernel_items i ON i.id = c.item_id AND i.deleted_at IS NULL + JOIN kernel_search_items s ON s.item_id = c.item_id + WHERE kernel_search_fts MATCH ${match} + AND s.vector_status IN ('ready', 'failed') + AND NOT EXISTS ( + SELECT 1 + FROM kernel_search_pending pending + WHERE pending.item_id = i.id + ) + AND ( + ${localScope.kind} = 'workspace' + OR ( + ${localScope.kind} = 'item' + AND c.item_id IN (SELECT value FROM json_each(${localScope.idsJson})) + ) + OR ( + ${localScope.kind} = 'folder' + AND i.parent_id IN (SELECT value FROM json_each(${localScope.idsJson})) + ) + ) + AND i.type IN (SELECT value FROM json_each(${typesJson})) + ORDER BY bm25(kernel_search_fts, 8.0, 1.0) ASC + LIMIT ${input.candidateLimit} + `; + + return rows.map(mapSearchCandidate); + } + + private async searchSemantic(input: { + candidateLimit: number; + query: string; + scope: WorkspaceSearchScope; + types: WorkspaceSearchItemType[]; + }) { + if (input.scope.vectorFilters.length === 0) { + return { candidates: [], degraded: false }; + } + const [embedding] = await embedWorkspaceSearchTexts(this.ai, [input.query]); + if (!embedding) { + throw new Error("Workspace search embedding response did not include the query vector."); + } + + const semanticMatches = await this.querySemanticMatches({ + candidateLimit: input.candidateLimit, + embedding, + filters: input.scope.vectorFilters, + }); + const { matches } = semanticMatches; + const vectorIds = matches.map((match) => match.id); + if (vectorIds.length === 0) { + return { candidates: [], degraded: semanticMatches.degraded }; + } + + const rows = this.loadSemanticCandidates({ + scope: input.scope.local, + types: input.types, + vectorIds, + }); + const byVectorId = new Map(rows.map((row) => [row.chunk_id, mapSearchCandidate(row)])); + + return { + candidates: vectorIds.flatMap((vectorId) => { + const candidate = byVectorId.get(vectorId); + return candidate ? [candidate] : []; + }), + degraded: semanticMatches.degraded, + }; + } + + private async querySemanticMatches(input: { + candidateLimit: number; + embedding: number[]; + filters: WorkspaceSearchVectorFilter[]; + }) { + let bestMatches = new Map(); + const filters = input.filters.slice(0, maximumSemanticQueries); + + for (const filterBatch of batchWorkspaceSearchValues(filters, semanticQueryConcurrency)) { + const results = await Promise.all( + filterBatch.map((filter) => + this.vectorize.query(input.embedding, { + ...(filter ? { filter } : {}), + namespace: this.workspaceId(), + returnMetadata: "none", + topK: input.candidateLimit, + }), + ), + ); + + for (const match of results.flatMap((result) => result.matches)) { + const existing = bestMatches.get(match.id); + if (!existing || match.score > existing.score) { + bestMatches.set(match.id, match); + } + } + bestMatches = new Map( + Array.from(bestMatches.values()) + .sort(compareSemanticMatches) + .slice(0, input.candidateLimit) + .map((match) => [match.id, match]), + ); + } + + return { + degraded: filters.length < input.filters.length, + matches: Array.from(bestMatches.values()).sort(compareSemanticMatches), + }; + } + + private loadSemanticCandidates(input: { + scope: WorkspaceSearchLocalScope; + types: WorkspaceSearchItemType[]; + vectorIds: string[]; + }) { + const localScope = serializeWorkspaceSearchLocalScope(input.scope); + const typesJson = JSON.stringify(input.types); + + return this.sql` + SELECT + c.chunk_id, + c.item_id, + kernel_search_fts.content, + c.page_number, + c.start_line, + c.end_line + FROM kernel_search_chunks c + JOIN kernel_search_fts ON kernel_search_fts.rowid = c.rowid + JOIN kernel_items i ON i.id = c.item_id AND i.deleted_at IS NULL + JOIN kernel_search_items s ON s.item_id = c.item_id AND s.vector_status = 'ready' + WHERE c.chunk_id IN (SELECT value FROM json_each(${JSON.stringify(input.vectorIds)})) + AND NOT EXISTS ( + SELECT 1 + FROM kernel_search_pending pending + WHERE pending.item_id = i.id + ) + AND ( + ${localScope.kind} = 'workspace' + OR ( + ${localScope.kind} = 'item' + AND c.item_id IN (SELECT value FROM json_each(${localScope.idsJson})) + ) + OR ( + ${localScope.kind} = 'folder' + AND i.parent_id IN (SELECT value FROM json_each(${localScope.idsJson})) + ) + ) + AND i.type IN (SELECT value FROM json_each(${typesJson})) + `; + } +} + +function serializeWorkspaceSearchLocalScope(scope: WorkspaceSearchLocalScope) { + return { + idsJson: JSON.stringify( + scope.kind === "workspace" ? [] : scope.kind === "item" ? [scope.itemId] : scope.folderIds, + ), + kind: scope.kind, + }; +} + +function compareSemanticMatches(left: VectorizeMatch, right: VectorizeMatch) { + return right.score - left.score || left.id.localeCompare(right.id); +} + +function mapSearchCandidate(row: SearchChunkRow): SearchCandidate { + return { + chunkId: row.chunk_id, + content: row.content, + endLine: row.end_line, + itemId: row.item_id, + pageNumber: row.page_number, + startLine: row.start_line, + }; +} + +function mapSearchResult( + candidate: SearchCandidate, + path: string, + item: WorkspaceItemSummary, + query: string, +): WorkspaceSearchResult | null { + if (item.type !== "document" && item.type !== "file") { + return null; + } + const fileType = item.type === "file" ? resolveWorkspaceFileTypeFromItem(item) : null; + + return { + ...(fileType ? { assetKind: fileType.assetKind } : {}), + excerpt: createSearchExcerpt(candidate.content, query), + itemId: candidate.itemId, + location: + candidate.pageNumber === null + ? { + endLine: candidate.endLine ?? 0, + kind: "lines", + startLine: candidate.startLine ?? 0, + } + : { + kind: "page", + pageNumber: candidate.pageNumber, + }, + path, + title: item.name, + type: item.type, + }; +} + +function createFtsMatchExpression(query: string) { + const tokens = getSearchQueryTokens(query); + return tokens.length > 0 + ? tokens.map((token) => `"${token.replaceAll('"', '""')}"*`).join(" OR ") + : null; +} + +function createSearchExcerpt(content: string, query: string) { + const maximumLength = 700; + if (content.length <= maximumLength) { + return content; + } + + const normalizedContent = content.toLocaleLowerCase(); + const matchIndex = getSearchQueryTokens(query).reduce((earliest, token) => { + const index = normalizedContent.indexOf(token.toLocaleLowerCase()); + return index === -1 || (earliest !== -1 && earliest <= index) ? earliest : index; + }, -1); + const start = Math.max(0, (matchIndex === -1 ? 0 : matchIndex) - Math.floor(maximumLength / 3)); + const end = Math.min(content.length, start + maximumLength); + + return `${start > 0 ? "…" : ""}${content.slice(start, end).trim()}${end < content.length ? "…" : ""}`; +} + +function getSearchQueryTokens(query: string) { + return query.match(/[\p{L}\p{N}_]+/gu)?.slice(0, 20) ?? []; +} diff --git a/src/features/workspaces/search/workspace-search-ranking.ts b/src/features/workspaces/search/workspace-search-ranking.ts new file mode 100644 index 00000000..8fc2482f --- /dev/null +++ b/src/features/workspaces/search/workspace-search-ranking.ts @@ -0,0 +1,53 @@ +export interface WorkspaceSearchRankCandidate { + chunkId: string; + itemId: string; +} + +const reciprocalRankConstant = 60; +const maximumChunksPerItem = 2; + +export function fuseWorkspaceSearchRanks(input: { + keyword: readonly T[]; + limit: number; + semantic: readonly T[]; +}): T[] { + const ranked = new Map(); + + addRankedList(ranked, input.keyword, 1.15); + addRankedList(ranked, input.semantic, 1); + + const itemCounts = new Map(); + const results: T[] = []; + for (const { candidate } of Array.from(ranked.values()).sort( + (left, right) => right.score - left.score, + )) { + const itemCount = itemCounts.get(candidate.itemId) ?? 0; + if (itemCount >= maximumChunksPerItem) { + continue; + } + + results.push(candidate); + itemCounts.set(candidate.itemId, itemCount + 1); + if (results.length === input.limit) { + break; + } + } + + return results; +} + +function addRankedList( + ranked: Map, + candidates: readonly T[], + weight: number, +) { + for (const [index, candidate] of candidates.entries()) { + const score = weight / (reciprocalRankConstant + index + 1); + const existing = ranked.get(candidate.chunkId); + if (existing) { + existing.score += score; + } else { + ranked.set(candidate.chunkId, { candidate, score }); + } + } +} diff --git a/src/features/workspaces/search/workspace-search-references.test.ts b/src/features/workspaces/search/workspace-search-references.test.ts new file mode 100644 index 00000000..d233bd48 --- /dev/null +++ b/src/features/workspaces/search/workspace-search-references.test.ts @@ -0,0 +1,55 @@ +import { describe, expect, it } from "vitest"; + +import type { WorkspaceSearchResult } from "#/features/workspaces/search/workspace-search-contract"; +import { + createWorkspaceSearchModelOutput, + createWorkspaceSearchReferences, +} from "#/features/workspaces/search/workspace-search-references"; + +describe("workspace search references", () => { + it("maps document hits to items and PDF hits to physical pages", () => { + const results = searchResults(); + const references = createWorkspaceSearchReferences(results); + + expect(references.map(({ location }) => location)).toEqual([ + { itemId: "document-1", kind: "item", version: 1 }, + { itemId: "file-1", kind: "pdf-page", pageNumber: 12, version: 1 }, + ]); + }); + + it("projects short references without exposing durable item IDs", () => { + const results = searchResults(); + const modelOutput = createWorkspaceSearchModelOutput({ + failed: [], + references: createWorkspaceSearchReferences(results), + results, + status: "ready", + }); + + expect(modelOutput.results.every((result) => "reference" in result)).toBe(true); + expect(JSON.stringify(modelOutput)).not.toContain("document-1"); + expect(JSON.stringify(modelOutput)).not.toContain("file-1"); + }); +}); + +function searchResults(): WorkspaceSearchResult[] { + return [ + { + excerpt: "Document hit", + itemId: "document-1", + location: { endLine: 8, kind: "lines", startLine: 4 }, + path: "/Notes", + title: "Notes", + type: "document", + }, + { + assetKind: "pdf", + excerpt: "PDF hit", + itemId: "file-1", + location: { kind: "page", pageNumber: 12 }, + path: "/Report.pdf", + title: "Report", + type: "file", + }, + ]; +} diff --git a/src/features/workspaces/search/workspace-search-references.ts b/src/features/workspaces/search/workspace-search-references.ts new file mode 100644 index 00000000..0a789a6e --- /dev/null +++ b/src/features/workspaces/search/workspace-search-references.ts @@ -0,0 +1,52 @@ +import { + createWorkspaceReferenceRecords, + getWorkspaceLocationKey, + type WorkspaceLocation, +} from "#/features/workspaces/locations/workspace-location"; +import type { + WorkspaceSearchOutput, + WorkspaceSearchResult, +} from "#/features/workspaces/search/workspace-search-contract"; + +export function createWorkspaceSearchReferences(results: readonly WorkspaceSearchResult[]) { + return createWorkspaceReferenceRecords(results.map(getSearchResultLocation)); +} + +export function createWorkspaceSearchModelOutput(output: WorkspaceSearchOutput) { + const refsByLocation = new Map( + output.references.map((record) => [getWorkspaceLocationKey(record.location), record.ref]), + ); + + return { + failed: output.failed, + status: output.status, + results: output.results.map((result) => { + const { itemId: _itemId, ...modelResult } = result; + const reference = refsByLocation.get( + getWorkspaceLocationKey(getSearchResultLocation(result)), + ); + + return { + ...modelResult, + ...(reference ? { reference } : {}), + }; + }), + }; +} + +function getSearchResultLocation(result: WorkspaceSearchResult): WorkspaceLocation { + if (result.type === "file" && result.assetKind === "pdf" && result.location.kind === "page") { + return { + itemId: result.itemId, + kind: "pdf-page", + pageNumber: result.location.pageNumber, + version: 1, + }; + } + + return { + itemId: result.itemId, + kind: "item", + version: 1, + }; +} diff --git a/src/features/workspaces/search/workspace-search-schema.ts b/src/features/workspaces/search/workspace-search-schema.ts new file mode 100644 index 00000000..2bdd6325 --- /dev/null +++ b/src/features/workspaces/search/workspace-search-schema.ts @@ -0,0 +1,86 @@ +import type { WorkspaceKernelSql } from "#/features/workspaces/kernel/workspace-kernel-schema"; +import { workspaceSearchIndexVersion } from "#/features/workspaces/search/workspace-search-version"; + +const workspaceSearchVersionKey = "workspace_search_version"; + +export function initializeWorkspaceSearchStorage(sql: WorkspaceKernelSql) { + sql` + CREATE TABLE IF NOT EXISTS kernel_search_vector_deletes ( + vector_id TEXT PRIMARY KEY, + requested_at INTEGER NOT NULL, + attempts INTEGER NOT NULL DEFAULT 0 + ) + `; + sql`CREATE INDEX IF NOT EXISTS kernel_search_vector_deletes_pending_idx + ON kernel_search_vector_deletes (requested_at)`; + + const storedVersion = sql<{ value: string }>` + SELECT value + FROM kernel_meta + WHERE key = ${workspaceSearchVersionKey} + LIMIT 1 + `[0]?.value; + if (storedVersion !== workspaceSearchIndexVersion) { + // Search is derived, so schema changes reset only this projection and + // queue its old vectors for deletion without touching workspace data. + const hasChunks = sql<{ name: string }>` + SELECT name + FROM sqlite_master + WHERE type = 'table' AND name = 'kernel_search_chunks' + LIMIT 1 + `[0]; + if (hasChunks) { + sql` + INSERT OR IGNORE INTO kernel_search_vector_deletes (vector_id, requested_at) + SELECT chunk_id, ${Date.now()} FROM kernel_search_chunks + `; + } + sql`DROP TABLE IF EXISTS kernel_search_fts`; + sql`DROP TABLE IF EXISTS kernel_search_chunks`; + sql`DROP TABLE IF EXISTS kernel_search_items`; + sql`DROP TABLE IF EXISTS kernel_search_pending`; + } + + sql` + CREATE TABLE IF NOT EXISTS kernel_search_items ( + item_id TEXT PRIMARY KEY, + source_version TEXT NOT NULL, + vector_status TEXT NOT NULL DEFAULT 'pending' + ) + `; + sql` + CREATE TABLE IF NOT EXISTS kernel_search_chunks ( + chunk_id TEXT PRIMARY KEY, + item_id TEXT NOT NULL, + page_number INTEGER, + start_line INTEGER, + end_line INTEGER + ) + `; + sql`CREATE INDEX IF NOT EXISTS kernel_search_chunks_item_idx + ON kernel_search_chunks (item_id)`; + sql` + CREATE VIRTUAL TABLE IF NOT EXISTS kernel_search_fts USING fts5( + title, + content, + tokenize = 'unicode61 remove_diacritics 2' + ) + `; + sql` + CREATE TABLE IF NOT EXISTS kernel_search_pending ( + item_id TEXT PRIMARY KEY, + requested_at INTEGER NOT NULL, + attempts INTEGER NOT NULL DEFAULT 0 + ) + `; + sql`CREATE INDEX IF NOT EXISTS kernel_search_pending_requested_idx + ON kernel_search_pending (requested_at)`; + + sql` + INSERT INTO kernel_meta (key, value, updated_at) + VALUES (${workspaceSearchVersionKey}, ${workspaceSearchIndexVersion}, ${Date.now()}) + ON CONFLICT(key) DO UPDATE SET + value = excluded.value, + updated_at = excluded.updated_at + `; +} diff --git a/src/features/workspaces/search/workspace-search-scope.ts b/src/features/workspaces/search/workspace-search-scope.ts new file mode 100644 index 00000000..8c7ae9a6 --- /dev/null +++ b/src/features/workspaces/search/workspace-search-scope.ts @@ -0,0 +1,192 @@ +import type { WorkspaceItemSummary } from "#/features/workspaces/contracts"; +import { + buildWorkspaceKernelTree, + normalizeWorkspacePath, + resolveWorkspaceKernelItemPath, + WorkspaceKernelPathError, +} from "#/features/workspaces/kernel/workspace-kernel-paths"; +import type { + WorkspaceSearchFailure, + WorkspaceSearchItemType, +} from "#/features/workspaces/search/workspace-search-contract"; + +const maximumVectorFilterBytes = 2_048; + +export type WorkspaceSearchVectorFilter = VectorizeVectorMetadataFilter | null; + +export type WorkspaceSearchLocalScope = + | { kind: "workspace" } + | { itemId: string; kind: "item" } + | { folderIds: string[]; kind: "folder" }; + +export interface WorkspaceSearchScope { + local: WorkspaceSearchLocalScope; + vectorFilters: WorkspaceSearchVectorFilter[]; +} + +export function resolveWorkspaceSearchScope(input: { + items: WorkspaceItemSummary[]; + path: string; + types: WorkspaceSearchItemType[]; +}): + | { scope: WorkspaceSearchScope; status: "ready" } + | { failure: WorkspaceSearchFailure; status: "failed" } { + let path: string; + try { + path = normalizeWorkspacePath(input.path); + } catch (error) { + if (error instanceof WorkspaceKernelPathError && error.code === "path_not_absolute") { + return { + failure: { code: error.code, path: input.path }, + status: "failed", + }; + } + throw error; + } + + if (path === "/") { + return { + scope: { + local: { kind: "workspace" }, + vectorFilters: createWorkspaceSearchVectorFilters({ + kind: "workspace", + types: input.types, + }), + }, + status: "ready", + }; + } + + const tree = buildWorkspaceKernelTree(input.items); + const item = resolveWorkspaceKernelItemPath(path, tree); + if (!item) { + return { + failure: { code: "path_not_found", path }, + status: "failed", + }; + } + if (item.type !== "folder") { + const itemType = isWorkspaceSearchItemType(item.type) ? item.type : null; + return { + scope: { + local: { itemId: item.id, kind: "item" }, + vectorFilters: + itemType && input.types.includes(itemType) + ? createWorkspaceSearchVectorFilters({ + itemId: item.id, + kind: "item", + }) + : [], + }, + status: "ready", + }; + } + + const folderIds = listWorkspaceSearchFolderIds(item.id, tree.childrenByParentId); + return { + scope: { + local: { folderIds, kind: "folder" }, + vectorFilters: createWorkspaceSearchVectorFilters({ + folderIds, + kind: "folder", + types: input.types, + }), + }, + status: "ready", + }; +} + +export function createWorkspaceSearchVectorMetadata(input: { + itemId: string; + parentId: string | null; + type: WorkspaceSearchItemType; +}): NonNullable { + return { + itemId: input.itemId, + ...(input.parentId ? { parentId: input.parentId } : {}), + type: input.type, + }; +} + +function createWorkspaceSearchVectorFilters( + input: + | { kind: "workspace"; types: WorkspaceSearchItemType[] } + | { itemId: string; kind: "item" } + | { folderIds: string[]; kind: "folder"; types: WorkspaceSearchItemType[] }, +): WorkspaceSearchVectorFilter[] { + if (input.kind === "item") { + return [{ itemId: input.itemId }]; + } + const type = input.types.length === 1 ? input.types[0] : undefined; + if (input.kind === "workspace") { + return type ? [{ type }] : [null]; + } + + return batchFolderFilters(input.folderIds, (folderIds) => { + const filter: VectorizeVectorMetadataFilter = { parentId: { $in: folderIds } }; + if (type) { + filter.type = type; + } + return filter; + }); +} + +function batchFolderFilters( + folderIds: string[], + createFilter: (folderIds: string[]) => VectorizeVectorMetadataFilter, +) { + const filters: VectorizeVectorMetadataFilter[] = []; + let batch: string[] = []; + + for (const folderId of folderIds) { + const nextBatch = [...batch, folderId]; + if (getFilterByteLength(createFilter(nextBatch)) < maximumVectorFilterBytes) { + batch = nextBatch; + continue; + } + if (batch.length === 0) { + throw new Error("Workspace folder identifier exceeds the Vectorize filter limit."); + } + filters.push(createFilter(batch)); + batch = [folderId]; + } + + if (batch.length > 0) { + filters.push(createFilter(batch)); + } + return filters; +} + +function getFilterByteLength(filter: VectorizeVectorMetadataFilter) { + return new TextEncoder().encode(JSON.stringify(filter)).byteLength; +} + +function isWorkspaceSearchItemType(type: WorkspaceItemSummary["type"]) { + return type === "document" || type === "file"; +} + +function listWorkspaceSearchFolderIds( + folderId: string, + childrenByParentId: Map, +) { + const folderIds: string[] = []; + const pending = [folderId]; + const visited = new Set(); + + while (pending.length > 0) { + const parentId = pending.pop(); + if (!parentId || visited.has(parentId)) { + continue; + } + visited.add(parentId); + folderIds.push(parentId); + + for (const item of childrenByParentId.get(parentId) ?? []) { + if (item.type === "folder") { + pending.push(item.id); + } + } + } + + return folderIds; +} diff --git a/src/features/workspaces/search/workspace-search-version.ts b/src/features/workspaces/search/workspace-search-version.ts new file mode 100644 index 00000000..f7f0a8dd --- /dev/null +++ b/src/features/workspaces/search/workspace-search-version.ts @@ -0,0 +1,20 @@ +/** + * Bump whenever the search schema, embedding model, chunking, or embedding text changes. + * The SQL freshness check in seedPendingItems must mirror the format built here. + */ +export const workspaceSearchIndexVersion = "v3-bge-m3-1800-scoped"; + +export function buildWorkspaceSearchSourceVersion( + input: + | { type: "document"; updatedAt: number } + | { + projectionUpdatedAt: number; + sourceHash: string; + type: "file"; + updatedAt: number; + }, +) { + return input.type === "document" + ? `${workspaceSearchIndexVersion}:document:${input.updatedAt}` + : `${workspaceSearchIndexVersion}:file:${input.updatedAt}:${input.projectionUpdatedAt}:${input.sourceHash}`; +} diff --git a/src/features/workspaces/search/workspace-search.test.ts b/src/features/workspaces/search/workspace-search.test.ts new file mode 100644 index 00000000..2993ab80 --- /dev/null +++ b/src/features/workspaces/search/workspace-search.test.ts @@ -0,0 +1,226 @@ +import { describe, expect, it } from "vitest"; + +import type { WorkspaceItemSummary } from "#/features/workspaces/contracts"; +import { iterateWorkspaceSearchTextChunks } from "#/features/workspaces/search/workspace-search-chunks"; +import { iteratePreparedWorkspaceSearchChunks } from "#/features/workspaces/search/workspace-search-content"; +import { workspaceSearchInputSchema } from "#/features/workspaces/search/workspace-search-contract"; +import { fuseWorkspaceSearchRanks } from "#/features/workspaces/search/workspace-search-ranking"; +import { + createWorkspaceSearchVectorMetadata, + resolveWorkspaceSearchScope, +} from "#/features/workspaces/search/workspace-search-scope"; +import { buildWorkspaceSearchSourceVersion } from "#/features/workspaces/search/workspace-search-version"; + +describe("workspace search", () => { + it("chunks long content with bounded overlap and line locations", () => { + const text = Array.from( + { length: 80 }, + (_, index) => `Line ${index + 1}: ${"searchable content ".repeat(4)}`, + ).join("\n"); + const chunks = Array.from(iterateWorkspaceSearchTextChunks(text)); + + expect(chunks.length).toBeGreaterThan(2); + expect(chunks[0]).toMatchObject({ startLine: 1 }); + expect(chunks.at(-1)?.endLine).toBe(80); + expect(chunks.every((chunk) => chunk.content.length <= 1_800)).toBe(true); + expect(chunks[1]?.startLine).toBeLessThanOrEqual((chunks[0]?.endLine ?? 0) + 1); + }); + + it("does not extend a chunk past its hard boundary for a separator", () => { + const [chunk] = Array.from( + iterateWorkspaceSearchTextChunks(`${"a".repeat(1_800)}. ${"b".repeat(100)}`), + ); + + expect(chunk?.content).toHaveLength(1_800); + }); + + it("deduplicates repeated content types", () => { + const input = workspaceSearchInputSchema.parse({ + query: "search files", + types: ["file", "file"], + }); + + expect(input.types).toEqual(["file"]); + }); + + it("builds canonical source versions for documents and files", () => { + expect( + buildWorkspaceSearchSourceVersion({ + type: "document", + updatedAt: 12, + }), + ).toBe("v3-bge-m3-1800-scoped:document:12"); + expect( + buildWorkspaceSearchSourceVersion({ + projectionUpdatedAt: 34, + sourceHash: "hash", + type: "file", + updatedAt: 12, + }), + ).toBe("v3-bge-m3-1800-scoped:file:12:34:hash"); + }); + + it("indexes document content beyond the former per-item cutoff", async () => { + const marker = "end-of-large-document"; + const checkpoint = JSON.stringify({ + type: "doc", + content: [ + { + type: "paragraph", + content: [{ type: "text", text: `${"searchable ".repeat(26_000)}${marker}` }], + }, + ], + }); + const chunks = []; + + for await (const chunk of iteratePreparedWorkspaceSearchChunks({ + bucket: {} as R2Bucket, + source: { + itemId: "large-document", + name: "Large document", + shellPath: "/large-document.json", + sourceVersion: "test", + type: "document", + }, + workspace: { + readFile: async () => checkpoint, + }, + })) { + chunks.push(chunk); + } + + expect(checkpoint.length).toBeGreaterThan(250_000); + expect(chunks.at(-1)?.content).toContain(marker); + }); + + it("fuses lexical and semantic ranks while diversifying items", () => { + const shared = { chunkId: "shared", itemId: "a" }; + const results = fuseWorkspaceSearchRanks({ + keyword: [ + shared, + { chunkId: "a-2", itemId: "a" }, + { chunkId: "a-3", itemId: "a" }, + { chunkId: "b-1", itemId: "b" }, + ], + limit: 4, + semantic: [shared, { chunkId: "c-1", itemId: "c" }], + }); + + expect(results[0]).toEqual(shared); + expect(results.filter((result) => result.itemId === "a")).toHaveLength(2); + expect(results.map((result) => result.itemId)).toContain("b"); + expect(results.map((result) => result.itemId)).toContain("c"); + }); + + it("resolves recursive folder search through stable parent identifiers", () => { + const items = [ + createItem({ id: "projects", name: "Projects", type: "folder" }), + createItem({ + id: "overview", + name: "Overview", + parentId: "projects", + type: "document", + }), + createItem({ id: "acme", name: "Acme", parentId: "projects", type: "folder" }), + createItem({ id: "plan", name: "Plan", parentId: "acme", type: "file" }), + createItem({ id: "elsewhere", name: "Elsewhere", type: "folder" }), + createItem({ id: "other", name: "Other", parentId: "elsewhere", type: "file" }), + ]; + const resolved = resolveWorkspaceSearchScope({ + items, + path: "/Projects", + types: ["document", "file"], + }); + + expect(resolved.status).toBe("ready"); + if (resolved.status !== "ready") { + return; + } + expect(resolved.scope.local.kind).toBe("folder"); + if (resolved.scope.local.kind !== "folder") { + return; + } + expect(new Set(resolved.scope.local.folderIds)).toEqual(new Set(["projects", "acme"])); + expect(new Set(readFilteredParentIds(resolved.scope.vectorFilters))).toEqual( + new Set(["projects", "acme"]), + ); + expect( + createWorkspaceSearchVectorMetadata({ + itemId: "plan", + parentId: "acme", + type: "file", + }), + ).toEqual({ itemId: "plan", parentId: "acme", type: "file" }); + }); + + it("keeps every recursive folder filter within Vectorize's byte limit", () => { + const items = [ + createItem({ id: "archive", name: "Archive", type: "folder" }), + ...Array.from({ length: 180 }, (_, index) => + createItem({ + id: `folder-${index}-${"x".repeat(24)}`, + name: `Folder ${index}`, + parentId: "archive", + type: "folder", + }), + ), + ]; + const resolved = resolveWorkspaceSearchScope({ + items, + path: "/Archive", + types: ["file"], + }); + + expect(resolved.status).toBe("ready"); + if (resolved.status !== "ready") { + return; + } + expect(resolved.scope.vectorFilters.length).toBeGreaterThan(1); + expect( + resolved.scope.vectorFilters.every( + (filter) => new TextEncoder().encode(JSON.stringify(filter)).byteLength < 2_048, + ), + ).toBe(true); + expect(new Set(readFilteredParentIds(resolved.scope.vectorFilters)).size).toBe(181); + }); +}); + +function createItem( + input: Pick & { + parentId?: string | null; + }, +): WorkspaceItemSummary { + return { + color: null, + createdAt: "2026-01-01T00:00:00.000Z", + deletedAt: null, + id: input.id, + meta: "", + metadataJson: {}, + name: input.name, + parentId: input.parentId ?? null, + sortOrder: 0, + title: input.name, + type: input.type, + updatedAt: "2026-01-01T00:00:00.000Z", + workspaceId: "workspace-1", + }; +} + +function readFilteredParentIds(filters: Array) { + return filters.flatMap((filter) => { + if (!filter || !("parentId" in filter)) { + return []; + } + const parentId = filter.parentId; + if ( + parentId !== null && + typeof parentId === "object" && + "$in" in parentId && + Array.isArray(parentId.$in) + ) { + return parentId.$in.filter((value): value is string => typeof value === "string"); + } + return []; + }); +} diff --git a/worker-configuration.d.ts b/worker-configuration.d.ts index 5ff33da3..c4c1824b 100644 --- a/worker-configuration.d.ts +++ b/worker-configuration.d.ts @@ -1,9 +1,10 @@ /* eslint-disable */ -// Generated by Wrangler by running `wrangler types` (hash: 5348c83093b741e804d43d2ed057d3fa) +// Generated by Wrangler by running `wrangler types` (hash: a0dd3d735037ed637999df03419f817d) // Runtime types generated with workerd@1.20260722.1 2026-03-24 nodejs_compat interface __BaseEnv_Env { WORKSPACE_KERNEL_FILES: R2Bucket; DB: D1Database; + WORKSPACE_SEARCH: VectorizeIndex; EMAIL: SendEmail; LOADER: WorkerLoader; BROWSER: BrowserRun; @@ -42,6 +43,7 @@ declare namespace Cloudflare { interface StagingEnv { WORKSPACE_KERNEL_FILES: R2Bucket; DB: D1Database; + WORKSPACE_SEARCH: VectorizeIndex; EMAIL: SendEmail; LOADER: WorkerLoader; BROWSER: BrowserRun; @@ -75,6 +77,7 @@ declare namespace Cloudflare { interface ProductionEnv { WORKSPACE_KERNEL_FILES: R2Bucket; DB: D1Database; + WORKSPACE_SEARCH: VectorizeIndex; EMAIL: SendEmail; LOADER: WorkerLoader; BROWSER: BrowserRun; diff --git a/wrangler.jsonc b/wrangler.jsonc index 1eed3b57..415251d4 100644 --- a/wrangler.jsonc +++ b/wrangler.jsonc @@ -34,6 +34,14 @@ "binding": "AI", "remote": true, }, + // Provisioning: docs/configuration/deployments.mdx#workspace-search-vector-indexes + "vectorize": [ + { + "binding": "WORKSPACE_SEARCH", + "index_name": "thinkex-workspace-search-staging", + "remote": true, + }, + ], "images": { "binding": "IMAGES", }, @@ -226,6 +234,13 @@ "binding": "AI", "remote": true, }, + "vectorize": [ + { + "binding": "WORKSPACE_SEARCH", + "index_name": "thinkex-workspace-search-staging", + "remote": true, + }, + ], "images": { "binding": "IMAGES", }, @@ -350,6 +365,13 @@ "binding": "AI", "remote": true, }, + "vectorize": [ + { + "binding": "WORKSPACE_SEARCH", + "index_name": "thinkex-workspace-search", + "remote": true, + }, + ], "images": { "binding": "IMAGES", },