History

Build tool registry from serialized artifacts

Rhys Sullivan committed 8 months ago

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

Unified · Split

3 files
  1. executor/packages/convex/runtime/tool_call_resolution.ts +15 −5
  2. executor/packages/convex/runtime/tool_invocation.ts +12 −5
  3. 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
88 import { getDecisionForContext, getToolDecision } from "./policy";
99 import { normalizeToolPathForLookup, toPreferredToolPath } from "./tool_paths";
1010 import { baseTools } from "./workspace_tools";
11+import { sourceSignature } from "./tool_source_loading";
1112
1213 export function getGraphqlDecision(
1314 task: TaskRecord,
@@ -69,13 +70,22 @@ async function resolveRegistryBuildId( · 55 unmodified lines
6970 ctx: ActionCtx,
7071 workspaceId: TaskRecord["workspaceId"],
7172 ): Promise<string> {
72− const state = await ctx.runQuery(internal.toolRegistry.getState, { workspaceId }) as
73− | null
74− | { readyBuildId?: string };
73+ const [state, sources] = await Promise.all([
74+ ctx.runQuery(internal.toolRegistry.getState, { workspaceId }) as Promise<null | { signature: string; 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}`;
7581 const buildId = state?.readyBuildId;
76− if (!buildId) {
77− throw new Error("Tool registry is not ready yet. Open Tools in the UI to build the registry.");
82+
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+ );
7887 }
88+
7989 return buildId;
8090 }
8191
executor/packages/convex/runtime/tool_invocation.ts+12−5
@@ -21,6 +21,7 @@ import { publishTaskEvent } from "./events"; · 20 unmodified lines
2121 import { completeToolCall, denyToolCall, failToolCall } from "./tool_call_lifecycle";
2222 import { assertPersistedCallRunnable, resolveCredentialHeaders } from "./tool_call_credentials";
2323 import { getGraphqlDecision, resolveToolForCall } from "./tool_call_resolution";
24+import { sourceSignature } from "./tool_source_loading";
2425
2526 function createApprovalId(): string {
2627 return `approval_${crypto.randomUUID()}`;
@@ -61,12 +62,18 @@ export async function invokeTool(ctx: ActionCtx, task: TaskRecord, call: ToolCal · 34 unmodified lines
6162 // Fast system tools are handled server-side from the registry.
6263 if (toolPath === "discover" || toolPath === "catalog.namespaces" || toolPath === "catalog.tools") {
6364 const toolRegistry = internal.toolRegistry;
64− const state = await ctx.runQuery(toolRegistry.getState, { workspaceId: task.workspaceId }) as
65− | null
66− | { readyBuildId?: string };
65+ const [state, sources] = await Promise.all([
66+ ctx.runQuery(toolRegistry.getState, { workspaceId: task.workspaceId }) as Promise<null | { signature: string; readyBuildId?: string }>,
67+ ctx.runQuery(internal.database.listToolSources, { workspaceId: task.workspaceId }) as Promise<Array<{ id: string; updatedAt: number; enabled: boolean }>>,
68+ ]);
69+ const enabledSources = sources.filter((source) => source.enabled);
70+ const signature = sourceSignature(task.workspaceId, enabledSources);
71+ const expectedSignature = `toolreg_v1|${signature}`;
6772 const buildId = state?.readyBuildId;
68− if (!buildId) {
69− throw new Error("Tool registry is not ready yet. Open Tools in the UI to build the registry.");
73+ if (!buildId || state.signature !== expectedSignature) {
74+ throw new Error(
75+ "Tool registry is not ready (or is stale). Open Tools to refresh, or call listToolsWithWarnings to rebuild.",
76+ );
7077 }
7178
7279 const payload = input && typeof input === "object" ? (input as Record<string, unknown>) : {};
executor/packages/convex/runtime/workspace_tools.ts+39−10
@@ -12,7 +12,7 @@ import { · 11 unmodified lines
1212 type CompiledToolSourceArtifact,
1313 type WorkspaceToolSnapshot,
1414 } from "../../core/src/tool-sources";
15−import { serializeTools } from "../../core/src/tool/source-serialization";
15+import type { SerializedTool } from "../../core/src/tool/source-serialization";
1616 import type { ExternalToolSourceConfig } from "../../core/src/tool/source-types";
1717 import type {
1818 AccessPolicyRecord,
@@ -370,20 +370,18 @@ async function buildWorkspaceToolRegistry( · 351 unmodified lines
370370 ctx: ActionCtx,
371371 args: {
372372 workspaceId: Id<"workspaces">;
373− signature: string;
374− tools: ToolDefinition[];
373+ registrySignature: string;
374+ serializedTools: SerializedTool[];
375375 },
376376 ): Promise<{ buildId: string }> {
377377 const buildId = `toolreg_${crypto.randomUUID()}`;
378378 await ctx.runMutation(internal.toolRegistry.beginBuild, {
379379 workspaceId: args.workspaceId,
380− signature: args.signature,
380+ signature: args.registrySignature,
381381 buildId,
382382 });
383383
384− const serialized = serializeTools(args.tools);
385−
386− const entries = serialized.map((st) => {
384+ const entries = args.serializedTools.map((st) => {
387385 if (st.path === "discover" || st.path.startsWith("catalog.")) {
388386 return null;
389387 }
@@ -487,6 +485,25 @@ async function buildWorkspaceToolRegistry( · 97 unmodified lines
487485 return { buildId };
488486 }
489487
488+async function maybeBuildWorkspaceToolRegistry(
489+ ctx: ActionCtx,
490+ args: {
491+ workspaceId: Id<"workspaces">;
492+ registrySignature: string;
493+ serializedTools: SerializedTool[];
494+ },
495+): Promise<void> {
496+ const state = await ctx.runQuery(internal.toolRegistry.getState, { workspaceId: args.workspaceId }) as
497+ | null
498+ | { signature: string; readyBuildId?: string };
499+
500+ if (state?.readyBuildId && state.signature === args.registrySignature) {
501+ return;
502+ }
503+
504+ await buildWorkspaceToolRegistry(ctx, args);
505+}
506+
490507 export async function getWorkspaceTools(
491508 ctx: ActionCtx,
492509 workspaceId: Id<"workspaces">,
@@ -517,6 +534,7 @@ export async function getWorkspaceTools( · 24 unmodified lines
517534 traceStep("listToolSources", listSourcesStartedAt);
518535 const hasOpenApiSource = sources.some((source: { type: string }) => source.type === "openapi");
519536 const signature = sourceSignature(workspaceId, sources);
537+ const registrySignature = `toolreg_v1|${signature}`;
520538 const debugBase: Omit<WorkspaceToolsDebug, "mode" | "normalizedSourceCount" | "cacheHit" | "cacheFresh" | "timedOutSources" | "durationMs" | "trace"> = {
521539 includeDts,
522540 sourceTimeoutMs: sourceTimeoutMs ?? null,
@@ -545,6 +563,17 @@ export async function getWorkspaceTools( · 22 unmodified lines
545563 const typesStorageId = cacheEntry.typesStorageId as Id<"_storage"> | undefined;
546564 if (cacheEntry.isFresh) {
547565 if (typesStorageId) {
566+ // Registry is derived from the cached snapshot. Keep it in sync.
567+ try {
568+ await maybeBuildWorkspaceToolRegistry(ctx, {
569+ workspaceId,
570+ registrySignature,
571+ serializedTools: snapshot.externalArtifacts.flatMap((artifact) => artifact.tools),
572+ });
573+ } catch (error) {
574+ const msg = error instanceof Error ? error.message : String(error);
575+ console.warn(`[executor] workspace tool registry build failed for '${workspaceId}': ${msg}`);
576+ }
548577 return {
549578 tools: merged,
550579 warnings: snapshot.warnings,
@@ -677,10 +706,10 @@ export async function getWorkspaceTools( · 126 unmodified lines
677706
678707 // Build a per-tool registry for fast discover + invocation.
679708 const registryStartedAt = Date.now();
680− await buildWorkspaceToolRegistry(ctx, {
709+ await maybeBuildWorkspaceToolRegistry(ctx, {
681710 workspaceId,
682− signature,
683− tools: allTools,
711+ registrySignature,
712+ serializedTools: externalArtifacts.flatMap((artifact) => artifact.tools),
684713 });
685714 traceStep("toolRegistryWrite", registryStartedAt);
686715