From 1e5d47e5cbf864e306886f00208c1f791e2d4379 Mon Sep 17 00:00:00 2001 From: Urjit Chakraborty <135136842+urjitc@users.noreply.github.com> Date: Thu, 30 Jul 2026 11:59:27 -0400 Subject: [PATCH 01/13] feat(workspaces): add semantic workspace search Add a Vectorize-backed workspace search module (indexer, embeddings, chunking, ranking, query) with a search-workspace operation and AI tool wiring, plus the WORKSPACE_SEARCH Vectorize binding. Co-Authored-By: Claude Opus 4.8 --- .../workspaces/ai/ai-tool-registry.ts | 1 + .../workspaces/ai/workspace-citations.ts | 13 +- .../ai/workspace-tool-result-adapters.ts | 47 ++ src/features/workspaces/ai/workspace-tools.ts | 24 +- .../extraction/workspace-page-projection.ts | 34 ++ .../kernel/workspace-kernel-access.ts | 8 + .../kernel/workspace-kernel-events.ts | 4 + .../kernel/workspace-kernel-item-commands.ts | 23 +- .../workspaces/kernel/workspace-kernel.ts | 48 +- .../workspaces/operations/search-workspace.ts | 23 + .../workspace-operation-observability.ts | 7 + .../operations/workspace-tool-definitions.ts | 21 +- .../operations/workspace-tool-schemas.ts | 25 +- .../search/workspace-search-chunks.ts | 90 ++++ .../search/workspace-search-content.ts | 94 ++++ .../search/workspace-search-contract.ts | 74 +++ .../search/workspace-search-embeddings.ts | 41 ++ .../search/workspace-search-indexer.ts | 468 ++++++++++++++++++ .../search/workspace-search-projection.ts | 91 ++++ .../search/workspace-search-query.ts | 351 +++++++++++++ .../search/workspace-search-ranking.ts | 53 ++ .../search/workspace-search-references.ts | 51 ++ .../search/workspace-search-schema.ts | 44 ++ .../search/workspace-search.test.ts | 39 ++ worker-configuration.d.ts | 5 +- wrangler.jsonc | 21 + 26 files changed, 1655 insertions(+), 45 deletions(-) create mode 100644 src/features/workspaces/ai/workspace-tool-result-adapters.ts create mode 100644 src/features/workspaces/operations/search-workspace.ts create mode 100644 src/features/workspaces/search/workspace-search-chunks.ts create mode 100644 src/features/workspaces/search/workspace-search-content.ts create mode 100644 src/features/workspaces/search/workspace-search-contract.ts create mode 100644 src/features/workspaces/search/workspace-search-embeddings.ts create mode 100644 src/features/workspaces/search/workspace-search-indexer.ts create mode 100644 src/features/workspaces/search/workspace-search-projection.ts create mode 100644 src/features/workspaces/search/workspace-search-query.ts create mode 100644 src/features/workspaces/search/workspace-search-ranking.ts create mode 100644 src/features/workspaces/search/workspace-search-references.ts create mode 100644 src/features/workspaces/search/workspace-search-schema.ts create mode 100644 src/features/workspaces/search/workspace-search.test.ts 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..036c153a --- /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; + }, + }; +} + +export const workspaceReadItemsResultAdapter = defineWorkspaceToolResultAdapter({ + collectReferences: (output) => output.references, + outputSchema: workspaceReadItemsOutputSchema, + projectOutput: createWorkspaceReadItemsModelOutput, +}); + +export 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/extraction/workspace-page-projection.ts b/src/features/workspaces/extraction/workspace-page-projection.ts index 19a55124..3edc6e0c 100644 --- a/src/features/workspaces/extraction/workspace-page-projection.ts +++ b/src/features/workspaces/extraction/workspace-page-projection.ts @@ -211,6 +211,40 @@ 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 input.bucket.get(getWorkspacePageObjectKey(prefix, pageNumber)); + if (!object) { + throw new Error(`Extracted page ${pageNumber} was not found.`); + } + + 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.`); + } + + yield { + markdown: await object.text(), + pageNumber, + }; + } +} + function requireManifestPage( pagesByNumber: ReadonlyMap, pageNumber: number, diff --git a/src/features/workspaces/kernel/workspace-kernel-access.ts b/src/features/workspaces/kernel/workspace-kernel-access.ts index 2644b61d..edd46f8e 100644 --- a/src/features/workspaces/kernel/workspace-kernel-access.ts +++ b/src/features/workspaces/kernel/workspace-kernel-access.ts @@ -36,6 +36,11 @@ 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, +} from "#/features/workspaces/search/workspace-search-contract"; import { assertCanMutateWorkspace, assertCanReadWorkspace, @@ -114,6 +119,9 @@ export interface WorkspaceKernelClient { actorUserId?: string | null; clientMutationId?: string | null; }): Promise>; + searchWorkspace( + input: WorkspaceSearchInput, + ): Promise<{ failed: WorkspaceSearchFailure[]; results: WorkspaceSearchResult[] }>; 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..8f38de7c 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 { @@ -356,21 +355,13 @@ export class WorkspaceKernelItemCommands { const now = Date.now(); - 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: item.metadata_json, + sql: this.sql, + updatedAt: now, + }); 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..d855560e 100644 --- a/src/features/workspaces/kernel/workspace-kernel.ts +++ b/src/features/workspaces/kernel/workspace-kernel.ts @@ -63,6 +63,8 @@ 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; @@ -82,11 +84,22 @@ 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) => this.search.observe(event), }); private readonly relations = new WorkspaceKernelRelations(this.kernelSql); private readonly itemCommands = new WorkspaceKernelItemCommands({ @@ -105,8 +118,12 @@ export class WorkspaceKernel extends Agent { workspaceId: () => this.name, }); - onStart() { + async onStart() { initializeWorkspaceKernelStorage(this.kernelSql); + this.search.initialize(); + if (this.search.hasRetryablePending()) { + await this.scheduleWorkspaceSearchIndexing(); + } } onConnect(connection: Connection, context: ConnectionContext) { @@ -283,6 +300,16 @@ export class WorkspaceKernel extends Agent { ); } + async searchWorkspace(input: WorkspaceSearchInput) { + return await this.search.search(input); + } + + async processWorkspaceSearchIndex() { + if (await this.search.processBatch()) { + await this.scheduleWorkspaceSearchIndexing(); + } + } + private async runMutation( operation: string, input: { actorUserId?: string | null; clientMutationId?: string | null }, @@ -319,6 +346,17 @@ export class WorkspaceKernel extends Agent { 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, { @@ -350,7 +388,13 @@ export class WorkspaceKernel extends Agent { ]); await this.ctx.storage.deleteAll(); - return { attempted: documentItemIds.length + 1, failed }; + return { attempted: documentItemIds.length + 2, failed }; + } + + private async scheduleWorkspaceSearchIndexing() { + await this.schedule(1, "processWorkspaceSearchIndex", undefined, { + idempotent: true, + }); } 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-operation-observability.ts b/src/features/workspaces/operations/workspace-operation-observability.ts index 3b5a62db..eaa880d9 100644 --- a/src/features/workspaces/operations/workspace-operation-observability.ts +++ b/src/features/workspaces/operations/workspace-operation-observability.ts @@ -73,6 +73,13 @@ export function summarizeWorkspaceReadResult(input: { return summarizeWorkspaceResult(succeededCount, failures, pendingCount); } +export function summarizeWorkspaceSearchResult(input: { + failed: ReadonlyArray<{ code: string }>; + results: readonly unknown[]; +}) { + return summarizeWorkspaceResult(input.results.length, input.failed); +} + export function summarizeWorkspaceItemResult(input: { failed: ReadonlyArray<{ code: string }>; item?: unknown; diff --git a/src/features/workspaces/operations/workspace-tool-definitions.ts b/src/features/workspaces/operations/workspace-tool-definitions.ts index 9927caf4..865e7124 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, @@ -46,6 +50,7 @@ import { summarizeWorkspaceCollectionResult, summarizeWorkspaceItemResult, summarizeWorkspaceReadResult, + summarizeWorkspaceSearchResult, type WorkspaceOperationSummary, } from "#/features/workspaces/operations/workspace-operation-observability"; @@ -146,7 +151,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 +161,20 @@ 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: summarizeWorkspaceSearchResult, + 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-chunks.ts b/src/features/workspaces/search/workspace-search-chunks.ts new file mode 100644 index 00000000..8ae942be --- /dev/null +++ b/src/features/workspaces/search/workspace-search-chunks.ts @@ -0,0 +1,90 @@ +const targetChunkCharacters = 1_800; +const minimumChunkCharacters = 900; +const overlapCharacters = 220; + +export interface WorkspaceSearchTextChunk { + content: string; + endLine: number; + startLine: number; +} + +export function chunkWorkspaceSearchText(text: string): WorkspaceSearchTextChunk[] { + const normalized = text.replace(/\r\n?/g, "\n"); + if (normalized.length === 0) { + return [{ content: "", endLine: 0, startLine: 0 }]; + } + + const chunks: WorkspaceSearchTextChunk[] = []; + 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; + } + chunks.push({ + 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); + } + + return chunks.length > 0 ? chunks : [{ 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); + 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..ff05a399 --- /dev/null +++ b/src/features/workspaces/search/workspace-search-content.ts @@ -0,0 +1,94 @@ +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 { chunkWorkspaceSearchText } from "#/features/workspaces/search/workspace-search-chunks"; + +const maximumIndexedCharactersPerItem = 8_000_000; + +export interface WorkspaceSearchFileSystem { + readFile(path: string): Promise; +} + +interface WorkspaceSearchIndexSourceBase { + itemId: string; + name: string; + path: string; + sourceVersion: string; +} + +export type WorkspaceSearchIndexSource = WorkspaceSearchIndexSourceBase & + ( + | { shellPath: string; type: "document" } + | { objectKey: string; sourceHash: string; type: "file" } + ); + +export interface PreparedWorkspaceSearchChunk { + content: string; + endLine: number | null; + index: number; + pageNumber: number | null; + startLine: number | null; +} + +export async function prepareWorkspaceSearchChunks(input: { + bucket: R2Bucket; + source: WorkspaceSearchIndexSource; + workspace: WorkspaceSearchFileSystem; +}): Promise { + 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)); + const searchable = markdown.slice(0, maximumIndexedCharactersPerItem); + + return chunkWorkspaceSearchText(searchable).map((chunk, index) => ({ + content: chunk.content, + endLine: chunk.endLine, + index, + pageNumber: null, + startLine: chunk.startLine, + })); + } + + const chunks: PreparedWorkspaceSearchChunk[] = []; + let indexedCharacters = 0; + + for await (const page of iterateWorkspacePageProjection({ + bucket: input.bucket, + expectedSourceHash: input.source.sourceHash, + manifestObjectKey: input.source.objectKey, + })) { + const remaining = maximumIndexedCharactersPerItem - indexedCharacters; + if (remaining <= 0) { + break; + } + + const markdown = page.markdown.slice(0, remaining); + for (const chunk of chunkWorkspaceSearchText(markdown)) { + chunks.push({ + content: chunk.content, + endLine: null, + index: chunks.length, + pageNumber: page.pageNumber, + startLine: null, + }); + } + indexedCharacters += markdown.length; + + if (markdown.length < page.markdown.length) { + break; + } + } + + return chunks; +} + +export function createWorkspaceSearchEmbeddingText(input: { + content: string; + path: string; + title: string; +}) { + return [`Title: ${input.title}`, `Path: ${input.path}`, "", 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..f838e4e0 --- /dev/null +++ b/src/features/workspaces/search/workspace-search-contract.ts @@ -0,0 +1,74 @@ +import { z } from "zod"; + +import { workspaceReferenceRecordSchema } from "#/features/workspaces/locations/workspace-location"; +import { workspaceFileAssetKindSchema } from "#/features/workspaces/model/workspace-file"; + +export 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) + .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(), + }), +]); + +export 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, +}); + +export const workspaceSearchFailureSchema = z.object({ + code: z.enum(["path_not_absolute", "path_not_found"]), + path: z.string(), +}); + +export const workspaceSearchOutputSchema = z.object({ + failed: z.array(workspaceSearchFailureSchema), + references: z.array(workspaceReferenceRecordSchema), + results: z.array(workspaceSearchResultSchema), +}); + +export type WorkspaceSearchInput = z.output; +export type WorkspaceSearchItemType = z.output; +export type WorkspaceSearchResult = z.output; +export type WorkspaceSearchFailure = 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..f9f21475 --- /dev/null +++ b/src/features/workspaces/search/workspace-search-embeddings.ts @@ -0,0 +1,41 @@ +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; +} + +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; +} + +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..f6610f6e --- /dev/null +++ b/src/features/workspaces/search/workspace-search-indexer.ts @@ -0,0 +1,468 @@ +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 { + createWorkspaceSearchEmbeddingText, + prepareWorkspaceSearchChunks, + type WorkspaceSearchFileSystem, + type WorkspaceSearchIndexSource, +} from "#/features/workspaces/search/workspace-search-content"; +import { + batchWorkspaceSearchValues, + embedWorkspaceSearchTexts, +} from "#/features/workspaces/search/workspace-search-embeddings"; +import { recordOperationalFailure } from "#/integrations/observability/operational-events"; +import { sha256Base64UrlText } from "#/lib/binary"; + +const searchIndexBatchSize = 4; +const vectorMutationBatchSize = 1_000; +const vectorDeleteBatchSize = 100; +const maximumIndexAttempts = 5; + +interface SearchSourceRow { + id: string; + name: string; + projection_object_key: string | null; + projection_source_hash: string | null; + projection_updated_at: number | null; + shell_path: string; + type: string; + updated_at: number; +} + +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 getItems: () => WorkspaceItemSummary[]; + private readonly sql: WorkspaceKernelSql; + private readonly vectorize: VectorizeIndex; + private readonly workspace: WorkspaceSearchFileSystem; + private readonly workspaceId: () => string; + + constructor(input: { + ai: Ai; + bucket: R2Bucket; + getItems: () => WorkspaceItemSummary[]; + sql: WorkspaceKernelSql; + vectorize: VectorizeIndex; + workspace: WorkspaceSearchFileSystem; + workspaceId: () => string; + }) { + this.ai = input.ai; + this.bucket = input.bucket; + this.getItems = input.getItems; + 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, attempts) + SELECT i.id, ${now}, 0 + 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 'document:' || i.updated_at + ELSE 'file:' || i.updated_at || ':' || p.updated_at || ':' || p.source_hash + END + OR s.vector_ready = 0 + ) + ON CONFLICT(item_id) DO NOTHING + `; + this.sql` + INSERT INTO kernel_search_pending (item_id, requested_at, attempts) + SELECT s.item_id, ${now}, 0 + 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, attempts) + VALUES (${itemId}, ${Date.now()}, 0) + ON CONFLICT(item_id) DO UPDATE SET + requested_at = excluded.requested_at, + attempts = 0 + `; + } + + markTreePending(itemId: string) { + this.sql` + WITH RECURSIVE search_tree(id) AS ( + SELECT ${itemId} + UNION ALL + SELECT i.id + FROM kernel_items i + JOIN search_tree parent ON i.parent_id = parent.id + WHERE i.deleted_at IS NULL + ) + INSERT INTO kernel_search_pending (item_id, requested_at, attempts) + SELECT i.id, ${Date.now()}, 0 + FROM kernel_items i + JOIN search_tree tree ON tree.id = i.id + WHERE i.deleted_at IS NULL AND i.type IN ('document', 'file') + ON CONFLICT(item_id) DO UPDATE SET + requested_at = excluded.requested_at, + attempts = 0 + `; + } + + hasRetryablePending() { + return Boolean( + this.sql<{ item_id: string }>` + SELECT item_id + FROM kernel_search_pending + WHERE attempts < ${maximumIndexAttempts} + LIMIT 1 + `[0], + ); + } + + async processBatch() { + await this.flushVectorDeletes(); + const pending = this.sql<{ item_id: string }>` + SELECT item_id + FROM kernel_search_pending + WHERE attempts < ${maximumIndexAttempts} + ORDER BY requested_at ASC + LIMIT ${searchIndexBatchSize} + `; + + for (const row of pending) { + try { + await this.indexItem(row.item_id); + } catch (error) { + this.recordIndexFailure(row.item_id, error); + } + } + + await this.flushVectorDeletes(); + return this.hasRetryablePending(); + } + + 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), + ); + + for (const batch of batchWorkspaceSearchValues(Array.from(ids), vectorDeleteBatchSize)) { + await this.vectorize.deleteByIds(batch); + } + } + + private async indexItem(itemId: string) { + const source = this.getIndexSource(itemId); + if (!source) { + this.removeIndexedItem(itemId); + return; + } + + const preparedChunks = await prepareWorkspaceSearchChunks({ + bucket: this.bucket, + source, + workspace: this.workspace, + }); + if (!this.isCurrentSource(source)) { + return; + } + + const revisionKey = await sha256Base64UrlText( + `${this.workspaceId()}:${source.itemId}:${source.sourceVersion}`, + ); + const chunks: SearchIndexChunk[] = preparedChunks.map((chunk) => ({ + ...chunk, + chunkId: `s${revisionKey}-${chunk.index}`, + })); + this.replaceKeywordIndex({ + chunks, + source, + }); + + const embeddings = await embedWorkspaceSearchTexts( + this.ai, + chunks.map((chunk) => + createWorkspaceSearchEmbeddingText({ + content: chunk.content, + path: source.path, + title: source.name, + }), + ), + ); + if (embeddings.length !== chunks.length) { + throw new Error("Workspace search embedding response did not match the indexed chunks."); + } + + const vectors = chunks.map( + (chunk, index): VectorizeVector => ({ + id: chunk.chunkId, + namespace: this.workspaceId(), + values: embeddings[index] ?? [], + }), + ); + for (const batch of batchWorkspaceSearchValues(vectors, vectorMutationBatchSize)) { + await this.vectorize.upsert(batch); + } + + if (!this.isCurrentSource(source)) { + for (const vector of vectors) { + this.queueVectorDelete(vector.id); + } + return; + } + + this.markVectorIndexReady( + source, + chunks.map((chunk) => chunk.chunkId), + ); + } + + private getIndexSource(itemId: string): WorkspaceSearchIndexSource | null { + const row = this.sql` + SELECT + i.id, + i.type, + i.name, + 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 path = buildWorkspaceKernelItemPathIndex(this.getItems()).get(row.id); + if (!path) { + return null; + } + + const source = { + itemId: row.id, + name: row.name, + path, + }; + + if (row.type === "document") { + return { + ...source, + shellPath: row.shell_path, + sourceVersion: `document:${row.updated_at}`, + type: "document", + }; + } + if (!row.projection_object_key || !row.projection_source_hash) { + return null; + } + + return { + ...source, + objectKey: row.projection_object_key, + sourceHash: row.projection_source_hash, + sourceVersion: `file:${row.updated_at}:${row.projection_updated_at}:${row.projection_source_hash}`, + type: "file", + }; + } + + private isCurrentSource(source: WorkspaceSearchIndexSource) { + const current = this.getIndexSource(source.itemId); + return ( + current?.sourceVersion === source.sourceVersion && + current.name === source.name && + current.path === source.path + ); + } + + private replaceKeywordIndex(input: { + chunks: SearchIndexChunk[]; + source: WorkspaceSearchIndexSource; + }) { + const retainedChunkIds = new Set(input.chunks.map((chunk) => chunk.chunkId)); + for (const row of this.sql<{ chunk_id: string }>` + SELECT chunk_id + FROM kernel_search_chunks + WHERE item_id = ${input.source.itemId} + `) { + if (!retainedChunkIds.has(row.chunk_id)) { + this.queueVectorDelete(row.chunk_id); + } + } + + this.deleteLocalChunks(input.source.itemId); + for (const chunk of input.chunks) { + this.insertChunk(input.source, chunk); + } + this.sql` + INSERT INTO kernel_search_items ( + item_id, + source_version, + vector_ready + ) + VALUES ( + ${input.source.itemId}, + ${input.source.sourceVersion}, + 0 + ) + ON CONFLICT(item_id) DO UPDATE SET + source_version = excluded.source_version, + vector_ready = 0 + `; + } + + private insertChunk(source: WorkspaceSearchIndexSource, 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 (chunk_id, title, path, content) + VALUES (${chunk.chunkId}, ${source.name}, ${source.path}, ${chunk.content}) + `; + } + + private markVectorIndexReady(source: WorkspaceSearchIndexSource, vectorIds: readonly string[]) { + this.sql` + UPDATE kernel_search_items + SET vector_ready = 1 + WHERE item_id = ${source.itemId} AND source_version = ${source.sourceVersion} + `; + this.sql`DELETE FROM kernel_search_pending WHERE item_id = ${source.itemId}`; + if (vectorIds.length > 0) { + this.sql` + DELETE FROM kernel_search_vector_deletes + WHERE vector_id IN ( + SELECT value FROM json_each(${JSON.stringify(vectorIds)}) + ) + `; + } + } + + private removeIndexedItem(itemId: string) { + for (const row of this.sql<{ chunk_id: string }>` + SELECT chunk_id + FROM kernel_search_chunks + WHERE item_id = ${itemId} + `) { + this.queueVectorDelete(row.chunk_id); + } + this.deleteLocalChunks(itemId); + this.sql`DELETE FROM kernel_search_items WHERE item_id = ${itemId}`; + this.sql`DELETE FROM kernel_search_pending WHERE item_id = ${itemId}`; + } + + private deleteLocalChunks(itemId: string) { + const chunkIds = this.sql<{ chunk_id: string }>` + SELECT chunk_id + FROM kernel_search_chunks + WHERE item_id = ${itemId} + `.map((row) => row.chunk_id); + if (chunkIds.length > 0) { + this.sql` + DELETE FROM kernel_search_fts + WHERE chunk_id IN (SELECT value FROM json_each(${JSON.stringify(chunkIds)})) + `; + } + this.sql`DELETE FROM kernel_search_chunks WHERE item_id = ${itemId}`; + } + + 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 async flushVectorDeletes() { + const ids = this.sql<{ vector_id: string }>` + SELECT vector_id + FROM kernel_search_vector_deletes + 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) { + recordOperationalFailure({ + error, + event: "workspace_search_vector_cleanup", + fields: { workspace_id: this.workspaceId() }, + }); + } + } + + private recordIndexFailure(itemId: string, error: unknown) { + this.sql` + UPDATE kernel_search_pending + SET attempts = attempts + 1 + WHERE item_id = ${itemId} + `; + recordOperationalFailure({ + error, + event: "workspace_search_indexing", + fields: { + item_id: itemId, + workspace_id: this.workspaceId(), + }, + }); + } +} 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..d8ff7242 --- /dev/null +++ b/src/features/workspaces/search/workspace-search-projection.ts @@ -0,0 +1,91 @@ +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": + case "workspace.item.moved": + this.indexer.markTreePending(event.payload.item.id); + break; + case "workspace.items.moved": + for (const item of event.payload.items) { + this.indexer.markTreePending(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; + } + + if (this.indexer.hasRetryablePending()) { + this.requestRun(); + } + } + + hasRetryablePending() { + return this.indexer.hasRetryablePending(); + } + + 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..ee37ed7d --- /dev/null +++ b/src/features/workspaces/search/workspace-search-query.ts @@ -0,0 +1,351 @@ +import type { WorkspaceItemSummary } from "#/features/workspaces/contracts"; +import { + buildWorkspaceKernelItemPathIndex, + buildWorkspaceKernelTree, + normalizeWorkspacePath, + resolveWorkspaceKernelItemPath, + WorkspaceKernelPathError, + type WorkspaceKernelTree, +} 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, +} from "#/features/workspaces/search/workspace-search-contract"; +import { embedWorkspaceSearchTexts } from "#/features/workspaces/search/workspace-search-embeddings"; +import { + fuseWorkspaceSearchRanks, + type WorkspaceSearchRankCandidate, +} from "#/features/workspaces/search/workspace-search-ranking"; +import { recordOperationalFailure } from "#/integrations/observability/operational-events"; + +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[] }> { + const scope = this.resolveScope(input.path ?? "/"); + if (scope.status === "failed") { + return { failed: [scope.failure], results: [] }; + } + + const limit = input.limit ?? 10; + const types = input.types ?? ["document", "file"]; + const candidateLimit = Math.min(100, Math.max(50, limit * 6)); + const keyword = this.searchKeyword({ + candidateLimit, + query: input.query, + scopeItemIds: scope.itemIds, + types, + }); + const semantic = await this.searchSemanticWithFallback({ + candidateLimit, + query: input.query, + scopeItemIds: scope.itemIds, + types, + }); + const items = this.getItems(); + const itemsById = new Map(items.map((item) => [item.id, item])); + const paths = buildWorkspaceKernelItemPathIndex(items); + const ranked = fuseWorkspaceSearchRanks({ keyword, limit, semantic }); + + return { + failed: [], + 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 async searchSemanticWithFallback(input: { + candidateLimit: number; + query: string; + scopeItemIds: string[] | null; + types: WorkspaceSearchItemType[]; + }) { + try { + return await this.searchSemantic(input); + } catch (error) { + recordOperationalFailure({ + error, + event: "workspace_search_semantic", + fields: { workspace_id: this.workspaceId() }, + }); + return []; + } + } + + private resolveScope( + requestedPath: string, + ): + | { itemIds: string[] | null; status: "ready" } + | { failure: WorkspaceSearchFailure; status: "failed" } { + let path: string; + try { + path = normalizeWorkspacePath(requestedPath); + } catch (error) { + if (error instanceof WorkspaceKernelPathError && error.code === "path_not_absolute") { + return { + failure: { code: error.code, path: requestedPath }, + status: "failed", + }; + } + throw error; + } + + if (path === "/") { + return { itemIds: null, status: "ready" }; + } + + const items = this.getItems(); + const tree = buildWorkspaceKernelTree(items); + const item = resolveWorkspaceKernelItemPath(path, tree); + if (!item) { + return { + failure: { code: "path_not_found", path }, + status: "failed", + }; + } + if (item.type !== "folder") { + return { itemIds: [item.id], status: "ready" }; + } + + return { itemIds: listDescendantItemIds(item.id, tree), status: "ready" }; + } + + private searchKeyword(input: { + candidateLimit: number; + query: string; + scopeItemIds: string[] | null; + types: WorkspaceSearchItemType[]; + }): SearchCandidate[] { + const match = createFtsMatchExpression(input.query); + if (!match) { + return []; + } + + const scopeJson = input.scopeItemIds ? JSON.stringify(input.scopeItemIds) : null; + 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.chunk_id = kernel_search_fts.chunk_id + 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 + LEFT JOIN kernel_item_projections p + ON p.item_id = i.id + AND p.format = 'pages' + AND p.status = 'ready' + WHERE kernel_search_fts MATCH ${match} + AND s.source_version = CASE + WHEN i.type = 'document' THEN 'document:' || i.updated_at + ELSE 'file:' || i.updated_at || ':' || p.updated_at || ':' || p.source_hash + END + AND (${scopeJson} IS NULL OR c.item_id IN ( + SELECT value FROM json_each(${scopeJson}) + )) + AND i.type IN (SELECT value FROM json_each(${typesJson})) + ORDER BY bm25(kernel_search_fts, 0.0, 8.0, 4.0, 1.0) ASC + LIMIT ${input.candidateLimit} + `; + + return rows.map(mapSearchCandidate); + } + + private async searchSemantic(input: { + candidateLimit: number; + query: string; + scopeItemIds: string[] | null; + types: WorkspaceSearchItemType[]; + }): Promise { + const [embedding] = await embedWorkspaceSearchTexts(this.ai, [input.query]); + if (!embedding) { + return []; + } + + const matches = await this.vectorize.query(embedding, { + namespace: this.workspaceId(), + returnMetadata: "none", + topK: input.candidateLimit, + }); + const vectorIds = matches.matches.map((match) => match.id); + if (vectorIds.length === 0) { + return []; + } + + const rows = this.loadSemanticCandidates({ + scopeItemIds: input.scopeItemIds, + types: input.types, + vectorIds, + }); + const byVectorId = new Map(rows.map((row) => [row.chunk_id, mapSearchCandidate(row)])); + + return vectorIds.flatMap((vectorId) => { + const candidate = byVectorId.get(vectorId); + return candidate ? [candidate] : []; + }); + } + + private loadSemanticCandidates(input: { + scopeItemIds: string[] | null; + types: WorkspaceSearchItemType[]; + vectorIds: string[]; + }) { + const scopeJson = input.scopeItemIds ? JSON.stringify(input.scopeItemIds) : null; + 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.chunk_id = c.chunk_id + 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_ready = 1 + LEFT JOIN kernel_item_projections p + ON p.item_id = i.id + AND p.format = 'pages' + AND p.status = 'ready' + WHERE c.chunk_id IN (SELECT value FROM json_each(${JSON.stringify(input.vectorIds)})) + AND s.source_version = CASE + WHEN i.type = 'document' THEN 'document:' || i.updated_at + ELSE 'file:' || i.updated_at || ':' || p.updated_at || ':' || p.source_hash + END + AND (${scopeJson} IS NULL OR c.item_id IN ( + SELECT value FROM json_each(${scopeJson}) + )) + AND i.type IN (SELECT value FROM json_each(${typesJson})) + `; + } +} + +function listDescendantItemIds(folderId: string, tree: WorkspaceKernelTree) { + const ids: string[] = []; + const pending = [folderId]; + while (pending.length > 0) { + const itemId = pending.pop(); + if (itemId) { + ids.push(itemId); + pending.push(...(tree.childrenByParentId.get(itemId) ?? []).map((item) => item.id)); + } + } + return ids; +} + +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 = query.match(/[\p{L}\p{N}_]+/gu)?.slice(0, 20) ?? []; + 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 term = query.match(/[\p{L}\p{N}_]{3,}/u)?.[0]?.toLocaleLowerCase(); + const matchIndex = term ? content.toLocaleLowerCase().indexOf(term) : -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 ? "…" : ""}`; +} 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.ts b/src/features/workspaces/search/workspace-search-references.ts new file mode 100644 index 00000000..4f9d1ff8 --- /dev/null +++ b/src/features/workspaces/search/workspace-search-references.ts @@ -0,0 +1,51 @@ +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, + 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..8ddd9d61 --- /dev/null +++ b/src/features/workspaces/search/workspace-search-schema.ts @@ -0,0 +1,44 @@ +import type { WorkspaceKernelSql } from "#/features/workspaces/kernel/workspace-kernel-schema"; + +export function initializeWorkspaceSearchStorage(sql: WorkspaceKernelSql) { + sql` + CREATE TABLE IF NOT EXISTS kernel_search_items ( + item_id TEXT PRIMARY KEY, + source_version TEXT NOT NULL, + vector_ready INTEGER NOT NULL DEFAULT 0 + ) + `; + 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( + chunk_id UNINDEXED, + title, + path, + 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 TABLE IF NOT EXISTS kernel_search_vector_deletes ( + vector_id TEXT PRIMARY KEY, + requested_at INTEGER NOT NULL + ) + `; +} 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..e27add4d --- /dev/null +++ b/src/features/workspaces/search/workspace-search.test.ts @@ -0,0 +1,39 @@ +import { describe, expect, it } from "vitest"; + +import { chunkWorkspaceSearchText } from "#/features/workspaces/search/workspace-search-chunks"; +import { fuseWorkspaceSearchRanks } from "#/features/workspaces/search/workspace-search-ranking"; + +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 = chunkWorkspaceSearchText(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("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"); + }); +}); 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..14e5ffc0 100644 --- a/wrangler.jsonc +++ b/wrangler.jsonc @@ -34,6 +34,13 @@ "binding": "AI", "remote": true, }, + "vectorize": [ + { + "binding": "WORKSPACE_SEARCH", + "index_name": "thinkex-workspace-search-staging", + "remote": true, + }, + ], "images": { "binding": "IMAGES", }, @@ -226,6 +233,13 @@ "binding": "AI", "remote": true, }, + "vectorize": [ + { + "binding": "WORKSPACE_SEARCH", + "index_name": "thinkex-workspace-search-staging", + "remote": true, + }, + ], "images": { "binding": "IMAGES", }, @@ -350,6 +364,13 @@ "binding": "AI", "remote": true, }, + "vectorize": [ + { + "binding": "WORKSPACE_SEARCH", + "index_name": "thinkex-workspace-search", + "remote": true, + }, + ], "images": { "binding": "IMAGES", }, From df6952e75bb9f0cef124070b807cafe8664868ef Mon Sep 17 00:00:00 2001 From: Urjit Chakraborty <135136842+urjitc@users.noreply.github.com> Date: Thu, 30 Jul 2026 22:00:01 -0400 Subject: [PATCH 02/13] fix(workspace-search): harden indexing and scoped retrieval --- .../extraction/workspace-page-projection.ts | 52 +-- .../kernel/workspace-kernel-access.ts | 9 +- .../workspaces/kernel/workspace-kernel.ts | 33 +- .../search/workspace-search-batches.ts | 7 + .../search/workspace-search-chunks.ts | 18 +- .../search/workspace-search-content.ts | 64 ++-- .../search/workspace-search-contract.ts | 6 + .../search/workspace-search-embeddings.ts | 10 +- .../search/workspace-search-indexer.ts | 338 ++++++++++-------- .../search/workspace-search-projection.ts | 19 +- .../search/workspace-search-query.ts | 250 ++++++++----- .../search/workspace-search-references.ts | 1 + .../search/workspace-search-schema.ts | 53 ++- .../search/workspace-search-scope.ts | 192 ++++++++++ .../search/workspace-search-version.ts | 20 ++ .../search/workspace-search.test.ts | 173 ++++++++- wrangler.jsonc | 1 + 17 files changed, 906 insertions(+), 340 deletions(-) create mode 100644 src/features/workspaces/search/workspace-search-batches.ts create mode 100644 src/features/workspaces/search/workspace-search-scope.ts create mode 100644 src/features/workspaces/search/workspace-search-version.ts diff --git a/src/features/workspaces/extraction/workspace-page-projection.ts b/src/features/workspaces/extraction/workspace-page-projection.ts index 3edc6e0c..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, @@ -227,16 +223,12 @@ export async function* iterateWorkspacePageProjection(input: { const prefix = getManifestPrefix(input.manifestObjectKey); for (let pageNumber = 1; pageNumber <= manifest.pageCount; pageNumber += 1) { - const object = await input.bucket.get(getWorkspacePageObjectKey(prefix, pageNumber)); - if (!object) { - throw new Error(`Extracted page ${pageNumber} was not found.`); - } - - 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.`); - } + const object = await getWorkspacePageProjectionObject({ + bucket: input.bucket, + pageMetadataByNumber, + pageNumber, + prefix, + }); yield { markdown: await object.text(), @@ -245,6 +237,26 @@ export async function* iterateWorkspacePageProjection(input: { } } +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/kernel/workspace-kernel-access.ts b/src/features/workspaces/kernel/workspace-kernel-access.ts index edd46f8e..04976123 100644 --- a/src/features/workspaces/kernel/workspace-kernel-access.ts +++ b/src/features/workspaces/kernel/workspace-kernel-access.ts @@ -40,6 +40,7 @@ import type { WorkspaceSearchFailure, WorkspaceSearchInput, WorkspaceSearchResult, + WorkspaceSearchStatus, } from "#/features/workspaces/search/workspace-search-contract"; import { assertCanMutateWorkspace, @@ -119,9 +120,11 @@ export interface WorkspaceKernelClient { actorUserId?: string | null; clientMutationId?: string | null; }): Promise>; - searchWorkspace( - input: WorkspaceSearchInput, - ): Promise<{ failed: WorkspaceSearchFailure[]; results: WorkspaceSearchResult[] }>; + searchWorkspace(input: WorkspaceSearchInput): Promise<{ + failed: WorkspaceSearchFailure[]; + results: WorkspaceSearchResult[]; + status: WorkspaceSearchStatus; + }>; purgeForDeletion(): Promise; } diff --git a/src/features/workspaces/kernel/workspace-kernel.ts b/src/features/workspaces/kernel/workspace-kernel.ts index d855560e..fb724250 100644 --- a/src/features/workspaces/kernel/workspace-kernel.ts +++ b/src/features/workspaces/kernel/workspace-kernel.ts @@ -99,7 +99,20 @@ export class WorkspaceKernel extends Agent { workspaceId: () => this.name, getNextRevision: () => this.store.getNextRevision(), broadcast: (message) => this.broadcastRealtimeMessage(message), - onCommit: (event) => this.search.observe(event), + 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({ @@ -121,7 +134,7 @@ export class WorkspaceKernel extends Agent { async onStart() { initializeWorkspaceKernelStorage(this.kernelSql); this.search.initialize(); - if (this.search.hasRetryablePending()) { + if (this.search.hasPending()) { await this.scheduleWorkspaceSearchIndexing(); } } @@ -301,12 +314,17 @@ export class WorkspaceKernel extends Agent { } async searchWorkspace(input: WorkspaceSearchInput) { + if (this.search.hasPending()) { + this.ctx.waitUntil(this.scheduleWorkspaceSearchIndexing()); + } return await this.search.search(input); } async processWorkspaceSearchIndex() { if (await this.search.processBatch()) { - await this.scheduleWorkspaceSearchIndexing(); + // 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); } } @@ -391,9 +409,14 @@ export class WorkspaceKernel extends Agent { return { attempted: documentItemIds.length + 2, failed }; } - private async scheduleWorkspaceSearchIndexing() { + private async scheduleWorkspaceSearchIndexing(idempotent = true) { await this.schedule(1, "processWorkspaceSearchIndex", undefined, { - idempotent: true, + idempotent, + retry: { + baseDelayMs: 250, + maxAttempts: 5, + maxDelayMs: 3_000, + }, }); } 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 index 8ae942be..4cf77506 100644 --- a/src/features/workspaces/search/workspace-search-chunks.ts +++ b/src/features/workspaces/search/workspace-search-chunks.ts @@ -8,13 +8,16 @@ export interface WorkspaceSearchTextChunk { startLine: number; } -export function chunkWorkspaceSearchText(text: string): WorkspaceSearchTextChunk[] { +export function* iterateWorkspaceSearchTextChunks( + text: string, +): Generator { const normalized = text.replace(/\r\n?/g, "\n"); if (normalized.length === 0) { - return [{ content: "", endLine: 0, startLine: 0 }]; + yield { content: "", endLine: 0, startLine: 0 }; + return; } - const chunks: WorkspaceSearchTextChunk[] = []; + let yielded = false; let start = 0; let startLine = 1; let startLineCursor = 0; @@ -38,11 +41,12 @@ export function chunkWorkspaceSearchText(text: string): WorkspaceSearchTextChunk startLine += 1; startLineCursor = newline + 1; } - chunks.push({ + yielded = true; + yield { content, endLine: startLine + countLineBreaks(normalized, contentStart, contentEnd), startLine, - }); + }; } if (end >= normalized.length) { @@ -61,7 +65,9 @@ export function chunkWorkspaceSearchText(text: string): WorkspaceSearchTextChunk start = Math.min(boundary, end); } - return chunks.length > 0 ? chunks : [{ content: "", endLine: 0, startLine: 0 }]; + if (!yielded) { + yield { content: "", endLine: 0, startLine: 0 }; + } } function findChunkEnd(text: string, start: number, hardEnd: number) { diff --git a/src/features/workspaces/search/workspace-search-content.ts b/src/features/workspaces/search/workspace-search-content.ts index ff05a399..e618a343 100644 --- a/src/features/workspaces/search/workspace-search-content.ts +++ b/src/features/workspaces/search/workspace-search-content.ts @@ -1,9 +1,7 @@ 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 { chunkWorkspaceSearchText } from "#/features/workspaces/search/workspace-search-chunks"; - -const maximumIndexedCharactersPerItem = 8_000_000; +import { iterateWorkspaceSearchTextChunks } from "#/features/workspaces/search/workspace-search-chunks"; export interface WorkspaceSearchFileSystem { readFile(path: string): Promise; @@ -12,7 +10,6 @@ export interface WorkspaceSearchFileSystem { interface WorkspaceSearchIndexSourceBase { itemId: string; name: string; - path: string; sourceVersion: string; } @@ -30,65 +27,50 @@ export interface PreparedWorkspaceSearchChunk { startLine: number | null; } -export async function prepareWorkspaceSearchChunks(input: { +export async function* iteratePreparedWorkspaceSearchChunks(input: { bucket: R2Bucket; source: WorkspaceSearchIndexSource; workspace: WorkspaceSearchFileSystem; -}): Promise { +}): 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)); - const searchable = markdown.slice(0, maximumIndexedCharactersPerItem); - - return chunkWorkspaceSearchText(searchable).map((chunk, index) => ({ - content: chunk.content, - endLine: chunk.endLine, - index, - pageNumber: null, - startLine: chunk.startLine, - })); + for (const chunk of iterateWorkspaceSearchTextChunks(markdown)) { + yield { + content: chunk.content, + endLine: chunk.endLine, + index, + pageNumber: null, + startLine: chunk.startLine, + }; + index += 1; + } + return; } - const chunks: PreparedWorkspaceSearchChunk[] = []; - let indexedCharacters = 0; - for await (const page of iterateWorkspacePageProjection({ bucket: input.bucket, expectedSourceHash: input.source.sourceHash, manifestObjectKey: input.source.objectKey, })) { - const remaining = maximumIndexedCharactersPerItem - indexedCharacters; - if (remaining <= 0) { - break; - } - - const markdown = page.markdown.slice(0, remaining); - for (const chunk of chunkWorkspaceSearchText(markdown)) { - chunks.push({ + for (const chunk of iterateWorkspaceSearchTextChunks(page.markdown)) { + yield { content: chunk.content, endLine: null, - index: chunks.length, + index, pageNumber: page.pageNumber, startLine: null, - }); - } - indexedCharacters += markdown.length; - - if (markdown.length < page.markdown.length) { - break; + }; + index += 1; } } - - return chunks; } -export function createWorkspaceSearchEmbeddingText(input: { - content: string; - path: string; - title: string; -}) { - return [`Title: ${input.title}`, `Path: ${input.path}`, "", input.content].join("\n"); +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 index f838e4e0..02d68cc8 100644 --- a/src/features/workspaces/search/workspace-search-contract.ts +++ b/src/features/workspaces/search/workspace-search-contract.ts @@ -61,14 +61,20 @@ export const workspaceSearchFailureSchema = z.object({ path: z.string(), }); +export 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 index f9f21475..5790a609 100644 --- a/src/features/workspaces/search/workspace-search-embeddings.ts +++ b/src/features/workspaces/search/workspace-search-embeddings.ts @@ -1,3 +1,5 @@ +import { batchWorkspaceSearchValues } from "#/features/workspaces/search/workspace-search-batches"; + const workspaceSearchEmbeddingModel = "@cf/baai/bge-m3"; const embeddingBatchSize = 16; @@ -13,14 +15,6 @@ export async function embedWorkspaceSearchTexts(ai: Ai, texts: string[]) { return embeddings; } -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; -} - function readEmbeddingData(output: unknown): number[][] { if (!isRecord(output) || !Array.isArray(output.data)) { throw new Error("Workspace search embedding response is missing vector data."); diff --git a/src/features/workspaces/search/workspace-search-indexer.ts b/src/features/workspaces/search/workspace-search-indexer.ts index f6610f6e..0513fb08 100644 --- a/src/features/workspaces/search/workspace-search-indexer.ts +++ b/src/features/workspaces/search/workspace-search-indexer.ts @@ -1,27 +1,29 @@ -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 { batchWorkspaceSearchValues } from "#/features/workspaces/search/workspace-search-batches"; import { createWorkspaceSearchEmbeddingText, - prepareWorkspaceSearchChunks, + 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 { - batchWorkspaceSearchValues, - embedWorkspaceSearchTexts, -} from "#/features/workspaces/search/workspace-search-embeddings"; + buildWorkspaceSearchSourceVersion, + workspaceSearchIndexVersion, +} from "#/features/workspaces/search/workspace-search-version"; import { recordOperationalFailure } from "#/integrations/observability/operational-events"; import { sha256Base64UrlText } from "#/lib/binary"; -const searchIndexBatchSize = 4; -const vectorMutationBatchSize = 1_000; -const vectorDeleteBatchSize = 100; +const searchIndexBatchSize = 2; +const searchChunkProcessingBatchSize = 32; const maximumIndexAttempts = 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; @@ -30,6 +32,10 @@ interface SearchSourceRow { updated_at: number; } +type ScopedWorkspaceSearchIndexSource = WorkspaceSearchIndexSource & { + parentId: string | null; +}; + interface SearchIndexChunk { chunkId: string; content: string; @@ -42,7 +48,6 @@ interface SearchIndexChunk { export class WorkspaceSearchIndexer { private readonly ai: Ai; private readonly bucket: R2Bucket; - private readonly getItems: () => WorkspaceItemSummary[]; private readonly sql: WorkspaceKernelSql; private readonly vectorize: VectorizeIndex; private readonly workspace: WorkspaceSearchFileSystem; @@ -51,7 +56,6 @@ export class WorkspaceSearchIndexer { constructor(input: { ai: Ai; bucket: R2Bucket; - getItems: () => WorkspaceItemSummary[]; sql: WorkspaceKernelSql; vectorize: VectorizeIndex; workspace: WorkspaceSearchFileSystem; @@ -59,7 +63,6 @@ export class WorkspaceSearchIndexer { }) { this.ai = input.ai; this.bucket = input.bucket; - this.getItems = input.getItems; this.sql = input.sql; this.vectorize = input.vectorize; this.workspace = input.workspace; @@ -69,32 +72,33 @@ export class WorkspaceSearchIndexer { seedPendingItems() { const now = Date.now(); this.sql` - INSERT INTO kernel_search_pending (item_id, requested_at, attempts) - SELECT i.id, ${now}, 0 + 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 'document:' || i.updated_at - ELSE 'file:' || i.updated_at || ':' || p.updated_at || ':' || p.source_hash - END - OR s.vector_ready = 0 - ) + 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 + OR s.vector_status != 'ready' + ) ON CONFLICT(item_id) DO NOTHING `; this.sql` - INSERT INTO kernel_search_pending (item_id, requested_at, attempts) - SELECT s.item_id, ${now}, 0 + 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 @@ -104,41 +108,20 @@ export class WorkspaceSearchIndexer { markPending(itemId: string) { this.sql` - INSERT INTO kernel_search_pending (item_id, requested_at, attempts) - VALUES (${itemId}, ${Date.now()}, 0) + 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 `; } - markTreePending(itemId: string) { - this.sql` - WITH RECURSIVE search_tree(id) AS ( - SELECT ${itemId} - UNION ALL - SELECT i.id - FROM kernel_items i - JOIN search_tree parent ON i.parent_id = parent.id - WHERE i.deleted_at IS NULL - ) - INSERT INTO kernel_search_pending (item_id, requested_at, attempts) - SELECT i.id, ${Date.now()}, 0 - FROM kernel_items i - JOIN search_tree tree ON tree.id = i.id - WHERE i.deleted_at IS NULL AND i.type IN ('document', 'file') - ON CONFLICT(item_id) DO UPDATE SET - requested_at = excluded.requested_at, - attempts = 0 - `; - } - - hasRetryablePending() { + hasPending() { return Boolean( this.sql<{ item_id: string }>` - SELECT item_id - FROM kernel_search_pending - WHERE attempts < ${maximumIndexAttempts} + SELECT item_id FROM kernel_search_pending + UNION ALL + SELECT vector_id AS item_id FROM kernel_search_vector_deletes LIMIT 1 `[0], ); @@ -149,21 +132,27 @@ export class WorkspaceSearchIndexer { const pending = this.sql<{ item_id: string }>` SELECT item_id FROM kernel_search_pending - WHERE attempts < ${maximumIndexAttempts} 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(); - return this.hasRetryablePending(); + if (failures.length > 0) { + throw new AggregateError(failures, "Workspace search indexing failed."); + } + return this.hasPending(); } async purgeVectors() { @@ -187,71 +176,51 @@ export class WorkspaceSearchIndexer { return; } - const preparedChunks = await prepareWorkspaceSearchChunks({ - bucket: this.bucket, - source, - workspace: this.workspace, - }); + const revisionKey = await sha256Base64UrlText( + `${this.workspaceId()}:${source.itemId}:${source.sourceVersion}`, + ); if (!this.isCurrentSource(source)) { return; } + this.beginIndex(source); - const revisionKey = await sha256Base64UrlText( - `${this.workspaceId()}:${source.itemId}:${source.sourceVersion}`, - ); - const chunks: SearchIndexChunk[] = preparedChunks.map((chunk) => ({ - ...chunk, - chunkId: `s${revisionKey}-${chunk.index}`, - })); - this.replaceKeywordIndex({ - chunks, + let chunks: SearchIndexChunk[] = []; + for await (const prepared of iteratePreparedWorkspaceSearchChunks({ + bucket: this.bucket, source, - }); - - const embeddings = await embedWorkspaceSearchTexts( - this.ai, - chunks.map((chunk) => - createWorkspaceSearchEmbeddingText({ - content: chunk.content, - path: source.path, - title: source.name, - }), - ), - ); - if (embeddings.length !== chunks.length) { - throw new Error("Workspace search embedding response did not match the indexed chunks."); + workspace: this.workspace, + })) { + const chunk = { + ...prepared, + chunkId: `s${revisionKey}-${prepared.index}`, + }; + this.insertChunk(source, chunk); + chunks.push(chunk); + + if (chunks.length === searchChunkProcessingBatchSize) { + if (!(await this.indexChunkBatch(source, chunks))) { + return; + } + chunks = []; + } } - const vectors = chunks.map( - (chunk, index): VectorizeVector => ({ - id: chunk.chunkId, - namespace: this.workspaceId(), - values: embeddings[index] ?? [], - }), - ); - for (const batch of batchWorkspaceSearchValues(vectors, vectorMutationBatchSize)) { - await this.vectorize.upsert(batch); + if (chunks.length > 0 && !(await this.indexChunkBatch(source, chunks))) { + return; } - if (!this.isCurrentSource(source)) { - for (const vector of vectors) { - this.queueVectorDelete(vector.id); - } return; } - - this.markVectorIndexReady( - source, - chunks.map((chunk) => chunk.chunkId), - ); + this.markVectorIndexReady(source); } - private getIndexSource(itemId: string): WorkspaceSearchIndexSource | null { + 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, @@ -271,26 +240,28 @@ export class WorkspaceSearchIndexer { return null; } - const path = buildWorkspaceKernelItemPathIndex(this.getItems()).get(row.id); - if (!path) { - return null; - } - const source = { itemId: row.id, name: row.name, - path, + parentId: row.parent_id, }; if (row.type === "document") { return { ...source, shellPath: row.shell_path, - sourceVersion: `document:${row.updated_at}`, + sourceVersion: buildWorkspaceSearchSourceVersion({ + type: "document", + updatedAt: row.updated_at, + }), type: "document", }; } - if (!row.projection_object_key || !row.projection_source_hash) { + if ( + !row.projection_object_key || + !row.projection_source_hash || + row.projection_updated_at === null + ) { return null; } @@ -298,57 +269,101 @@ export class WorkspaceSearchIndexer { ...source, objectKey: row.projection_object_key, sourceHash: row.projection_source_hash, - sourceVersion: `file:${row.updated_at}:${row.projection_updated_at}:${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: WorkspaceSearchIndexSource) { + private isCurrentSource(source: ScopedWorkspaceSearchIndexSource) { const current = this.getIndexSource(source.itemId); return ( current?.sourceVersion === source.sourceVersion && current.name === source.name && - current.path === source.path + current.parentId === source.parentId ); } - private replaceKeywordIndex(input: { - chunks: SearchIndexChunk[]; - source: WorkspaceSearchIndexSource; - }) { - const retainedChunkIds = new Set(input.chunks.map((chunk) => chunk.chunkId)); + private beginIndex(source: ScopedWorkspaceSearchIndexSource) { for (const row of this.sql<{ chunk_id: string }>` SELECT chunk_id FROM kernel_search_chunks - WHERE item_id = ${input.source.itemId} + WHERE item_id = ${source.itemId} `) { - if (!retainedChunkIds.has(row.chunk_id)) { - this.queueVectorDelete(row.chunk_id); - } + this.queueVectorDelete(row.chunk_id); } - this.deleteLocalChunks(input.source.itemId); - for (const chunk of input.chunks) { - this.insertChunk(input.source, chunk); - } + this.deleteLocalChunks(source.itemId); this.sql` INSERT INTO kernel_search_items ( item_id, source_version, - vector_ready + vector_status ) VALUES ( - ${input.source.itemId}, - ${input.source.sourceVersion}, - 0 + ${source.itemId}, + ${source.sourceVersion}, + 'pending' ) ON CONFLICT(item_id) DO UPDATE SET source_version = excluded.source_version, - vector_ready = 0 + vector_status = 'pending' `; } - private insertChunk(source: WorkspaceSearchIndexSource, chunk: SearchIndexChunk) { + 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 => ({ + id: chunk.chunkId, + metadata: createWorkspaceSearchVectorMetadata({ + itemId: source.itemId, + parentId: source.parentId, + type: source.type, + }), + namespace: this.workspaceId(), + values: embeddings[index] ?? [], + }), + ); + 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, @@ -366,26 +381,18 @@ export class WorkspaceSearchIndexer { ) `; this.sql` - INSERT INTO kernel_search_fts (chunk_id, title, path, content) - VALUES (${chunk.chunkId}, ${source.name}, ${source.path}, ${chunk.content}) + INSERT INTO kernel_search_fts (chunk_id, title, content) + VALUES (${chunk.chunkId}, ${source.name}, ${chunk.content}) `; } - private markVectorIndexReady(source: WorkspaceSearchIndexSource, vectorIds: readonly string[]) { + private markVectorIndexReady(source: ScopedWorkspaceSearchIndexSource) { this.sql` UPDATE kernel_search_items - SET vector_ready = 1 + 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}`; - if (vectorIds.length > 0) { - this.sql` - DELETE FROM kernel_search_vector_deletes - WHERE vector_id IN ( - SELECT value FROM json_each(${JSON.stringify(vectorIds)}) - ) - `; - } } private removeIndexedItem(itemId: string) { @@ -424,6 +431,18 @@ export class WorkspaceSearchIndexer { `; } + 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 @@ -451,11 +470,6 @@ export class WorkspaceSearchIndexer { } private recordIndexFailure(itemId: string, error: unknown) { - this.sql` - UPDATE kernel_search_pending - SET attempts = attempts + 1 - WHERE item_id = ${itemId} - `; recordOperationalFailure({ error, event: "workspace_search_indexing", @@ -465,4 +479,24 @@ export class WorkspaceSearchIndexer { }, }); } + + 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 index d8ff7242..105d0585 100644 --- a/src/features/workspaces/search/workspace-search-projection.ts +++ b/src/features/workspaces/search/workspace-search-projection.ts @@ -45,12 +45,21 @@ export class WorkspaceSearchProjection { 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": - this.indexer.markTreePending(event.payload.item.id); + // 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) { - this.indexer.markTreePending(item.id); + if (item.type === "document" || item.type === "file") { + this.indexer.markPending(item.id); + } } break; case "workspace.item.projection.updated": @@ -68,13 +77,13 @@ export class WorkspaceSearchProjection { return; } - if (this.indexer.hasRetryablePending()) { + if (this.indexer.hasPending()) { this.requestRun(); } } - hasRetryablePending() { - return this.indexer.hasRetryablePending(); + hasPending() { + return this.indexer.hasPending(); } async processBatch() { diff --git a/src/features/workspaces/search/workspace-search-query.ts b/src/features/workspaces/search/workspace-search-query.ts index ee37ed7d..56d9b25e 100644 --- a/src/features/workspaces/search/workspace-search-query.ts +++ b/src/features/workspaces/search/workspace-search-query.ts @@ -1,12 +1,5 @@ import type { WorkspaceItemSummary } from "#/features/workspaces/contracts"; -import { - buildWorkspaceKernelItemPathIndex, - buildWorkspaceKernelTree, - normalizeWorkspacePath, - resolveWorkspaceKernelItemPath, - WorkspaceKernelPathError, - type WorkspaceKernelTree, -} from "#/features/workspaces/kernel/workspace-kernel-paths"; +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 { @@ -14,14 +7,25 @@ import type { 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 { workspaceSearchIndexVersion } from "#/features/workspaces/search/workspace-search-version"; import { recordOperationalFailure } from "#/integrations/observability/operational-events"; +const semanticQueryConcurrency = 4; + interface SearchChunkRow { chunk_id: string; content: string; @@ -59,36 +63,44 @@ export class WorkspaceSearchQuery { this.workspaceId = input.workspaceId; } - async search( - input: WorkspaceSearchInput, - ): Promise<{ failed: WorkspaceSearchFailure[]; results: WorkspaceSearchResult[] }> { - const scope = this.resolveScope(input.path ?? "/"); - if (scope.status === "failed") { - return { failed: [scope.failure], results: [] }; - } - + 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, - scopeItemIds: scope.itemIds, + scope: scope.local, types, }); const semantic = await this.searchSemanticWithFallback({ candidateLimit, query: input.query, - scopeItemIds: scope.itemIds, + scope, types, }); - const items = this.getItems(); const itemsById = new Map(items.map((item) => [item.id, item])); const paths = buildWorkspaceKernelItemPathIndex(items); const ranked = fuseWorkspaceSearchRanks({ keyword, limit, semantic }); return { failed: [], + status, results: ranked.flatMap((candidate) => { const path = paths.get(candidate.itemId); const item = itemsById.get(candidate.itemId); @@ -101,10 +113,45 @@ export class WorkspaceSearchQuery { }; } + 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; - scopeItemIds: string[] | null; + scope: WorkspaceSearchScope; types: WorkspaceSearchItemType[]; }) { try { @@ -119,48 +166,10 @@ export class WorkspaceSearchQuery { } } - private resolveScope( - requestedPath: string, - ): - | { itemIds: string[] | null; status: "ready" } - | { failure: WorkspaceSearchFailure; status: "failed" } { - let path: string; - try { - path = normalizeWorkspacePath(requestedPath); - } catch (error) { - if (error instanceof WorkspaceKernelPathError && error.code === "path_not_absolute") { - return { - failure: { code: error.code, path: requestedPath }, - status: "failed", - }; - } - throw error; - } - - if (path === "/") { - return { itemIds: null, status: "ready" }; - } - - const items = this.getItems(); - const tree = buildWorkspaceKernelTree(items); - const item = resolveWorkspaceKernelItemPath(path, tree); - if (!item) { - return { - failure: { code: "path_not_found", path }, - status: "failed", - }; - } - if (item.type !== "folder") { - return { itemIds: [item.id], status: "ready" }; - } - - return { itemIds: listDescendantItemIds(item.id, tree), status: "ready" }; - } - private searchKeyword(input: { candidateLimit: number; query: string; - scopeItemIds: string[] | null; + scope: WorkspaceSearchLocalScope; types: WorkspaceSearchItemType[]; }): SearchCandidate[] { const match = createFtsMatchExpression(input.query); @@ -168,7 +177,7 @@ export class WorkspaceSearchQuery { return []; } - const scopeJson = input.scopeItemIds ? JSON.stringify(input.scopeItemIds) : null; + const localScope = serializeWorkspaceSearchLocalScope(input.scope); const typesJson = JSON.stringify(input.types); const rows = this.sql` SELECT @@ -188,14 +197,23 @@ export class WorkspaceSearchQuery { AND p.status = 'ready' WHERE kernel_search_fts MATCH ${match} AND s.source_version = CASE - WHEN i.type = 'document' THEN 'document:' || i.updated_at - ELSE 'file:' || i.updated_at || ':' || p.updated_at || ':' || p.source_hash + WHEN i.type = 'document' + THEN ${workspaceSearchIndexVersion} || ':document:' || i.updated_at + ELSE ${workspaceSearchIndexVersion} || ':file:' || i.updated_at || ':' || p.updated_at || ':' || p.source_hash END - AND (${scopeJson} IS NULL OR c.item_id IN ( - SELECT value FROM json_each(${scopeJson}) - )) + 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, 0.0, 8.0, 4.0, 1.0) ASC + ORDER BY bm25(kernel_search_fts, 0.0, 8.0, 1.0) ASC LIMIT ${input.candidateLimit} `; @@ -205,26 +223,29 @@ export class WorkspaceSearchQuery { private async searchSemantic(input: { candidateLimit: number; query: string; - scopeItemIds: string[] | null; + scope: WorkspaceSearchScope; types: WorkspaceSearchItemType[]; }): Promise { + if (input.scope.vectorFilters.length === 0) { + return []; + } const [embedding] = await embedWorkspaceSearchTexts(this.ai, [input.query]); if (!embedding) { return []; } - const matches = await this.vectorize.query(embedding, { - namespace: this.workspaceId(), - returnMetadata: "none", - topK: input.candidateLimit, + const matches = await this.querySemanticMatches({ + candidateLimit: input.candidateLimit, + embedding, + filters: input.scope.vectorFilters, }); - const vectorIds = matches.matches.map((match) => match.id); + const vectorIds = matches.map((match) => match.id); if (vectorIds.length === 0) { return []; } const rows = this.loadSemanticCandidates({ - scopeItemIds: input.scopeItemIds, + scope: input.scope.local, types: input.types, vectorIds, }); @@ -236,12 +257,48 @@ export class WorkspaceSearchQuery { }); } + private async querySemanticMatches(input: { + candidateLimit: number; + embedding: number[]; + filters: WorkspaceSearchVectorFilter[]; + }) { + let bestMatches = new Map(); + + for (const filterBatch of batchWorkspaceSearchValues(input.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 Array.from(bestMatches.values()).sort(compareSemanticMatches); + } + private loadSemanticCandidates(input: { - scopeItemIds: string[] | null; + scope: WorkspaceSearchLocalScope; types: WorkspaceSearchItemType[]; vectorIds: string[]; }) { - const scopeJson = input.scopeItemIds ? JSON.stringify(input.scopeItemIds) : null; + const localScope = serializeWorkspaceSearchLocalScope(input.scope); const typesJson = JSON.stringify(input.types); return this.sql` @@ -255,35 +312,44 @@ export class WorkspaceSearchQuery { FROM kernel_search_chunks c JOIN kernel_search_fts ON kernel_search_fts.chunk_id = c.chunk_id 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_ready = 1 + JOIN kernel_search_items s ON s.item_id = c.item_id AND s.vector_status = 'ready' LEFT JOIN kernel_item_projections p ON p.item_id = i.id AND p.format = 'pages' AND p.status = 'ready' WHERE c.chunk_id IN (SELECT value FROM json_each(${JSON.stringify(input.vectorIds)})) AND s.source_version = CASE - WHEN i.type = 'document' THEN 'document:' || i.updated_at - ELSE 'file:' || i.updated_at || ':' || p.updated_at || ':' || p.source_hash + WHEN i.type = 'document' + THEN ${workspaceSearchIndexVersion} || ':document:' || i.updated_at + ELSE ${workspaceSearchIndexVersion} || ':file:' || i.updated_at || ':' || p.updated_at || ':' || p.source_hash END - AND (${scopeJson} IS NULL OR c.item_id IN ( - SELECT value FROM json_each(${scopeJson}) - )) + 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 listDescendantItemIds(folderId: string, tree: WorkspaceKernelTree) { - const ids: string[] = []; - const pending = [folderId]; - while (pending.length > 0) { - const itemId = pending.pop(); - if (itemId) { - ids.push(itemId); - pending.push(...(tree.childrenByParentId.get(itemId) ?? []).map((item) => item.id)); - } - } - return ids; +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 { diff --git a/src/features/workspaces/search/workspace-search-references.ts b/src/features/workspaces/search/workspace-search-references.ts index 4f9d1ff8..0a789a6e 100644 --- a/src/features/workspaces/search/workspace-search-references.ts +++ b/src/features/workspaces/search/workspace-search-references.ts @@ -19,6 +19,7 @@ export function createWorkspaceSearchModelOutput(output: WorkspaceSearchOutput) return { failed: output.failed, + status: output.status, results: output.results.map((result) => { const { itemId: _itemId, ...modelResult } = result; const reference = refsByLocation.get( diff --git a/src/features/workspaces/search/workspace-search-schema.ts b/src/features/workspaces/search/workspace-search-schema.ts index 8ddd9d61..79e74b8f 100644 --- a/src/features/workspaces/search/workspace-search-schema.ts +++ b/src/features/workspaces/search/workspace-search-schema.ts @@ -1,11 +1,53 @@ import type { WorkspaceKernelSql } from "#/features/workspaces/kernel/workspace-kernel-schema"; +const workspaceSearchStorageVersion = 1; + export function initializeWorkspaceSearchStorage(sql: WorkspaceKernelSql) { + sql` + CREATE TABLE IF NOT EXISTS kernel_search_metadata ( + key TEXT PRIMARY KEY, + value INTEGER NOT NULL + ) + `; + sql` + CREATE TABLE IF NOT EXISTS kernel_search_vector_deletes ( + vector_id TEXT PRIMARY KEY, + requested_at INTEGER NOT NULL + ) + `; + + const storedVersion = sql<{ value: number }>` + SELECT value + FROM kernel_search_metadata + WHERE key = 'storage_version' + LIMIT 1 + `[0]?.value; + if (storedVersion !== workspaceSearchStorageVersion) { + // 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_ready INTEGER NOT NULL DEFAULT 0 + vector_status TEXT NOT NULL DEFAULT 'pending' ) `; sql` @@ -23,7 +65,6 @@ export function initializeWorkspaceSearchStorage(sql: WorkspaceKernelSql) { CREATE VIRTUAL TABLE IF NOT EXISTS kernel_search_fts USING fts5( chunk_id UNINDEXED, title, - path, content, tokenize = 'unicode61 remove_diacritics 2' ) @@ -35,10 +76,10 @@ export function initializeWorkspaceSearchStorage(sql: WorkspaceKernelSql) { attempts INTEGER NOT NULL DEFAULT 0 ) `; + sql` - CREATE TABLE IF NOT EXISTS kernel_search_vector_deletes ( - vector_id TEXT PRIMARY KEY, - requested_at INTEGER NOT NULL - ) + INSERT INTO kernel_search_metadata (key, value) + VALUES ('storage_version', ${workspaceSearchStorageVersion}) + ON CONFLICT(key) DO UPDATE SET value = excluded.value `; } 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..86fa714b --- /dev/null +++ b/src/features/workspaces/search/workspace-search-version.ts @@ -0,0 +1,20 @@ +/** + * Bump whenever the embedding model, chunking, or embedding text format changes. + * The SQL freshness checks in the indexer and query 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 index e27add4d..75d1f70f 100644 --- a/src/features/workspaces/search/workspace-search.test.ts +++ b/src/features/workspaces/search/workspace-search.test.ts @@ -1,7 +1,14 @@ import { describe, expect, it } from "vitest"; -import { chunkWorkspaceSearchText } from "#/features/workspaces/search/workspace-search-chunks"; +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 { 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", () => { @@ -9,7 +16,7 @@ describe("workspace search", () => { { length: 80 }, (_, index) => `Line ${index + 1}: ${"searchable content ".repeat(4)}`, ).join("\n"); - const chunks = chunkWorkspaceSearchText(text); + const chunks = Array.from(iterateWorkspaceSearchTextChunks(text)); expect(chunks.length).toBeGreaterThan(2); expect(chunks[0]).toMatchObject({ startLine: 1 }); @@ -18,6 +25,56 @@ describe("workspace search", () => { expect(chunks[1]?.startLine).toBeLessThanOrEqual((chunks[0]?.endLine ?? 0) + 1); }); + 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({ @@ -36,4 +93,116 @@ describe("workspace search", () => { 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/wrangler.jsonc b/wrangler.jsonc index 14e5ffc0..8029c6df 100644 --- a/wrangler.jsonc +++ b/wrangler.jsonc @@ -34,6 +34,7 @@ "binding": "AI", "remote": true, }, + // Both indexes require string metadata indexes for itemId, parentId, and type. "vectorize": [ { "binding": "WORKSPACE_SEARCH", From bef73b8bdf87c90f831256e9fefcc8f92eee6eac Mon Sep 17 00:00:00 2001 From: Urjit Chakraborty <135136842+urjitc@users.noreply.github.com> Date: Thu, 30 Jul 2026 22:00:28 -0400 Subject: [PATCH 03/13] fix(workspaces): heal stale file extractions --- .../request-workspace-file-extraction.ts | 19 +--- .../workspace-file-extraction-reconciler.ts | 102 ++++++++++++++++++ .../workspace-file-extraction-workflow-id.ts | 17 +++ .../workspace-projection-readiness.ts | 4 +- .../workspaces/kernel/workspace-kernel.ts | 20 ++++ 5 files changed, 146 insertions(+), 16 deletions(-) create mode 100644 src/features/workspaces/extraction/workspace-file-extraction-reconciler.ts create mode 100644 src/features/workspaces/extraction/workspace-file-extraction-workflow-id.ts 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..7810a099 --- /dev/null +++ b/src/features/workspaces/extraction/workspace-file-extraction-reconciler.ts @@ -0,0 +1,102 @@ +import type { WorkspaceItemSummary } from "#/features/workspaces/contracts"; +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 { resolveWorkspaceFileTypeFromItem } from "#/features/workspaces/model/workspace-file"; + +const extractionHealingVersion = "extraction-healing-v1"; +const failedExtractionCooldownMs = 15 * 60_000; +const missingProjectionGraceMs = workspaceExtractionStallThresholdMs; +const workflowBatchSize = 100; + +export async function reconcileWorkspaceFileExtractions(input: { + items: readonly WorkspaceItemSummary[]; + sql: WorkspaceKernelSql; + workflow: Workflow; + workspaceId: string; +}) { + const now = Date.now(); + const candidates = input.sql<{ + id: string; + object_key: string; + projection_status: string | null; + projection_updated_at: number | null; + }>` + SELECT + i.id, + i.object_key, + p.status AS projection_status, + 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' + 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 - missingProjectionGraceMs}) + 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 - missingProjectionGraceMs} + ) + ) + ORDER BY i.created_at ASC + LIMIT ${workflowBatchSize} + `; + const itemsById = new Map(input.items.map((item) => [item.id, item])); + const workflows: Array<{ + id: string; + params: WorkspaceFileExtractionWorkflowParams; + }> = []; + + for (const candidate of candidates) { + const item = itemsById.get(candidate.id); + if (!item) { + continue; + } + + const fileType = resolveWorkspaceFileTypeFromItem(item); + if (!fileType) { + continue; + } + + const runKey = [ + extractionHealingVersion, + candidate.object_key, + candidate.projection_status ?? "missing", + candidate.projection_updated_at ?? 0, + ].join(":"); + const params = { + actorUserId: null, + assetKind: fileType.assetKind, + itemId: item.id, + requestId: extractionHealingVersion, + workspaceId: input.workspaceId, + } satisfies WorkspaceFileExtractionWorkflowParams; + workflows.push({ + id: await getWorkspaceFileExtractionWorkflowId({ + assetKind: params.assetKind, + itemId: item.id, + runKey, + workspaceId: input.workspaceId, + }), + params, + }); + } + + if (workflows.length === 0) { + return; + } + + 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-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.ts b/src/features/workspaces/kernel/workspace-kernel.ts index fb724250..e69d0986 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"; @@ -137,6 +138,7 @@ export class WorkspaceKernel extends Agent { if (this.search.hasPending()) { await this.scheduleWorkspaceSearchIndexing(); } + this.requestWorkspaceFileExtractionHealing(); } onConnect(connection: Connection, context: ConnectionContext) { @@ -150,6 +152,7 @@ export class WorkspaceKernel extends Agent { connection.setState({ user, }); + this.requestWorkspaceFileExtractionHealing(); this.broadcastPresenceSnapshot(); } @@ -420,6 +423,23 @@ export class WorkspaceKernel extends Agent { }); } + private requestWorkspaceFileExtractionHealing() { + this.ctx.waitUntil( + reconcileWorkspaceFileExtractions({ + items: this.store.getPageItems(), + 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() { this.broadcastRealtimeMessage({ type: "presence.snapshot", From 73ebe6491e94ba4fc1a2a5740218e275c00ef0a4 Mon Sep 17 00:00:00 2001 From: Urjit Chakraborty <135136842+urjitc@users.noreply.github.com> Date: Thu, 30 Jul 2026 22:30:56 -0400 Subject: [PATCH 04/13] fix(workspace-search): make indexing cleanup recoverable Persist cleanup progress across partial Vectorize failures and bound background deletion retries. Remove abandoned partial chunks and retry failed workspace purges without erasing their recovery inventory. Keep document source revisions monotonic so indexing cannot accept stale checkpoint content. --- .../kernel/workspace-kernel-item-commands.ts | 7 +- .../workspaces/kernel/workspace-kernel.ts | 59 +++++++++--- .../search/workspace-search-indexer.ts | 94 +++++++++++++------ .../search/workspace-search-schema.ts | 7 +- 4 files changed, 119 insertions(+), 48 deletions(-) diff --git a/src/features/workspaces/kernel/workspace-kernel-item-commands.ts b/src/features/workspaces/kernel/workspace-kernel-item-commands.ts index 8f38de7c..e8247ca5 100644 --- a/src/features/workspaces/kernel/workspace-kernel-item-commands.ts +++ b/src/features/workspaces/kernel/workspace-kernel-item-commands.ts @@ -353,14 +353,15 @@ 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); persistDocumentItemContentUpdate({ content: input.content, itemId: input.itemId, - metadataJson: item.metadata_json, + metadataJson: currentItem.metadata_json, sql: this.sql, - updatedAt: now, + updatedAt, }); return this.commitItemEvent({ diff --git a/src/features/workspaces/kernel/workspace-kernel.ts b/src/features/workspaces/kernel/workspace-kernel.ts index e69d0986..5c8b6b56 100644 --- a/src/features/workspaces/kernel/workspace-kernel.ts +++ b/src/features/workspaces/kernel/workspace-kernel.ts @@ -68,10 +68,13 @@ import { WorkspaceSearchProjection } from "#/features/workspaces/search/workspac 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({ @@ -362,7 +365,7 @@ 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; @@ -397,18 +400,43 @@ export class WorkspaceKernel extends Agent { } } - await Promise.all([ - deleteR2Prefix( - this.env.WORKSPACE_KERNEL_FILES, - getChatAttachmentWorkspacePrefix(workspaceId), - ), - deleteR2Prefix(this.env.WORKSPACE_KERNEL_FILES, `uploads/workspaces/${workspaceId}/`), - deleteR2Prefix(this.env.WORKSPACE_KERNEL_FILES, `workspace_kernel_files/${workspaceId}/`), - deleteR2Prefix(this.env.WORKSPACE_KERNEL_FILES, `workspace_file_objects/${workspaceId}/`), - deleteR2Prefix(this.env.WORKSPACE_KERNEL_FILES, `workspace_file_uploads/${workspaceId}/`), - ]); - - await this.ctx.storage.deleteAll(); + try { + await Promise.all([ + deleteR2Prefix( + this.env.WORKSPACE_KERNEL_FILES, + getChatAttachmentWorkspacePrefix(workspaceId), + ), + deleteR2Prefix(this.env.WORKSPACE_KERNEL_FILES, `uploads/workspaces/${workspaceId}/`), + deleteR2Prefix(this.env.WORKSPACE_KERNEL_FILES, `workspace_kernel_files/${workspaceId}/`), + deleteR2Prefix(this.env.WORKSPACE_KERNEL_FILES, `workspace_file_objects/${workspaceId}/`), + deleteR2Prefix(this.env.WORKSPACE_KERNEL_FILES, `workspace_file_uploads/${workspaceId}/`), + ]); + } catch (error) { + failed += 1; + recordOperationalFailure({ + error, + event: "workspace_r2_purge", + fields: { workspace_id: workspaceId }, + }); + } + + // Keep the local inventory when a remote purge fails so cleanup can be retried. + if (failed === 0) { + await this.ctx.storage.deleteAll(); + } 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 }; } @@ -424,6 +452,11 @@ export class WorkspaceKernel extends Agent { } private requestWorkspaceFileExtractionHealing() { + const now = Date.now(); + if (now - this.lastExtractionHealingRequestAt < workspaceExtractionHealingThrottleMs) { + return; + } + this.lastExtractionHealingRequestAt = now; this.ctx.waitUntil( reconcileWorkspaceFileExtractions({ items: this.store.getPageItems(), diff --git a/src/features/workspaces/search/workspace-search-indexer.ts b/src/features/workspaces/search/workspace-search-indexer.ts index 0513fb08..6329c565 100644 --- a/src/features/workspaces/search/workspace-search-indexer.ts +++ b/src/features/workspaces/search/workspace-search-indexer.ts @@ -18,6 +18,7 @@ import { sha256Base64UrlText } from "#/lib/binary"; const searchIndexBatchSize = 2; const searchChunkProcessingBatchSize = 32; const maximumIndexAttempts = 5; +const maximumVectorDeleteAttempts = 5; const vectorDeleteBatchSize = 100; interface SearchSourceRow { @@ -121,7 +122,9 @@ export class WorkspaceSearchIndexer { 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 + SELECT vector_id AS item_id + FROM kernel_search_vector_deletes + WHERE attempts < ${maximumVectorDeleteAttempts} LIMIT 1 `[0], ); @@ -164,8 +167,17 @@ export class WorkspaceSearchIndexer { `.map((row) => row.vector_id), ); + const failures: unknown[] = []; for (const batch of batchWorkspaceSearchValues(Array.from(ids), vectorDeleteBatchSize)) { - await this.vectorize.deleteByIds(batch); + 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."); } } @@ -289,15 +301,7 @@ export class WorkspaceSearchIndexer { } private beginIndex(source: ScopedWorkspaceSearchIndexSource) { - for (const row of this.sql<{ chunk_id: string }>` - SELECT chunk_id - FROM kernel_search_chunks - WHERE item_id = ${source.itemId} - `) { - this.queueVectorDelete(row.chunk_id); - } - - this.deleteLocalChunks(source.itemId); + this.discardItemChunks(source.itemId); this.sql` INSERT INTO kernel_search_items ( item_id, @@ -339,8 +343,12 @@ export class WorkspaceSearchIndexer { return false; } - const vectors = chunks.map( - (chunk, index): VectorizeVector => ({ + 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, @@ -348,9 +356,9 @@ export class WorkspaceSearchIndexer { type: source.type, }), namespace: this.workspaceId(), - values: embeddings[index] ?? [], - }), - ); + values, + }; + }); await this.vectorize.upsert(vectors); if (!this.isCurrentSource(source)) { for (const vector of vectors) { @@ -381,8 +389,10 @@ export class WorkspaceSearchIndexer { ) `; this.sql` - INSERT INTO kernel_search_fts (chunk_id, title, content) - VALUES (${chunk.chunkId}, ${source.name}, ${chunk.content}) + INSERT INTO kernel_search_fts (rowid, chunk_id, title, content) + SELECT rowid, chunk_id, ${source.name}, ${chunk.content} + FROM kernel_search_chunks + WHERE chunk_id = ${chunk.chunkId} `; } @@ -396,6 +406,12 @@ export class WorkspaceSearchIndexer { } 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) { for (const row of this.sql<{ chunk_id: string }>` SELECT chunk_id FROM kernel_search_chunks @@ -404,25 +420,35 @@ export class WorkspaceSearchIndexer { this.queueVectorDelete(row.chunk_id); } this.deleteLocalChunks(itemId); - this.sql`DELETE FROM kernel_search_items WHERE item_id = ${itemId}`; - this.sql`DELETE FROM kernel_search_pending WHERE item_id = ${itemId}`; } private deleteLocalChunks(itemId: string) { - const chunkIds = this.sql<{ chunk_id: string }>` - SELECT chunk_id - FROM kernel_search_chunks - WHERE item_id = ${itemId} - `.map((row) => row.chunk_id); - if (chunkIds.length > 0) { - this.sql` - DELETE FROM kernel_search_fts - WHERE chunk_id IN (SELECT value FROM json_each(${JSON.stringify(chunkIds)})) - `; - } + 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) @@ -461,6 +487,11 @@ export class WorkspaceSearchIndexer { 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", @@ -491,6 +522,7 @@ export class WorkspaceSearchIndexer { return false; } + this.discardItemChunks(itemId); this.sql` UPDATE kernel_search_items SET vector_status = 'failed' diff --git a/src/features/workspaces/search/workspace-search-schema.ts b/src/features/workspaces/search/workspace-search-schema.ts index 79e74b8f..ff21b6b9 100644 --- a/src/features/workspaces/search/workspace-search-schema.ts +++ b/src/features/workspaces/search/workspace-search-schema.ts @@ -9,10 +9,13 @@ export function initializeWorkspaceSearchStorage(sql: WorkspaceKernelSql) { value INTEGER NOT NULL ) `; + sql`CREATE INDEX IF NOT EXISTS kernel_search_vector_deletes_pending_idx + ON kernel_search_vector_deletes (attempts, requested_at)`; sql` CREATE TABLE IF NOT EXISTS kernel_search_vector_deletes ( vector_id TEXT PRIMARY KEY, - requested_at INTEGER NOT NULL + requested_at INTEGER NOT NULL, + attempts INTEGER NOT NULL DEFAULT 0 ) `; @@ -76,6 +79,8 @@ export function initializeWorkspaceSearchStorage(sql: WorkspaceKernelSql) { 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_search_metadata (key, value) From 9f60c29cb47fdecbc354cdbbf67dad7f63bbb348 Mon Sep 17 00:00:00 2001 From: Urjit Chakraborty <135136842+urjitc@users.noreply.github.com> Date: Thu, 30 Jul 2026 22:31:06 -0400 Subject: [PATCH 05/13] fix(workspace-search): bound and clarify retrieval Cap semantic calls for very large folder scopes and report reduced semantic coverage as partial. Keep stale or incomplete chunks out of results, preserve type filters for repeated values, center excerpts on any matching term, and cover chunk bounds plus model-facing references. --- .../workspace-operation-observability.ts | 7 -- .../operations/workspace-tool-definitions.ts | 7 +- .../search/workspace-search-chunks.ts | 2 +- .../search/workspace-search-projection.ts | 3 + .../search/workspace-search-query.ts | 84 +++++++++++-------- .../workspace-search-references.test.ts | 55 ++++++++++++ .../search/workspace-search-scope.ts | 3 +- .../search/workspace-search.test.ts | 21 +++++ 8 files changed, 135 insertions(+), 47 deletions(-) create mode 100644 src/features/workspaces/search/workspace-search-references.test.ts diff --git a/src/features/workspaces/operations/workspace-operation-observability.ts b/src/features/workspaces/operations/workspace-operation-observability.ts index eaa880d9..3b5a62db 100644 --- a/src/features/workspaces/operations/workspace-operation-observability.ts +++ b/src/features/workspaces/operations/workspace-operation-observability.ts @@ -73,13 +73,6 @@ export function summarizeWorkspaceReadResult(input: { return summarizeWorkspaceResult(succeededCount, failures, pendingCount); } -export function summarizeWorkspaceSearchResult(input: { - failed: ReadonlyArray<{ code: string }>; - results: readonly unknown[]; -}) { - return summarizeWorkspaceResult(input.results.length, input.failed); -} - export function summarizeWorkspaceItemResult(input: { failed: ReadonlyArray<{ code: string }>; item?: unknown; diff --git a/src/features/workspaces/operations/workspace-tool-definitions.ts b/src/features/workspaces/operations/workspace-tool-definitions.ts index 865e7124..ee07d811 100644 --- a/src/features/workspaces/operations/workspace-tool-definitions.ts +++ b/src/features/workspaces/operations/workspace-tool-definitions.ts @@ -50,7 +50,6 @@ import { summarizeWorkspaceCollectionResult, summarizeWorkspaceItemResult, summarizeWorkspaceReadResult, - summarizeWorkspaceSearchResult, type WorkspaceOperationSummary, } from "#/features/workspaces/operations/workspace-operation-observability"; @@ -169,7 +168,11 @@ export const workspaceToolDefinitions = [ inputSchema: workspaceSearchInputSchema, inputExamples: workspaceSearchInputExamples, outputSchema: workspaceSearchOutputSchema, - summarizeResult: summarizeWorkspaceSearchResult, + summarizeResult: (result) => + summarizeWorkspaceCollectionResult({ + failed: result.failed, + items: result.results, + }), effects: { destructive: false, idempotent: true }, execute: async (args, context) => { return await searchWorkspaceOperation(context, args); diff --git a/src/features/workspaces/search/workspace-search-chunks.ts b/src/features/workspaces/search/workspace-search-chunks.ts index 4cf77506..384c2eeb 100644 --- a/src/features/workspaces/search/workspace-search-chunks.ts +++ b/src/features/workspaces/search/workspace-search-chunks.ts @@ -77,7 +77,7 @@ function findChunkEnd(text: string, start: number, hardEnd: number) { const minimumEnd = start + minimumChunkCharacters; for (const separator of ["\n\n", "\n", ". "]) { - const boundary = text.lastIndexOf(separator, hardEnd); + const boundary = text.lastIndexOf(separator, hardEnd - separator.length); if (boundary >= minimumEnd) { return boundary + separator.length; } diff --git a/src/features/workspaces/search/workspace-search-projection.ts b/src/features/workspaces/search/workspace-search-projection.ts index 105d0585..d49c3716 100644 --- a/src/features/workspaces/search/workspace-search-projection.ts +++ b/src/features/workspaces/search/workspace-search-projection.ts @@ -75,6 +75,9 @@ export class WorkspaceSearchProjection { case "workspace.item.color.updated": case "workspace.relations.updated": return; + default: + event satisfies never; + return; } if (this.indexer.hasPending()) { diff --git a/src/features/workspaces/search/workspace-search-query.ts b/src/features/workspaces/search/workspace-search-query.ts index 56d9b25e..56e3890d 100644 --- a/src/features/workspaces/search/workspace-search-query.ts +++ b/src/features/workspaces/search/workspace-search-query.ts @@ -21,10 +21,10 @@ import { type WorkspaceSearchScope, type WorkspaceSearchVectorFilter, } from "#/features/workspaces/search/workspace-search-scope"; -import { workspaceSearchIndexVersion } from "#/features/workspaces/search/workspace-search-version"; import { recordOperationalFailure } from "#/integrations/observability/operational-events"; const semanticQueryConcurrency = 4; +const maximumSemanticQueries = 16; interface SearchChunkRow { chunk_id: string; @@ -96,11 +96,15 @@ export class WorkspaceSearchQuery { }); const itemsById = new Map(items.map((item) => [item.id, item])); const paths = buildWorkspaceKernelItemPathIndex(items); - const ranked = fuseWorkspaceSearchRanks({ keyword, limit, semantic }); + const ranked = fuseWorkspaceSearchRanks({ + keyword, + limit, + semantic: semantic.candidates, + }); return { failed: [], - status, + status: semantic.degraded ? "partial" : status, results: ranked.flatMap((candidate) => { const path = paths.get(candidate.itemId); const item = itemsById.get(candidate.itemId); @@ -162,7 +166,7 @@ export class WorkspaceSearchQuery { event: "workspace_search_semantic", fields: { workspace_id: this.workspaceId() }, }); - return []; + return { candidates: [], degraded: true }; } } @@ -191,16 +195,13 @@ export class WorkspaceSearchQuery { JOIN kernel_search_chunks c ON c.chunk_id = kernel_search_fts.chunk_id 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 - LEFT JOIN kernel_item_projections p - ON p.item_id = i.id - AND p.format = 'pages' - AND p.status = 'ready' WHERE kernel_search_fts MATCH ${match} - AND 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 + AND s.vector_status = 'ready' + AND NOT EXISTS ( + SELECT 1 + FROM kernel_search_pending pending + WHERE pending.item_id = i.id + ) AND ( ${localScope.kind} = 'workspace' OR ( @@ -225,23 +226,24 @@ export class WorkspaceSearchQuery { query: string; scope: WorkspaceSearchScope; types: WorkspaceSearchItemType[]; - }): Promise { + }) { if (input.scope.vectorFilters.length === 0) { - return []; + return { candidates: [], degraded: false }; } const [embedding] = await embedWorkspaceSearchTexts(this.ai, [input.query]); if (!embedding) { - return []; + throw new Error("Workspace search embedding response did not include the query vector."); } - const matches = await this.querySemanticMatches({ + 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 []; + return { candidates: [], degraded: semanticMatches.degraded }; } const rows = this.loadSemanticCandidates({ @@ -251,10 +253,13 @@ export class WorkspaceSearchQuery { }); const byVectorId = new Map(rows.map((row) => [row.chunk_id, mapSearchCandidate(row)])); - return vectorIds.flatMap((vectorId) => { - const candidate = byVectorId.get(vectorId); - return candidate ? [candidate] : []; - }); + return { + candidates: vectorIds.flatMap((vectorId) => { + const candidate = byVectorId.get(vectorId); + return candidate ? [candidate] : []; + }), + degraded: semanticMatches.degraded, + }; } private async querySemanticMatches(input: { @@ -263,8 +268,9 @@ export class WorkspaceSearchQuery { filters: WorkspaceSearchVectorFilter[]; }) { let bestMatches = new Map(); + const filters = input.filters.slice(0, maximumSemanticQueries); - for (const filterBatch of batchWorkspaceSearchValues(input.filters, semanticQueryConcurrency)) { + for (const filterBatch of batchWorkspaceSearchValues(filters, semanticQueryConcurrency)) { const results = await Promise.all( filterBatch.map((filter) => this.vectorize.query(input.embedding, { @@ -290,7 +296,10 @@ export class WorkspaceSearchQuery { ); } - return Array.from(bestMatches.values()).sort(compareSemanticMatches); + return { + degraded: filters.length < input.filters.length, + matches: Array.from(bestMatches.values()).sort(compareSemanticMatches), + }; } private loadSemanticCandidates(input: { @@ -313,16 +322,12 @@ export class WorkspaceSearchQuery { JOIN kernel_search_fts ON kernel_search_fts.chunk_id = c.chunk_id 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' - LEFT JOIN kernel_item_projections p - ON p.item_id = i.id - AND p.format = 'pages' - AND p.status = 'ready' WHERE c.chunk_id IN (SELECT value FROM json_each(${JSON.stringify(input.vectorIds)})) - AND 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 + AND NOT EXISTS ( + SELECT 1 + FROM kernel_search_pending pending + WHERE pending.item_id = i.id + ) AND ( ${localScope.kind} = 'workspace' OR ( @@ -396,7 +401,7 @@ function mapSearchResult( } function createFtsMatchExpression(query: string) { - const tokens = query.match(/[\p{L}\p{N}_]+/gu)?.slice(0, 20) ?? []; + const tokens = getSearchQueryTokens(query); return tokens.length > 0 ? tokens.map((token) => `"${token.replaceAll('"', '""')}"*`).join(" OR ") : null; @@ -408,10 +413,17 @@ function createSearchExcerpt(content: string, query: string) { return content; } - const term = query.match(/[\p{L}\p{N}_]{3,}/u)?.[0]?.toLocaleLowerCase(); - const matchIndex = term ? content.toLocaleLowerCase().indexOf(term) : -1; + 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-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-scope.ts b/src/features/workspaces/search/workspace-search-scope.ts index 8c7ae9a6..cfd6647c 100644 --- a/src/features/workspaces/search/workspace-search-scope.ts +++ b/src/features/workspaces/search/workspace-search-scope.ts @@ -117,7 +117,8 @@ function createWorkspaceSearchVectorFilters( if (input.kind === "item") { return [{ itemId: input.itemId }]; } - const type = input.types.length === 1 ? input.types[0] : undefined; + const uniqueTypes = Array.from(new Set(input.types)); + const type = uniqueTypes.length === 1 ? uniqueTypes[0] : undefined; if (input.kind === "workspace") { return type ? [{ type }] : [null]; } diff --git a/src/features/workspaces/search/workspace-search.test.ts b/src/features/workspaces/search/workspace-search.test.ts index 75d1f70f..327723e0 100644 --- a/src/features/workspaces/search/workspace-search.test.ts +++ b/src/features/workspaces/search/workspace-search.test.ts @@ -25,6 +25,27 @@ describe("workspace search", () => { 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 resolved = resolveWorkspaceSearchScope({ + items: [], + path: "/", + types: ["file", "file"], + }); + + expect(resolved).toMatchObject({ + scope: { vectorFilters: [{ type: "file" }] }, + status: "ready", + }); + }); + it("builds canonical source versions for documents and files", () => { expect( buildWorkspaceSearchSourceVersion({ From 0d48acfa118a26f18c975e85b7b64f6afba6a03a Mon Sep 17 00:00:00 2001 From: Urjit Chakraborty <135136842+urjitc@users.noreply.github.com> Date: Thu, 30 Jul 2026 22:31:13 -0400 Subject: [PATCH 06/13] fix(workspaces): drain extraction healing backlog Build deterministic repair jobs concurrently, then submit every eligible file in service-sized batches. This prevents workspaces with more than 100 broken or missing projections from waiting for another connection to continue healing. --- .../workspace-file-extraction-reconciler.ts | 82 ++++++++++--------- 1 file changed, 42 insertions(+), 40 deletions(-) diff --git a/src/features/workspaces/extraction/workspace-file-extraction-reconciler.ts b/src/features/workspaces/extraction/workspace-file-extraction-reconciler.ts index 7810a099..4713397f 100644 --- a/src/features/workspaces/extraction/workspace-file-extraction-reconciler.ts +++ b/src/features/workspaces/extraction/workspace-file-extraction-reconciler.ts @@ -50,53 +50,55 @@ export async function reconcileWorkspaceFileExtractions(input: { AND p.updated_at <= ${now - missingProjectionGraceMs} ) ) - ORDER BY i.created_at ASC - LIMIT ${workflowBatchSize} - `; + ORDER BY i.created_at ASC + `; const itemsById = new Map(input.items.map((item) => [item.id, item])); - const workflows: Array<{ - id: string; - params: WorkspaceFileExtractionWorkflowParams; - }> = []; - - for (const candidate of candidates) { - const item = itemsById.get(candidate.id); - if (!item) { - continue; - } + const workflows = ( + await Promise.all( + candidates.map(async (candidate) => { + const item = itemsById.get(candidate.id); + if (!item) { + return null; + } - const fileType = resolveWorkspaceFileTypeFromItem(item); - if (!fileType) { - continue; - } + const fileType = resolveWorkspaceFileTypeFromItem(item); + if (!fileType) { + return null; + } - const runKey = [ - extractionHealingVersion, - candidate.object_key, - candidate.projection_status ?? "missing", - candidate.projection_updated_at ?? 0, - ].join(":"); - const params = { - actorUserId: null, - assetKind: fileType.assetKind, - itemId: item.id, - requestId: extractionHealingVersion, - workspaceId: input.workspaceId, - } satisfies WorkspaceFileExtractionWorkflowParams; - workflows.push({ - id: await getWorkspaceFileExtractionWorkflowId({ - assetKind: params.assetKind, - itemId: item.id, - runKey, - workspaceId: input.workspaceId, + const runKey = [ + extractionHealingVersion, + candidate.object_key, + candidate.projection_status ?? "missing", + candidate.projection_updated_at ?? 0, + ].join(":"); + const params = { + actorUserId: null, + assetKind: fileType.assetKind, + itemId: item.id, + requestId: extractionHealingVersion, + workspaceId: input.workspaceId, + } satisfies WorkspaceFileExtractionWorkflowParams; + return { + id: await getWorkspaceFileExtractionWorkflowId({ + assetKind: params.assetKind, + itemId: item.id, + runKey, + workspaceId: input.workspaceId, + }), + params, + }; }), - params, - }); - } + ) + ).filter((workflow) => workflow !== null); if (workflows.length === 0) { return; } - await input.workflow.createBatch(workflows); + // Workflow batches are capped at 100; submit every eligible file without + // opening unbounded concurrent calls to the service. + for (let offset = 0; offset < workflows.length; offset += workflowBatchSize) { + await input.workflow.createBatch(workflows.slice(offset, offset + workflowBatchSize)); + } } From 2d0d6af01043505d125e3d8dc863e5fb210d3290 Mon Sep 17 00:00:00 2001 From: Urjit Chakraborty <135136842+urjitc@users.noreply.github.com> Date: Thu, 30 Jul 2026 22:31:20 -0400 Subject: [PATCH 07/13] docs(workspaces): record Vectorize provisioning Document the required BGE-M3 dimensions, cosine metric, and metadata index commands for staging and production so scoped semantic search infrastructure is reproducible from the repository. --- docs/configuration/deployments.mdx | 18 ++++++++++++++++++ wrangler.jsonc | 2 +- 2 files changed, 19 insertions(+), 1 deletion(-) 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/wrangler.jsonc b/wrangler.jsonc index 8029c6df..415251d4 100644 --- a/wrangler.jsonc +++ b/wrangler.jsonc @@ -34,7 +34,7 @@ "binding": "AI", "remote": true, }, - // Both indexes require string metadata indexes for itemId, parentId, and type. + // Provisioning: docs/configuration/deployments.mdx#workspace-search-vector-indexes "vectorize": [ { "binding": "WORKSPACE_SEARCH", From ebf1e0f0e921e17a91e5bc47cb8a27d4bc5c100c Mon Sep 17 00:00:00 2001 From: Urjit Chakraborty <135136842+urjitc@users.noreply.github.com> Date: Thu, 30 Jul 2026 22:39:46 -0400 Subject: [PATCH 08/13] fix(workspace-search): normalize repeated type filters Use Zod overwrite so duplicate type values collapse at the tool boundary while remaining representable in the MCP JSON Schema. --- .../workspaces/search/workspace-search-contract.ts | 1 + .../workspaces/search/workspace-search-scope.ts | 3 +-- .../workspaces/search/workspace-search.test.ts | 11 ++++------- 3 files changed, 6 insertions(+), 9 deletions(-) diff --git a/src/features/workspaces/search/workspace-search-contract.ts b/src/features/workspaces/search/workspace-search-contract.ts index 02d68cc8..7c906901 100644 --- a/src/features/workspaces/search/workspace-search-contract.ts +++ b/src/features/workspaces/search/workspace-search-contract.ts @@ -23,6 +23,7 @@ export const workspaceSearchInputSchema = z.object({ .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 diff --git a/src/features/workspaces/search/workspace-search-scope.ts b/src/features/workspaces/search/workspace-search-scope.ts index cfd6647c..8c7ae9a6 100644 --- a/src/features/workspaces/search/workspace-search-scope.ts +++ b/src/features/workspaces/search/workspace-search-scope.ts @@ -117,8 +117,7 @@ function createWorkspaceSearchVectorFilters( if (input.kind === "item") { return [{ itemId: input.itemId }]; } - const uniqueTypes = Array.from(new Set(input.types)); - const type = uniqueTypes.length === 1 ? uniqueTypes[0] : undefined; + const type = input.types.length === 1 ? input.types[0] : undefined; if (input.kind === "workspace") { return type ? [{ type }] : [null]; } diff --git a/src/features/workspaces/search/workspace-search.test.ts b/src/features/workspaces/search/workspace-search.test.ts index 327723e0..2993ab80 100644 --- a/src/features/workspaces/search/workspace-search.test.ts +++ b/src/features/workspaces/search/workspace-search.test.ts @@ -3,6 +3,7 @@ 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, @@ -34,16 +35,12 @@ describe("workspace search", () => { }); it("deduplicates repeated content types", () => { - const resolved = resolveWorkspaceSearchScope({ - items: [], - path: "/", + const input = workspaceSearchInputSchema.parse({ + query: "search files", types: ["file", "file"], }); - expect(resolved).toMatchObject({ - scope: { vectorFilters: [{ type: "file" }] }, - status: "ready", - }); + expect(input.types).toEqual(["file"]); }); it("builds canonical source versions for documents and files", () => { From ad7b82d2bcf14cdbe8de3835276b3289ce235231 Mon Sep 17 00:00:00 2001 From: Urjit Chakraborty <135136842+urjitc@users.noreply.github.com> Date: Fri, 31 Jul 2026 01:15:49 -0400 Subject: [PATCH 09/13] fix(workspace-search): simplify projection schema --- .../search/workspace-search-indexer.ts | 5 ++-- .../search/workspace-search-query.ts | 6 ++-- .../search/workspace-search-schema.ts | 30 ++++++++----------- .../search/workspace-search-version.ts | 4 +-- 4 files changed, 21 insertions(+), 24 deletions(-) diff --git a/src/features/workspaces/search/workspace-search-indexer.ts b/src/features/workspaces/search/workspace-search-indexer.ts index 6329c565..d1c8b1b3 100644 --- a/src/features/workspaces/search/workspace-search-indexer.ts +++ b/src/features/workspaces/search/workspace-search-indexer.ts @@ -389,8 +389,8 @@ export class WorkspaceSearchIndexer { ) `; this.sql` - INSERT INTO kernel_search_fts (rowid, chunk_id, title, content) - SELECT rowid, chunk_id, ${source.name}, ${chunk.content} + INSERT INTO kernel_search_fts (rowid, title, content) + SELECT rowid, ${source.name}, ${chunk.content} FROM kernel_search_chunks WHERE chunk_id = ${chunk.chunkId} `; @@ -473,6 +473,7 @@ export class WorkspaceSearchIndexer { 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); diff --git a/src/features/workspaces/search/workspace-search-query.ts b/src/features/workspaces/search/workspace-search-query.ts index 56e3890d..3bedf9f0 100644 --- a/src/features/workspaces/search/workspace-search-query.ts +++ b/src/features/workspaces/search/workspace-search-query.ts @@ -192,7 +192,7 @@ export class WorkspaceSearchQuery { c.start_line, c.end_line FROM kernel_search_fts - JOIN kernel_search_chunks c ON c.chunk_id = kernel_search_fts.chunk_id + 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} @@ -214,7 +214,7 @@ export class WorkspaceSearchQuery { ) ) AND i.type IN (SELECT value FROM json_each(${typesJson})) - ORDER BY bm25(kernel_search_fts, 0.0, 8.0, 1.0) ASC + ORDER BY bm25(kernel_search_fts, 8.0, 1.0) ASC LIMIT ${input.candidateLimit} `; @@ -319,7 +319,7 @@ export class WorkspaceSearchQuery { c.start_line, c.end_line FROM kernel_search_chunks c - JOIN kernel_search_fts ON kernel_search_fts.chunk_id = c.chunk_id + 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)})) diff --git a/src/features/workspaces/search/workspace-search-schema.ts b/src/features/workspaces/search/workspace-search-schema.ts index ff21b6b9..2bdd6325 100644 --- a/src/features/workspaces/search/workspace-search-schema.ts +++ b/src/features/workspaces/search/workspace-search-schema.ts @@ -1,16 +1,9 @@ import type { WorkspaceKernelSql } from "#/features/workspaces/kernel/workspace-kernel-schema"; +import { workspaceSearchIndexVersion } from "#/features/workspaces/search/workspace-search-version"; -const workspaceSearchStorageVersion = 1; +const workspaceSearchVersionKey = "workspace_search_version"; export function initializeWorkspaceSearchStorage(sql: WorkspaceKernelSql) { - sql` - CREATE TABLE IF NOT EXISTS kernel_search_metadata ( - key TEXT PRIMARY KEY, - value INTEGER NOT NULL - ) - `; - sql`CREATE INDEX IF NOT EXISTS kernel_search_vector_deletes_pending_idx - ON kernel_search_vector_deletes (attempts, requested_at)`; sql` CREATE TABLE IF NOT EXISTS kernel_search_vector_deletes ( vector_id TEXT PRIMARY KEY, @@ -18,14 +11,16 @@ export function initializeWorkspaceSearchStorage(sql: WorkspaceKernelSql) { 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: number }>` + const storedVersion = sql<{ value: string }>` SELECT value - FROM kernel_search_metadata - WHERE key = 'storage_version' + FROM kernel_meta + WHERE key = ${workspaceSearchVersionKey} LIMIT 1 `[0]?.value; - if (storedVersion !== workspaceSearchStorageVersion) { + 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 }>` @@ -66,7 +61,6 @@ export function initializeWorkspaceSearchStorage(sql: WorkspaceKernelSql) { ON kernel_search_chunks (item_id)`; sql` CREATE VIRTUAL TABLE IF NOT EXISTS kernel_search_fts USING fts5( - chunk_id UNINDEXED, title, content, tokenize = 'unicode61 remove_diacritics 2' @@ -83,8 +77,10 @@ export function initializeWorkspaceSearchStorage(sql: WorkspaceKernelSql) { ON kernel_search_pending (requested_at)`; sql` - INSERT INTO kernel_search_metadata (key, value) - VALUES ('storage_version', ${workspaceSearchStorageVersion}) - ON CONFLICT(key) DO UPDATE SET value = excluded.value + 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-version.ts b/src/features/workspaces/search/workspace-search-version.ts index 86fa714b..f7f0a8dd 100644 --- a/src/features/workspaces/search/workspace-search-version.ts +++ b/src/features/workspaces/search/workspace-search-version.ts @@ -1,6 +1,6 @@ /** - * Bump whenever the embedding model, chunking, or embedding text format changes. - * The SQL freshness checks in the indexer and query must mirror the format built here. + * 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"; From 0391865f3f6770fd38b467ff753a7735ab29c0c7 Mon Sep 17 00:00:00 2001 From: Urjit Chakraborty <135136842+urjitc@users.noreply.github.com> Date: Fri, 31 Jul 2026 01:23:45 -0400 Subject: [PATCH 10/13] refactor(workspaces): simplify extraction healing --- .../workspace-file-extraction-reconciler.ts | 44 +++++-------------- .../workspaces/kernel/workspace-kernel.ts | 1 - 2 files changed, 12 insertions(+), 33 deletions(-) diff --git a/src/features/workspaces/extraction/workspace-file-extraction-reconciler.ts b/src/features/workspaces/extraction/workspace-file-extraction-reconciler.ts index 4713397f..7e70f092 100644 --- a/src/features/workspaces/extraction/workspace-file-extraction-reconciler.ts +++ b/src/features/workspaces/extraction/workspace-file-extraction-reconciler.ts @@ -1,33 +1,28 @@ -import type { WorkspaceItemSummary } from "#/features/workspaces/contracts"; 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 { resolveWorkspaceFileTypeFromItem } from "#/features/workspaces/model/workspace-file"; +import { workspaceFileAssetKindSchema } from "#/features/workspaces/model/workspace-file"; const extractionHealingVersion = "extraction-healing-v1"; const failedExtractionCooldownMs = 15 * 60_000; -const missingProjectionGraceMs = workspaceExtractionStallThresholdMs; const workflowBatchSize = 100; export async function reconcileWorkspaceFileExtractions(input: { - items: readonly WorkspaceItemSummary[]; 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_status: string | null; - projection_updated_at: number | null; }>` SELECT + json_extract(i.metadata_json, '$.assetKind') AS asset_kind, i.id, - i.object_key, - p.status AS projection_status, - p.updated_at AS projection_updated_at + i.object_key FROM kernel_items i LEFT JOIN kernel_item_projections p ON p.item_id = i.id AND p.format = 'pages' @@ -35,7 +30,7 @@ export async function reconcileWorkspaceFileExtractions(input: { AND i.type = 'file' AND i.object_key IS NOT NULL AND ( - (p.item_id IS NULL AND i.created_at <= ${now - missingProjectionGraceMs}) + (p.item_id IS NULL AND i.created_at <= ${now - workspaceExtractionStallThresholdMs}) OR ( p.status = 'failed' AND p.updated_at <= ${now - failedExtractionCooldownMs} @@ -47,42 +42,31 @@ export async function reconcileWorkspaceFileExtractions(input: { OR ( p.status = 'ready' AND (p.object_key IS NULL OR p.source_hash IS NULL) - AND p.updated_at <= ${now - missingProjectionGraceMs} + AND p.updated_at <= ${now - workspaceExtractionStallThresholdMs} ) ) ORDER BY i.created_at ASC `; - const itemsById = new Map(input.items.map((item) => [item.id, item])); const workflows = ( await Promise.all( candidates.map(async (candidate) => { - const item = itemsById.get(candidate.id); - if (!item) { - return null; - } - - const fileType = resolveWorkspaceFileTypeFromItem(item); - if (!fileType) { + const assetKind = workspaceFileAssetKindSchema.safeParse(candidate.asset_kind); + if (!assetKind.success) { return null; } - const runKey = [ - extractionHealingVersion, - candidate.object_key, - candidate.projection_status ?? "missing", - candidate.projection_updated_at ?? 0, - ].join(":"); + const runKey = `${extractionHealingVersion}:${candidate.object_key}`; const params = { actorUserId: null, - assetKind: fileType.assetKind, - itemId: item.id, + assetKind: assetKind.data, + itemId: candidate.id, requestId: extractionHealingVersion, workspaceId: input.workspaceId, } satisfies WorkspaceFileExtractionWorkflowParams; return { id: await getWorkspaceFileExtractionWorkflowId({ assetKind: params.assetKind, - itemId: item.id, + itemId: candidate.id, runKey, workspaceId: input.workspaceId, }), @@ -92,10 +76,6 @@ export async function reconcileWorkspaceFileExtractions(input: { ) ).filter((workflow) => workflow !== null); - if (workflows.length === 0) { - return; - } - // Workflow batches are capped at 100; submit every eligible file without // opening unbounded concurrent calls to the service. for (let offset = 0; offset < workflows.length; offset += workflowBatchSize) { diff --git a/src/features/workspaces/kernel/workspace-kernel.ts b/src/features/workspaces/kernel/workspace-kernel.ts index 5c8b6b56..b61a87e4 100644 --- a/src/features/workspaces/kernel/workspace-kernel.ts +++ b/src/features/workspaces/kernel/workspace-kernel.ts @@ -459,7 +459,6 @@ export class WorkspaceKernel extends Agent { this.lastExtractionHealingRequestAt = now; this.ctx.waitUntil( reconcileWorkspaceFileExtractions({ - items: this.store.getPageItems(), sql: this.kernelSql, workflow: this.env.WORKSPACE_FILE_EXTRACTION_WORKFLOW, workspaceId: this.name, From f648053abbd68ec7c5bd57b51182ed27a29f4357 Mon Sep 17 00:00:00 2001 From: Urjit Chakraborty <135136842+urjitc@users.noreply.github.com> Date: Fri, 31 Jul 2026 01:39:24 -0400 Subject: [PATCH 11/13] fix(workspace-search): preserve lexical fallback Keep the complete local text index after semantic retries are exhausted. Report semantic coverage as partial, and stop requeueing permanently failed vector projections on every object wake. Queue superseded vectors with one SQL write instead of one write per chunk. --- .../search/workspace-search-indexer.ts | 35 +++++++++++++------ .../search/workspace-search-query.ts | 2 +- 2 files changed, 25 insertions(+), 12 deletions(-) diff --git a/src/features/workspaces/search/workspace-search-indexer.ts b/src/features/workspaces/search/workspace-search-indexer.ts index d1c8b1b3..2b107ab5 100644 --- a/src/features/workspaces/search/workspace-search-indexer.ts +++ b/src/features/workspaces/search/workspace-search-indexer.ts @@ -93,7 +93,6 @@ export class WorkspaceSearchIndexer { THEN ${workspaceSearchIndexVersion} || ':document:' || i.updated_at ELSE ${workspaceSearchIndexVersion} || ':file:' || i.updated_at || ':' || p.updated_at || ':' || p.source_hash END - OR s.vector_status != 'ready' ) ON CONFLICT(item_id) DO NOTHING `; @@ -197,6 +196,7 @@ export class WorkspaceSearchIndexer { this.beginIndex(source); let chunks: SearchIndexChunk[] = []; + let indexError: unknown = null; for await (const prepared of iteratePreparedWorkspaceSearchChunks({ bucket: this.bucket, source, @@ -210,19 +210,34 @@ export class WorkspaceSearchIndexer { chunks.push(chunk); if (chunks.length === searchChunkProcessingBatchSize) { - if (!(await this.indexChunkBatch(source, chunks))) { - return; + if (indexError === null) { + try { + if (!(await this.indexChunkBatch(source, chunks))) { + return; + } + } catch (error) { + indexError = error; + } } chunks = []; } } - if (chunks.length > 0 && !(await this.indexChunkBatch(source, chunks))) { - return; + 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); } @@ -412,13 +427,12 @@ export class WorkspaceSearchIndexer { } private discardItemChunks(itemId: string) { - for (const row of this.sql<{ chunk_id: string }>` - SELECT chunk_id + 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.queueVectorDelete(row.chunk_id); - } + `; this.deleteLocalChunks(itemId); } @@ -523,7 +537,6 @@ export class WorkspaceSearchIndexer { return false; } - this.discardItemChunks(itemId); this.sql` UPDATE kernel_search_items SET vector_status = 'failed' diff --git a/src/features/workspaces/search/workspace-search-query.ts b/src/features/workspaces/search/workspace-search-query.ts index 3bedf9f0..301e7a3f 100644 --- a/src/features/workspaces/search/workspace-search-query.ts +++ b/src/features/workspaces/search/workspace-search-query.ts @@ -196,7 +196,7 @@ export class WorkspaceSearchQuery { 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 = 'ready' + AND s.vector_status IN ('ready', 'failed') AND NOT EXISTS ( SELECT 1 FROM kernel_search_pending pending From db041d956286afb2530260b938ad927d81493b5d Mon Sep 17 00:00:00 2001 From: Urjit Chakraborty <135136842+urjitc@users.noreply.github.com> Date: Fri, 31 Jul 2026 01:39:36 -0400 Subject: [PATCH 12/13] fix(workspaces): tighten recovery lifecycle Hash and submit extraction repairs in service-sized batches, and key retries to the broken projection revision. Trigger healing from active searches as well as workspace opens. Use Agent.destroy for successful cleanup and await every R2 purge before deciding whether to retry. --- .../workspace-file-extraction-reconciler.ts | 63 ++++++++++--------- .../workspaces/kernel/workspace-kernel.ts | 39 ++++++------ 2 files changed, 53 insertions(+), 49 deletions(-) diff --git a/src/features/workspaces/extraction/workspace-file-extraction-reconciler.ts b/src/features/workspaces/extraction/workspace-file-extraction-reconciler.ts index 7e70f092..35178188 100644 --- a/src/features/workspaces/extraction/workspace-file-extraction-reconciler.ts +++ b/src/features/workspaces/extraction/workspace-file-extraction-reconciler.ts @@ -18,11 +18,13 @@ export async function reconcileWorkspaceFileExtractions(input: { 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 + 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' @@ -47,38 +49,37 @@ export async function reconcileWorkspaceFileExtractions(input: { ) ORDER BY i.created_at ASC `; - const workflows = ( - await Promise.all( - candidates.map(async (candidate) => { - const assetKind = workspaceFileAssetKindSchema.safeParse(candidate.asset_kind); - if (!assetKind.success) { - return null; - } + 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}`; - const params = { - actorUserId: null, - assetKind: assetKind.data, - itemId: candidate.id, - requestId: extractionHealingVersion, - workspaceId: input.workspaceId, - } satisfies WorkspaceFileExtractionWorkflowParams; - return { - id: await getWorkspaceFileExtractionWorkflowId({ - assetKind: params.assetKind, + const runKey = `${extractionHealingVersion}:${candidate.object_key}:${candidate.projection_updated_at}`; + const params = { + actorUserId: null, + assetKind: assetKind.data, itemId: candidate.id, - runKey, + requestId: extractionHealingVersion, workspaceId: input.workspaceId, - }), - params, - }; - }), - ) - ).filter((workflow) => workflow !== null); - - // Workflow batches are capped at 100; submit every eligible file without - // opening unbounded concurrent calls to the service. - for (let offset = 0; offset < workflows.length; offset += workflowBatchSize) { - await input.workflow.createBatch(workflows.slice(offset, offset + workflowBatchSize)); + } 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/kernel/workspace-kernel.ts b/src/features/workspaces/kernel/workspace-kernel.ts index b61a87e4..25835820 100644 --- a/src/features/workspaces/kernel/workspace-kernel.ts +++ b/src/features/workspaces/kernel/workspace-kernel.ts @@ -320,6 +320,7 @@ export class WorkspaceKernel extends Agent { } async searchWorkspace(input: WorkspaceSearchInput) { + this.requestWorkspaceFileExtractionHealing(); if (this.search.hasPending()) { this.ctx.waitUntil(this.scheduleWorkspaceSearchIndexing()); } @@ -400,29 +401,31 @@ export class WorkspaceKernel extends Agent { } } - try { - await Promise.all([ - deleteR2Prefix( - this.env.WORKSPACE_KERNEL_FILES, - getChatAttachmentWorkspacePrefix(workspaceId), - ), - deleteR2Prefix(this.env.WORKSPACE_KERNEL_FILES, `uploads/workspaces/${workspaceId}/`), - deleteR2Prefix(this.env.WORKSPACE_KERNEL_FILES, `workspace_kernel_files/${workspaceId}/`), - deleteR2Prefix(this.env.WORKSPACE_KERNEL_FILES, `workspace_file_objects/${workspaceId}/`), - deleteR2Prefix(this.env.WORKSPACE_KERNEL_FILES, `workspace_file_uploads/${workspaceId}/`), - ]); - } catch (error) { + const r2PurgeResults = await Promise.allSettled([ + deleteR2Prefix( + this.env.WORKSPACE_KERNEL_FILES, + getChatAttachmentWorkspacePrefix(workspaceId), + ), + deleteR2Prefix(this.env.WORKSPACE_KERNEL_FILES, `uploads/workspaces/${workspaceId}/`), + deleteR2Prefix(this.env.WORKSPACE_KERNEL_FILES, `workspace_kernel_files/${workspaceId}/`), + 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; - recordOperationalFailure({ - error, - event: "workspace_r2_purge", - fields: { workspace_id: workspaceId }, - }); + for (const result of r2Failures) { + recordOperationalFailure({ + error: result.reason, + event: "workspace_r2_purge", + fields: { workspace_id: workspaceId }, + }); + } } // Keep the local inventory when a remote purge fails so cleanup can be retried. if (failed === 0) { - await this.ctx.storage.deleteAll(); + await this.destroy(); } else { const attempt = input.attempt ?? 1; if (attempt < workspacePurgeMaximumAttempts) { From 248f37141aed273f8f1d23406ab6b63d775e6da2 Mon Sep 17 00:00:00 2001 From: Urjit Chakraborty <135136842+urjitc@users.noreply.github.com> Date: Fri, 31 Jul 2026 01:39:48 -0400 Subject: [PATCH 13/13] refactor(workspace-search): remove unused surfaces Delete the obsolete document timestamp helper and stop exporting search-only implementation types and adapters. Keep the public workspace search contract unchanged. --- .../ai/workspace-tool-result-adapters.ts | 4 ++-- .../workspaces/documents/document-item-content.ts | 14 -------------- .../operations/workspace-tool-definitions.ts | 2 +- .../workspaces/search/workspace-search-chunks.ts | 2 +- .../workspaces/search/workspace-search-content.ts | 2 +- .../workspaces/search/workspace-search-contract.ts | 8 ++++---- 6 files changed, 9 insertions(+), 23 deletions(-) diff --git a/src/features/workspaces/ai/workspace-tool-result-adapters.ts b/src/features/workspaces/ai/workspace-tool-result-adapters.ts index 036c153a..4e2333b3 100644 --- a/src/features/workspaces/ai/workspace-tool-result-adapters.ts +++ b/src/features/workspaces/ai/workspace-tool-result-adapters.ts @@ -23,13 +23,13 @@ function defineWorkspaceToolResultAdapter(input: { }; } -export const workspaceReadItemsResultAdapter = defineWorkspaceToolResultAdapter({ +const workspaceReadItemsResultAdapter = defineWorkspaceToolResultAdapter({ collectReferences: (output) => output.references, outputSchema: workspaceReadItemsOutputSchema, projectOutput: createWorkspaceReadItemsModelOutput, }); -export const workspaceSearchResultAdapter = defineWorkspaceToolResultAdapter({ +const workspaceSearchResultAdapter = defineWorkspaceToolResultAdapter({ collectReferences: (output) => output.references, outputSchema: workspaceSearchOutputSchema, projectOutput: createWorkspaceSearchModelOutput, 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/operations/workspace-tool-definitions.ts b/src/features/workspaces/operations/workspace-tool-definitions.ts index ee07d811..c7fbab64 100644 --- a/src/features/workspaces/operations/workspace-tool-definitions.ts +++ b/src/features/workspaces/operations/workspace-tool-definitions.ts @@ -66,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, diff --git a/src/features/workspaces/search/workspace-search-chunks.ts b/src/features/workspaces/search/workspace-search-chunks.ts index 384c2eeb..f4f0b262 100644 --- a/src/features/workspaces/search/workspace-search-chunks.ts +++ b/src/features/workspaces/search/workspace-search-chunks.ts @@ -2,7 +2,7 @@ const targetChunkCharacters = 1_800; const minimumChunkCharacters = 900; const overlapCharacters = 220; -export interface WorkspaceSearchTextChunk { +interface WorkspaceSearchTextChunk { content: string; endLine: number; startLine: number; diff --git a/src/features/workspaces/search/workspace-search-content.ts b/src/features/workspaces/search/workspace-search-content.ts index e618a343..33f5867a 100644 --- a/src/features/workspaces/search/workspace-search-content.ts +++ b/src/features/workspaces/search/workspace-search-content.ts @@ -19,7 +19,7 @@ export type WorkspaceSearchIndexSource = WorkspaceSearchIndexSourceBase & | { objectKey: string; sourceHash: string; type: "file" } ); -export interface PreparedWorkspaceSearchChunk { +interface PreparedWorkspaceSearchChunk { content: string; endLine: number | null; index: number; diff --git a/src/features/workspaces/search/workspace-search-contract.ts b/src/features/workspaces/search/workspace-search-contract.ts index 7c906901..3cf93cea 100644 --- a/src/features/workspaces/search/workspace-search-contract.ts +++ b/src/features/workspaces/search/workspace-search-contract.ts @@ -3,7 +3,7 @@ import { z } from "zod"; import { workspaceReferenceRecordSchema } from "#/features/workspaces/locations/workspace-location"; import { workspaceFileAssetKindSchema } from "#/features/workspaces/model/workspace-file"; -export const workspaceSearchItemTypeSchema = z.enum(["document", "file"]); +const workspaceSearchItemTypeSchema = z.enum(["document", "file"]); export const workspaceSearchInputSchema = z.object({ query: z @@ -47,7 +47,7 @@ const workspaceSearchLocationSchema = z.discriminatedUnion("kind", [ }), ]); -export const workspaceSearchResultSchema = z.object({ +const workspaceSearchResultSchema = z.object({ assetKind: workspaceFileAssetKindSchema.optional(), excerpt: z.string(), itemId: z.string().min(1), @@ -57,12 +57,12 @@ export const workspaceSearchResultSchema = z.object({ type: workspaceSearchItemTypeSchema, }); -export const workspaceSearchFailureSchema = z.object({ +const workspaceSearchFailureSchema = z.object({ code: z.enum(["path_not_absolute", "path_not_found"]), path: z.string(), }); -export const workspaceSearchStatusSchema = z.enum(["ready", "indexing", "partial"]); +const workspaceSearchStatusSchema = z.enum(["ready", "indexing", "partial"]); export const workspaceSearchOutputSchema = z.object({ failed: z.array(workspaceSearchFailureSchema),