History

fix(v2): stream tool listings via paged index reads

Rhys Sullivan committed 8 months ago

Commit 9607da41 parent 4be71f43 files · +297 −6Browse files

Unified · Split

3 files
  1. apps/convex/convex/control_plane/service.ts +41 −6
  2. apps/convex/convex/control_plane/tool_registry.ts +247 −0
  3. apps/convex/convex/control_plane/tools.ts +9 −0
apps/convex/convex/control_plane/service.ts+41−6
@@ -2,12 +2,50 @@ import { · 1 unmodified line
22 ControlPlaneService,
33 type ControlPlaneServiceShape,
44 } from "@executor-v2/management-api";
5+import type { SourceToolSummary } from "@executor-v2/management-api/tools/api";
56 import * as Effect from "effect/Effect";
67
7−import { api } from "../_generated/api";
8+import { api, internal } from "../_generated/api";
89 import type { ActionCtx } from "../_generated/server";
910 import { toSourceStoreError } from "./errors";
1011
12+const runtimeInternal = internal as any;
13+const toolPageSize = 250;
14+
15+const collectWorkspaceToolSummaries = async (
16+ ctx: ActionCtx,
17+ input: {
18+ workspaceId: string;
19+ sourceId?: string;
20+ },
21+): Promise<Array<SourceToolSummary>> => {
22+ await ctx.runAction(runtimeInternal.control_plane.tool_registry.ensureWorkspaceToolIndexCoverage, {
23+ workspaceId: input.workspaceId,
24+ ...(input.sourceId ? { sourceId: input.sourceId } : {}),
25+ });
26+
27+ const summaries: Array<SourceToolSummary> = [];
28+ let cursor: string | null = null;
29+
30+ do {
31+ const page: {
32+ summaries: Array<SourceToolSummary>;
33+ continueCursor: string | null;
34+ isDone: boolean;
35+ } = await ctx.runQuery(runtimeInternal.control_plane.tools.listWorkspaceToolSummaryPage, {
36+ workspaceId: input.workspaceId,
37+ ...(input.sourceId ? { sourceId: input.sourceId } : {}),
38+ cursor,
39+ pageSize: toolPageSize,
40+ });
41+
42+ summaries.push(...page.summaries);
43+ cursor = page.isDone ? null : page.continueCursor;
44+ } while (cursor);
45+
46+ return summaries;
47+};
48+
1149 export const makeConvexControlPlaneService = (
1250 ctx: ActionCtx,
1351 ): ControlPlaneServiceShape =>
@@ -121,16 +159,13 @@ export const makeConvexControlPlaneService = ( · 107 unmodified lines
121159 }),
122160 listWorkspaceTools: (workspaceId) =>
123161 Effect.tryPromise({
124− try: () =>
125− ctx.runAction(api.controlPlane.listWorkspaceTools, {
126− workspaceId,
127− }),
162+ try: () => collectWorkspaceToolSummaries(ctx, { workspaceId }),
128163 catch: (cause) => toSourceStoreError("controlPlane.listWorkspaceTools", cause),
129164 }),
130165 listSourceTools: (input) =>
131166 Effect.tryPromise({
132167 try: () =>
133− ctx.runAction(api.controlPlane.listSourceTools, {
168+ collectWorkspaceToolSummaries(ctx, {
134169 workspaceId: input.workspaceId,
135170 sourceId: input.sourceId,
136171 }),
apps/convex/convex/control_plane/tool_registry.ts+247−0
@@ -506,6 +506,52 @@ export const ingestOpenApiManifest = action({ · 505 unmodified lines
506506 artifactId: artifactMeta.artifactId,
507507 });
508508
509+ const namespace = sourceNamespace({ id: args.sourceId, name: args.sourceName });
510+ await ctx.runMutation(runtimeInternal.control_plane.tool_registry.clearWorkspaceSourceToolIndex, {
511+ workspaceId: args.workspaceId,
512+ sourceId: args.sourceId,
513+ });
514+
515+ for (const batch of chunkArray(args.tools, writeBatchSize)) {
516+ await ctx.runMutation(runtimeInternal.control_plane.tool_registry.appendWorkspaceSourceToolIndexChunk, {
517+ workspaceId: args.workspaceId,
518+ sourceId: args.sourceId,
519+ sourceName: args.sourceName,
520+ sourceKind: "openapi",
521+ artifactId: artifactMeta.artifactId,
522+ namespace,
523+ refHintTableJson: null,
524+ insertOnly: true,
525+ rows: batch.map((tool) => {
526+ const operationPath = `${tool.method.toUpperCase()} ${tool.path}`;
527+ const path = `${namespace}.${tool.toolId}`;
528+
529+ return {
530+ toolId: tool.toolId,
531+ protocol: "openapi",
532+ method: tool.method.toLowerCase(),
533+ path,
534+ operationPath,
535+ name: tool.name,
536+ description: tool.description,
537+ searchText: normalizedSearchText(
538+ args.sourceName,
539+ args.sourceEndpoint,
540+ namespace,
541+ tool.toolId,
542+ tool.name,
543+ tool.description,
544+ tool.method,
545+ tool.path,
546+ ),
547+ operationHash: tool.operationHash,
548+ approvalMode: "auto",
549+ status: args.sourceEnabled ? "active" : "disabled",
550+ };
551+ }),
552+ });
553+ }
554+
509555 return {
510556 artifactId: artifactMeta.artifactId,
511557 created: artifactMeta.created,
@@ -1085,6 +1131,20 @@ export const getSourceArtifactBinding = internalQuery({ · 573 unmodified lines
10851131 },
10861132 });
10871133
1134+export const listWorkspaceSourceArtifactBindings = internalQuery({
1135+ args: {
1136+ workspaceId: v.string(),
1137+ },
1138+ handler: async (ctx, args) => {
1139+ const rows = await ctx.db
1140+ .query("sourceArtifactBindings")
1141+ .withIndex("by_workspaceId", (q) => q.eq("workspaceId", args.workspaceId))
1142+ .collect();
1143+
1144+ return rows.map((row) => stripConvexSystemFields(row as unknown as Record<string, unknown>));
1145+ },
1146+});
1147+
10881148 export const getArtifactById = internalQuery({
10891149 args: {
10901150 artifactId: v.string(),
@@ -1115,6 +1175,189 @@ export const listArtifactTools = internalQuery({ · 24 unmodified lines
11151175 },
11161176 });
11171177
1178+export const listArtifactToolsPage = internalQuery({
1179+ args: {
1180+ artifactId: v.string(),
1181+ cursor: v.optional(v.union(v.string(), v.null())),
1182+ pageSize: v.optional(v.number()),
1183+ },
1184+ handler: async (ctx, args) => {
1185+ const pageSize = Math.max(1, Math.min(500, Math.floor(args.pageSize ?? 200)));
1186+ const page = await ctx.db
1187+ .query("artifactTools")
1188+ .withIndex("by_artifactId", (q) => q.eq("artifactId", args.artifactId))
1189+ .paginate({
1190+ cursor: args.cursor ?? null,
1191+ numItems: pageSize,
1192+ });
1193+
1194+ const rows = page.page.map((row) => stripConvexSystemFields(row as unknown as Record<string, unknown>));
1195+
1196+ return {
1197+ rows,
1198+ continueCursor: page.continueCursor,
1199+ isDone: page.isDone,
1200+ };
1201+ },
1202+});
1203+
1204+export const hasWorkspaceSourceToolIndexRows = internalQuery({
1205+ args: {
1206+ workspaceId: v.string(),
1207+ sourceId: v.string(),
1208+ },
1209+ handler: async (ctx, args) => {
1210+ const page = await ctx.db
1211+ .query("workspaceToolIndex")
1212+ .withIndex("by_workspaceId_sourceId", (q) =>
1213+ q.eq("workspaceId", args.workspaceId).eq("sourceId", args.sourceId),
1214+ )
1215+ .paginate({
1216+ cursor: null,
1217+ numItems: 1,
1218+ });
1219+
1220+ return page.page.length > 0;
1221+ },
1222+});
1223+
1224+export const ensureWorkspaceToolIndexCoverage = internalAction({
1225+ args: {
1226+ workspaceId: v.string(),
1227+ sourceId: v.optional(v.string()),
1228+ },
1229+ handler: async (ctx, args) => {
1230+ const sourceRows = await ctx.runQuery(runtimeInternal.control_plane.sources.listSources, {
1231+ workspaceId: args.workspaceId,
1232+ });
1233+ const sourceById = new Map<string, Record<string, unknown>>(
1234+ (sourceRows as Array<Record<string, unknown>>).map((row) => [String(row.id), row] as const),
1235+ );
1236+
1237+ const bindingRows = await ctx.runQuery(
1238+ runtimeInternal.control_plane.tool_registry.listWorkspaceSourceArtifactBindings,
1239+ {
1240+ workspaceId: args.workspaceId,
1241+ },
1242+ );
1243+
1244+ let rebuiltCount = 0;
1245+
1246+ for (const binding of bindingRows) {
1247+ const sourceId = String(binding.sourceId ?? "");
1248+ if (sourceId.length === 0) {
1249+ continue;
1250+ }
1251+
1252+ if (args.sourceId && sourceId !== args.sourceId) {
1253+ continue;
1254+ }
1255+
1256+ const source = sourceById.get(sourceId);
1257+ if (!source) {
1258+ continue;
1259+ }
1260+
1261+ const hasRows = await ctx.runQuery(runtimeInternal.control_plane.tool_registry.hasWorkspaceSourceToolIndexRows, {
1262+ workspaceId: args.workspaceId,
1263+ sourceId,
1264+ });
1265+ if (hasRows) {
1266+ continue;
1267+ }
1268+
1269+ const sourceName = String(source.name ?? "");
1270+ const sourceKind = String(source.kind ?? "openapi");
1271+ const sourceEndpoint = String(source.endpoint ?? "");
1272+ const sourceEnabled = source.enabled === true;
1273+ const artifactId = String(binding.artifactId ?? "");
1274+ if (artifactId.length === 0) {
1275+ continue;
1276+ }
1277+
1278+ const artifact = await ctx.runQuery(runtimeInternal.control_plane.tool_registry.getArtifactById, {
1279+ artifactId,
1280+ });
1281+ const namespace = sourceNamespace({ id: sourceId, name: sourceName });
1282+
1283+ await ctx.runMutation(runtimeInternal.control_plane.tool_registry.clearWorkspaceSourceToolIndex, {
1284+ workspaceId: args.workspaceId,
1285+ sourceId,
1286+ });
1287+
1288+ let cursor: string | null = null;
1289+ do {
1290+ const page: {
1291+ rows: Array<Record<string, unknown>>;
1292+ continueCursor: string | null;
1293+ isDone: boolean;
1294+ } = await ctx.runQuery(runtimeInternal.control_plane.tool_registry.listArtifactToolsPage, {
1295+ artifactId,
1296+ cursor,
1297+ pageSize: 200,
1298+ });
1299+
1300+ if (page.rows.length > 0) {
1301+ await ctx.runMutation(runtimeInternal.control_plane.tool_registry.appendWorkspaceSourceToolIndexChunk, {
1302+ workspaceId: args.workspaceId,
1303+ sourceId,
1304+ sourceName,
1305+ sourceKind,
1306+ artifactId,
1307+ namespace,
1308+ refHintTableJson: typeof artifact?.refHintTableJson === "string" ? artifact.refHintTableJson : null,
1309+ insertOnly: true,
1310+ rows: page.rows.map((row) => {
1311+ const toolId = String(row.toolId ?? "");
1312+ const protocol = String(row.protocol ?? "openapi");
1313+ const canonicalPath =
1314+ typeof row.canonicalPath === "string" && row.canonicalPath.trim().length > 0
1315+ ? row.canonicalPath.trim()
1316+ : null;
1317+ const method = toWorkspaceToolMethod(
1318+ protocol,
1319+ typeof row.metadataJson === "string" ? row.metadataJson : null,
1320+ canonicalPath,
1321+ );
1322+ const path = `${namespace}.${toolId}`;
1323+ const description = typeof row.description === "string" ? row.description : null;
1324+ const name = String(row.name ?? toolId);
1325+
1326+ return {
1327+ toolId,
1328+ protocol,
1329+ method,
1330+ path,
1331+ operationPath: canonicalPath,
1332+ name,
1333+ description,
1334+ searchText: normalizedSearchText(
1335+ sourceName,
1336+ sourceEndpoint,
1337+ namespace,
1338+ toolId,
1339+ name,
1340+ description,
1341+ ...metadataSearchTerms(typeof row.metadataJson === "string" ? row.metadataJson : null),
1342+ ),
1343+ operationHash: String(row.operationHash ?? ""),
1344+ approvalMode: "auto",
1345+ status: sourceEnabled ? "active" : "disabled",
1346+ };
1347+ }),
1348+ });
1349+ }
1350+
1351+ cursor = page.isDone ? null : page.continueCursor;
1352+ } while (cursor);
1353+
1354+ rebuiltCount += 1;
1355+ }
1356+
1357+ return { rebuiltCount };
1358+ },
1359+});
1360+
11181361 export const listArtifactToolIndexSlice = internalQuery({
11191362 args: {
11201363 artifactId: v.string(),
@@ -1292,6 +1535,10 @@ export const listWorkspaceNamespacesUncached = internalAction({ · 171 unmodified lines
12921535 limit: v.optional(v.number()),
12931536 },
12941537 handler: async (ctx, args) => {
1538+ await ctx.runAction(runtimeInternal.control_plane.tool_registry.ensureWorkspaceToolIndexCoverage, {
1539+ workspaceId: args.workspaceId,
1540+ });
1541+
12951542 const grouped = new Map<
12961543 string,
12971544 {
apps/convex/convex/control_plane/tools.ts+9−0
@@ -303,6 +303,10 @@ export const listWorkspaceTools = action({ · 302 unmodified lines
303303 workspaceId: v.string(),
304304 },
305305 handler: async (ctx, args): Promise<Array<SourceToolSummary>> => {
306+ await ctx.runAction(runtimeInternal.control_plane.tool_registry.ensureWorkspaceToolIndexCoverage, {
307+ workspaceId: args.workspaceId,
308+ });
309+
306310 const summaries: Array<SourceToolSummary> = [];
307311 let cursor: string | null = null;
308312
@@ -331,6 +335,11 @@ export const listSourceTools = action({ · 22 unmodified lines
331335 sourceId: v.string(),
332336 },
333337 handler: async (ctx, args): Promise<Array<SourceToolSummary>> => {
338+ await ctx.runAction(runtimeInternal.control_plane.tool_registry.ensureWorkspaceToolIndexCoverage, {
339+ workspaceId: args.workspaceId,
340+ sourceId: args.sourceId,
341+ });
342+
334343 const summaries: Array<SourceToolSummary> = [];
335344 let cursor: string | null = null;
336345