Add WIP search AI tool

Author: freekh

packages/shared/src/internal/index.ts

@@ -13,3 +13,4 @@ export * from "./parseRoutePattern";
 export * from "./getNextAppRouterSourceFolder";
 export * from "./routeValidation";
 export * from "./getErrorMessageFromUnknownJson";
+export * from "./zod/SerializedSchema";

packages/ui/spa/components/ValProvider.tsx

@@ -49,13 +49,14 @@ import { ValErrorProvider } from "./ValErrorProvider";
 import { ValPortalProvider } from "./ValPortalProvider";
 import { ValFieldProvider } from "./ValFieldProvider";
 import { ValRemoteProvider } from "./ValRemoteProvider";
+import { toolNames } from "../utils/toolNames";
 
 export const AITool = z.object({
-  name: z.string(),
+  name: z.enum(toolNames),
   description: z.string(),
   parameters: z.object({
     type: z.literal("object"),
-    properties: z.record(z.unknown()),
+    properties: z.record(z.string(), z.unknown()),
     required: z.array(z.string()).optional(),
   }),
 });
@@ -110,7 +111,9 @@ type ValContextValue = {
           | "unauthorized";
       };
   subscribeToWsMessages: (handler: WsMessageHandler) => () => void;
-  sendWsMessage: (data: z.infer<typeof AIPromptMessage> | AIToolResultMessage) => void;
+  sendWsMessage: (
+    data: z.infer<typeof AIPromptMessage> | AIToolResultMessage,
+  ) => void;
   isWsConnected: boolean;
 };
 const ValContext = React.createContext<ValContextValue>(

packages/ui/spa/hooks/useAI.ts

@@ -6,18 +6,37 @@ import {
   useWsMessages,
 } from "../components/ValProvider";
 import type { WsExtendedMessage } from "./useStatus";
+import { useAISearch } from "./useAISearch";
 
 const GET_ALL_SCHEMA_TOOL: AITool = {
   name: "get_all_schema",
   description:
     "Get all val schemas — returns the complete schema definitions for all val modules",
-  parameters: { type: "object", properties: {} },
+  parameters: {
+    type: "object",
+    properties: {},
+    required: [],
+  },
 };
+const SEARCH_CONTENT_TOOL: AITool = {
+  name: "search_content",
+  description:
+    "Search content — accepts a query and returns a list of matching content items",
+  parameters: {
+    type: "object",
+    properties: {
+      query: { type: "string", description: "The search query" },
+    },
+    required: ["query"],
+  },
+};
+const ALL_TOOLS: AITool[] = [GET_ALL_SCHEMA_TOOL, SEARCH_CONTENT_TOOL];
 
 export function useAI(chatRef: React.RefObject<AIChatHandle | null>) {
   const { subscribeToWsMessages, sendWsMessage, isWsConnected } =
     useWsMessages();
   const syncEngine = useSyncEngine();
+  const aiSearch = useAISearch();
   const [isStreaming, setIsStreaming] = useState(false);
   // Track active streaming ID — startAssistantMessage always appends a new
   // message (NOT idempotent), so we must only call it once per message ID.
@@ -50,6 +69,13 @@ export function useAI(chatRef: React.RefObject<AIChatHandle | null>) {
             toolCallId: message.toolCallId,
             result: schemas ?? {},
           });
+        } else if (message.name === "search_content") {
+          const args = message.arguments as { query: string };
+
+          aiSearch.query(args.query, message.toolCallId);
+        } else {
+          const exhaustiveCheck: never = message.name;
+          console.error("Received unknown tool call in useAI", exhaustiveCheck);
         }
       } else {
         console.error("Received unknown message type in useAI", message);
@@ -58,7 +84,7 @@ export function useAI(chatRef: React.RefObject<AIChatHandle | null>) {
     };
 
     return subscribeToWsMessages(handler);
-  }, [subscribeToWsMessages, sendWsMessage, syncEngine, chatRef]);
+  }, [subscribeToWsMessages, sendWsMessage, syncEngine, aiSearch, chatRef]);
 
   const sendMessage = useCallback(
     (text: string) => {
@@ -71,7 +97,7 @@ export function useAI(chatRef: React.RefObject<AIChatHandle | null>) {
 - Users are working in another code base where they have Val installed, and are using you to ask questions about their content schemas, and to get help writing content.
 - Be concise, accurate, and helpful.`,
         id: crypto.randomUUID(),
-        tools: [GET_ALL_SCHEMA_TOOL],
+        tools: ALL_TOOLS,
       });
     },
     [sendWsMessage],

packages/ui/spa/hooks/useAISearch.ts

@@ -0,0 +1,123 @@
+import { useEffect, useRef, useCallback, useSyncExternalStore } from "react";
+import { Json, ModuleFilePath, SerializedSchema } from "@valbuild/core";
+import type { WorkerRequest, WorkerResponse } from "../search/worker-types";
+import type { SearchResult } from "../search/useSearchWorker";
+import { useWsMessages, useSyncEngine } from "../components/ValProvider";
+
+export function useAISearch() {
+  const { sendWsMessage } = useWsMessages();
+  const syncEngine = useSyncEngine();
+
+  const schemas = useSyncExternalStore(
+    syncEngine.subscribe("schema"),
+    () => syncEngine.getAllSchemasSnapshot(),
+    () => syncEngine.getAllSchemasSnapshot(),
+  );
+  const sources = useSyncExternalStore(
+    syncEngine.subscribe("all-sources"),
+    () => syncEngine.getAllSourcesSnapshot(),
+    () => syncEngine.getAllSourcesSnapshot(),
+  );
+
+  const workerRef = useRef<Worker | null>(null);
+  const isIndexBuiltRef = useRef(false);
+  const requestIdCounter = useRef(0);
+  const pendingRequests = useRef<
+    Map<
+      string,
+      { resolve: (value: unknown) => void; reject: (error: Error) => void }
+    >
+  >(new Map());
+
+  useEffect(() => {
+    isIndexBuiltRef.current = false;
+  }, [schemas, sources]);
+
+  useEffect(() => {
+    const worker = new Worker(
+      new URL("../search/search.worker.ts", import.meta.url),
+      { type: "module" },
+    );
+    workerRef.current = worker;
+
+    worker.onmessage = (event: MessageEvent<WorkerResponse>) => {
+      const response = event.data;
+      const pending = pendingRequests.current.get(response.id);
+      if (!pending) return;
+      pendingRequests.current.delete(response.id);
+      if (response.type === "index-ready") {
+        isIndexBuiltRef.current = true;
+        pending.resolve(null);
+      } else if (response.type === "search-results") {
+        pending.resolve(response.results);
+      } else if (response.type === "error") {
+        pending.reject(new Error(response.error));
+      }
+    };
+
+    return () => {
+      worker.terminate();
+      workerRef.current = null;
+    };
+  }, []);
+
+  const query = useCallback(
+    async (searchQuery: string, toolCallId: string) => {
+      if (!workerRef.current) return;
+      try {
+        if (!isIndexBuiltRef.current) {
+          const modules: Record<
+            ModuleFilePath,
+            { source: Json; schema: SerializedSchema }
+          > = {};
+          for (const path in schemas) {
+            const moduleFilePath = path as ModuleFilePath;
+            const source = sources?.[moduleFilePath];
+            if (source !== undefined) {
+              modules[moduleFilePath] = {
+                source,
+                schema: schemas[moduleFilePath],
+              };
+            }
+          }
+          const buildId = `ai-build-${requestIdCounter.current++}`;
+          await new Promise((resolve, reject) => {
+            pendingRequests.current.set(buildId, { resolve, reject });
+            workerRef.current!.postMessage({
+              type: "build-index",
+              id: buildId,
+              modules,
+            } satisfies WorkerRequest);
+          });
+        }
+
+        const results = await new Promise<SearchResult[]>((resolve, reject) => {
+          pendingRequests.current.set(toolCallId, {
+            resolve: resolve as (value: unknown) => void,
+            reject,
+          });
+          workerRef.current!.postMessage({
+            type: "search",
+            id: toolCallId,
+            query: searchQuery,
+          } satisfies WorkerRequest);
+        });
+
+        console.log("Search results for query", searchQuery, results, {
+          toolCallId,
+        });
+        sendWsMessage({ type: "ai_tool_result", toolCallId, result: results });
+      } catch (error) {
+        sendWsMessage({
+          type: "ai_tool_result",
+          toolCallId,
+          result: String(error),
+          isError: true,
+        });
+      }
+    },
+    [schemas, sources, sendWsMessage],
+  );
+
+  return { query };
+}

packages/ui/spa/hooks/useStatus.ts

@@ -10,6 +10,7 @@ import React, {
   useCallback,
 } from "react";
 import { z } from "zod";
+import { toolNames } from "../utils/toolNames";
 
 const PatchId = z
   .string()
@@ -36,7 +37,7 @@ export const AIToolCallMessage = z.object({
   type: z.literal("ai_tool_call"),
   id: z.string(),
   toolCallId: z.string(),
-  name: z.string(),
+  name: z.enum(toolNames),
   arguments: z.unknown(),
 });
 export const AIServerMessage = z.union([

packages/ui/spa/utils/toolNames.ts

@@ -0,0 +1 @@
+export const toolNames = ["get_all_schema", "search_content"] as const;