Build tool registry from serialized artifacts

Rhys Sullivan committed 8 months ago

Commit 5ecb1f61 parent b096a823 files · +66 −20Browse files

Unified · Split

3 files

executor/packages/convex/runtime/tool_call_resolution.ts +15 −5
executor/packages/convex/runtime/tool_invocation.ts +12 −5
executor/packages/convex/runtime/workspace_tools.ts +39 −10

executor/packages/convex/runtime/tool_call_resolution.ts+15−5

@@ -8,6 +8,7 @@ import { rehydrateTools, type SerializedTool } from "../../core/src/tool/source- · 7 unmodified lines
8	import { getDecisionForContext, getToolDecision } from "./policy";	8	import { getDecisionForContext, getToolDecision } from "./policy";
9	import { normalizeToolPathForLookup, toPreferredToolPath } from "./tool_paths";	9	import { normalizeToolPathForLookup, toPreferredToolPath } from "./tool_paths";
10	import { baseTools } from "./workspace_tools";	10	import { baseTools } from "./workspace_tools";
		11	import { sourceSignature } from "./tool_source_loading";
11		12
12	export function getGraphqlDecision(	13	export function getGraphqlDecision(
13	task: TaskRecord,	14	task: TaskRecord,
@@ -69,13 +70,22 @@ async function resolveRegistryBuildId( · 55 unmodified lines
69	ctx: ActionCtx,	70	ctx: ActionCtx,
70	workspaceId: TaskRecord["workspaceId"],	71	workspaceId: TaskRecord["workspaceId"],
71	): Promise<string> {	72	): Promise<string> {
72	const state = await ctx.runQuery(internal.toolRegistry.getState, { workspaceId }) as	73	const [state, sources] = await Promise.all([
73	\| null	74	ctx.runQuery(internal.toolRegistry.getState, { workspaceId }) as Promise<null \| { signature: string; readyBuildId?: string }>,
74	\| { readyBuildId?: string };	75	ctx.runQuery(internal.database.listToolSources, { workspaceId }) as Promise<Array<{ id: string; updatedAt: number; enabled: boolean }>>,
		76	]);
		77
		78	const enabledSources = sources.filter((source) => source.enabled);
		79	const signature = sourceSignature(workspaceId, enabledSources);
		80	const expectedSignature = `toolreg_v1\|${signature}`;
75	const buildId = state?.readyBuildId;	81	const buildId = state?.readyBuildId;
76	if (!buildId) {	82
77	throw new Error("Tool registry is not ready yet. Open Tools in the UI to build the registry.");	83	if (!buildId \|\| state.signature !== expectedSignature) {
		84	throw new Error(
		85	"Tool registry is not ready (or is stale). Open Tools to refresh, or call listToolsWithWarnings to rebuild.",
		86	);
78	}	87	}
		88
79	return buildId;	89	return buildId;
80	}	90	}
81		91

executor/packages/convex/runtime/tool_invocation.ts+12−5

 import { completeToolCall, denyToolCall, failToolCall } from "./tool_call_lifecycle";
 import { assertPersistedCallRunnable, resolveCredentialHeaders } from "./tool_call_credentials";
 import { getGraphqlDecision, resolveToolForCall } from "./tool_call_resolution";
+import { sourceSignature } from "./tool_source_loading";
 function createApprovalId(): string {
   return `approval_${crypto.randomUUID()}`;
   // Fast system tools are handled server-side from the registry.
   if (toolPath === "discover" || toolPath === "catalog.namespaces" || toolPath === "catalog.tools") {
     const toolRegistry = internal.toolRegistry;
-    const state = await ctx.runQuery(toolRegistry.getState, { workspaceId: task.workspaceId }) as
+    const [state, sources] = await Promise.all([
-      | null
+      ctx.runQuery(toolRegistry.getState, { workspaceId: task.workspaceId }) as Promise<null | { signature: string; readyBuildId?: string }>,
-      | { readyBuildId?: string };
+      ctx.runQuery(internal.database.listToolSources, { workspaceId: task.workspaceId }) as Promise<Array<{ id: string; updatedAt: number; enabled: boolean }>>,
+    ]);
+    const enabledSources = sources.filter((source) => source.enabled);
+    const signature = sourceSignature(task.workspaceId, enabledSources);
+    const expectedSignature = `toolreg_v1|${signature}`;
     const buildId = state?.readyBuildId;
-    if (!buildId) {
+    if (!buildId || state.signature !== expectedSignature) {
-      throw new Error("Tool registry is not ready yet. Open Tools in the UI to build the registry.");
+      throw new Error(
+        "Tool registry is not ready (or is stale). Open Tools to refresh, or call listToolsWithWarnings to rebuild.",
+      );
     }
     const payload = input && typeof input === "object" ? (input as Record<string, unknown>) : {};

executor/packages/convex/runtime/workspace_tools.ts+39−10

   type CompiledToolSourceArtifact,
   type WorkspaceToolSnapshot,
 } from "../../core/src/tool-sources";
-import { serializeTools } from "../../core/src/tool/source-serialization";
+import type { SerializedTool } from "../../core/src/tool/source-serialization";
 import type { ExternalToolSourceConfig } from "../../core/src/tool/source-types";
 import type {
   AccessPolicyRecord,
   ctx: ActionCtx,
   args: {
     workspaceId: Id<"workspaces">;
-    signature: string;
+    registrySignature: string;
-    tools: ToolDefinition[];
+    serializedTools: SerializedTool[];
   },
 ): Promise<{ buildId: string }> {
   const buildId = `toolreg_${crypto.randomUUID()}`;
   await ctx.runMutation(internal.toolRegistry.beginBuild, {
     workspaceId: args.workspaceId,
-    signature: args.signature,
+    signature: args.registrySignature,
     buildId,
   });
-  const serialized = serializeTools(args.tools);
+  const entries = args.serializedTools.map((st) => {
-  const entries = serialized.map((st) => {
     if (st.path === "discover" || st.path.startsWith("catalog.")) {
       return null;
     }
   return { buildId };
 }
+async function maybeBuildWorkspaceToolRegistry(
+  ctx: ActionCtx,
+  args: {
+    workspaceId: Id<"workspaces">;
+    registrySignature: string;
+    serializedTools: SerializedTool[];
+  },
+): Promise<void> {
+  const state = await ctx.runQuery(internal.toolRegistry.getState, { workspaceId: args.workspaceId }) as
+    | null
+    | { signature: string; readyBuildId?: string };
+  if (state?.readyBuildId && state.signature === args.registrySignature) {
+    return;
+  }
+  await buildWorkspaceToolRegistry(ctx, args);
+}
 export async function getWorkspaceTools(
   ctx: ActionCtx,
   workspaceId: Id<"workspaces">,
   traceStep("listToolSources", listSourcesStartedAt);
   const hasOpenApiSource = sources.some((source: { type: string }) => source.type === "openapi");
   const signature = sourceSignature(workspaceId, sources);
+  const registrySignature = `toolreg_v1|${signature}`;
   const debugBase: Omit<WorkspaceToolsDebug, "mode" | "normalizedSourceCount" | "cacheHit" | "cacheFresh" | "timedOutSources" | "durationMs" | "trace"> = {
       includeDts,
       sourceTimeoutMs: sourceTimeoutMs ?? null,
         const typesStorageId = cacheEntry.typesStorageId as Id<"_storage"> | undefined;
         if (cacheEntry.isFresh) {
           if (typesStorageId) {
+            // Registry is derived from the cached snapshot. Keep it in sync.
+            try {
+              await maybeBuildWorkspaceToolRegistry(ctx, {
+                workspaceId,
+                registrySignature,
+                serializedTools: snapshot.externalArtifacts.flatMap((artifact) => artifact.tools),
+              });
+            } catch (error) {
+              const msg = error instanceof Error ? error.message : String(error);
+              console.warn(`[executor] workspace tool registry build failed for '${workspaceId}': ${msg}`);
+            }
             return {
               tools: merged,
               warnings: snapshot.warnings,
     // Build a per-tool registry for fast discover + invocation.
     const registryStartedAt = Date.now();
-    await buildWorkspaceToolRegistry(ctx, {
+    await maybeBuildWorkspaceToolRegistry(ctx, {
       workspaceId,
-      signature,
+      registrySignature,
-      tools: allTools,
+      serializedTools: externalArtifacts.flatMap((artifact) => artifact.tools),
     });
     traceStep("toolRegistryWrite", registryStartedAt);