feat(workspaces): add hybrid workspace search

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.
 </Warning>
 
+## 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.

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" }),

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) ?? []),
+			);
 		}
 	}
 

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<TSchema extends z.ZodTypeAny>(input: {
+	collectReferences: (output: z.output<TSchema>) => readonly WorkspaceReferenceRecord[];
+	outputSchema: TSchema;
+	projectOutput: (output: z.output<TSchema>) => unknown;
+}) {
+	return {
+		collectReferences: (output: unknown) => {
+			const parsed = input.outputSchema.safeParse(output);
+			return parsed.success ? input.collectReferences(parsed.data) : [];
+		},
+		projectOutput: (output: unknown) => {
+			return input.projectOutput(input.outputSchema.parse(output)) as JSONValue;
+		},
+	};
+}
+
+const workspaceReadItemsResultAdapter = defineWorkspaceToolResultAdapter({
+	collectReferences: (output) => output.references,
+	outputSchema: workspaceReadItemsOutputSchema,
+	projectOutput: createWorkspaceReadItemsModelOutput,
+});
+
+const workspaceSearchResultAdapter = defineWorkspaceToolResultAdapter({
+	collectReferences: (output) => output.references,
+	outputSchema: workspaceSearchOutputSchema,
+	projectOutput: createWorkspaceSearchModelOutput,
+});
+
+const workspaceToolResultAdapters = {
+	workspace_read_items: workspaceReadItemsResultAdapter,
+	workspace_search: workspaceSearchResultAdapter,
+} as const;
+
+export function getWorkspaceToolResultAdapter(name: string) {
+	return Object.hasOwn(workspaceToolResultAdapters, name)
+		? workspaceToolResultAdapters[name as keyof typeof workspaceToolResultAdapters]
+		: null;
+}

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,35 +15,33 @@ import {
 } from "#/features/workspaces/operations/workspace-access-context";
 
 type WorkspaceThreadToolConfig = {
-	definition: WorkspaceToolDefinition;
+	definition: (typeof workspaceToolDefinitions)[number];
 	getThreadContext: () => Promise<AIThreadContext | null>;
 	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,
 		inputSchema: definition.inputSchema,
 		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;

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
-	`;
-}

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)}`;
-}

src/features/workspaces/extraction/workspace-file-extraction-reconciler.ts

@@ -0,0 +1,85 @@
+import type { WorkspaceFileExtractionWorkflowParams } from "#/features/workspaces/extraction/types";
+import { getWorkspaceFileExtractionWorkflowId } from "#/features/workspaces/extraction/workspace-file-extraction-workflow-id";
+import { workspaceExtractionStallThresholdMs } from "#/features/workspaces/extraction/workspace-projection-readiness";
+import type { WorkspaceKernelSql } from "#/features/workspaces/kernel/workspace-kernel-schema";
+import { workspaceFileAssetKindSchema } from "#/features/workspaces/model/workspace-file";
+
+const extractionHealingVersion = "extraction-healing-v1";
+const failedExtractionCooldownMs = 15 * 60_000;
+const workflowBatchSize = 100;
+
+export async function reconcileWorkspaceFileExtractions(input: {
+	sql: WorkspaceKernelSql;
+	workflow: Workflow<WorkspaceFileExtractionWorkflowParams>;
+	workspaceId: string;
+}) {
+	const now = Date.now();
+	const candidates = input.sql<{
+		asset_kind: string | number | boolean | null;
+		id: string;
+		object_key: string;
+		projection_updated_at: number;
+	}>`
+		SELECT
+			json_extract(i.metadata_json, '$.assetKind') AS asset_kind,
+			i.id,
+			i.object_key,
+			COALESCE(p.updated_at, i.created_at) AS projection_updated_at
+		FROM kernel_items i
+		LEFT JOIN kernel_item_projections p
+			ON p.item_id = i.id AND p.format = 'pages'
+		WHERE i.deleted_at IS NULL
+			AND i.type = 'file'
+			AND i.object_key IS NOT NULL
+			AND (
+				(p.item_id IS NULL AND i.created_at <= ${now - workspaceExtractionStallThresholdMs})
+				OR (
+					p.status = 'failed'
+					AND p.updated_at <= ${now - failedExtractionCooldownMs}
+				)
+				OR (
+					p.status = 'processing'
+					AND p.updated_at <= ${now - workspaceExtractionStallThresholdMs}
+				)
+				OR (
+					p.status = 'ready'
+					AND (p.object_key IS NULL OR p.source_hash IS NULL)
+					AND p.updated_at <= ${now - workspaceExtractionStallThresholdMs}
+				)
+			)
+			ORDER BY i.created_at ASC
+		`;
+	for (let offset = 0; offset < candidates.length; offset += workflowBatchSize) {
+		const workflows = (
+			await Promise.all(
+				candidates.slice(offset, offset + workflowBatchSize).map(async (candidate) => {
+					const assetKind = workspaceFileAssetKindSchema.safeParse(candidate.asset_kind);
+					if (!assetKind.success) {
+						return null;
+					}
+
+					const runKey = `${extractionHealingVersion}:${candidate.object_key}:${candidate.projection_updated_at}`;
+					const params = {
+						actorUserId: null,
+						assetKind: assetKind.data,
+						itemId: candidate.id,
+						requestId: extractionHealingVersion,
+						workspaceId: input.workspaceId,
+					} satisfies WorkspaceFileExtractionWorkflowParams;
+					return {
+						id: await getWorkspaceFileExtractionWorkflowId({
+							assetKind: params.assetKind,
+							itemId: candidate.id,
+							runKey,
+							workspaceId: input.workspaceId,
+						}),
+						params,
+					};
+				}),
+			)
+		).filter((workflow) => workflow !== null);
+		if (workflows.length > 0) {
+			await input.workflow.createBatch(workflows);
+		}
+	}
+}

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)}`;
+}