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
2 ControlPlaneService,2 ControlPlaneService,
3 type ControlPlaneServiceShape,3 type ControlPlaneServiceShape,
4} from "@executor-v2/management-api";4} from "@executor-v2/management-api";
5import type { SourceToolSummary } from "@executor-v2/management-api/tools/api";
5import * as Effect from "effect/Effect";6import * as Effect from "effect/Effect";
6 7
7import { api } from "../_generated/api";8import { api, internal } from "../_generated/api";
8import type { ActionCtx } from "../_generated/server";9import type { ActionCtx } from "../_generated/server";
9import { toSourceStoreError } from "./errors";10import { toSourceStoreError } from "./errors";
10 11
12const runtimeInternal = internal as any;
13const toolPageSize = 250;
14
15const 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
11export const makeConvexControlPlaneService = (49export const makeConvexControlPlaneService = (
12 ctx: ActionCtx,50 ctx: ActionCtx,
13): ControlPlaneServiceShape =>51): ControlPlaneServiceShape =>
@@ -121,16 +159,13 @@ export const makeConvexControlPlaneService = ( · 107 unmodified lines
121 }),159 }),
122 listWorkspaceTools: (workspaceId) =>160 listWorkspaceTools: (workspaceId) =>
123 Effect.tryPromise({161 Effect.tryPromise({
124 try: () =>162 try: () => collectWorkspaceToolSummaries(ctx, { workspaceId }),
125 ctx.runAction(api.controlPlane.listWorkspaceTools, {
126 workspaceId,
127 }),
128 catch: (cause) => toSourceStoreError("controlPlane.listWorkspaceTools", cause),163 catch: (cause) => toSourceStoreError("controlPlane.listWorkspaceTools", cause),
129 }),164 }),
130 listSourceTools: (input) =>165 listSourceTools: (input) =>
131 Effect.tryPromise({166 Effect.tryPromise({
132 try: () =>167 try: () =>
133 ctx.runAction(api.controlPlane.listSourceTools, {168 collectWorkspaceToolSummaries(ctx, {
134 workspaceId: input.workspaceId,169 workspaceId: input.workspaceId,
135 sourceId: input.sourceId,170 sourceId: input.sourceId,
136 }),171 }),
apps/convex/convex/control_plane/tool_registry.ts+247−0
@@ -506,6 +506,52 @@ export const ingestOpenApiManifest = action({ · 505 unmodified lines
506 artifactId: artifactMeta.artifactId,506 artifactId: artifactMeta.artifactId,
507 });507 });
508 508
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
509 return {555 return {
510 artifactId: artifactMeta.artifactId,556 artifactId: artifactMeta.artifactId,
511 created: artifactMeta.created,557 created: artifactMeta.created,
@@ -1085,6 +1131,20 @@ export const getSourceArtifactBinding = internalQuery({ · 573 unmodified lines
1085 },1131 },
1086});1132});
1087 1133
1134export 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
1088export const getArtifactById = internalQuery({1148export const getArtifactById = internalQuery({
1089 args: {1149 args: {
1090 artifactId: v.string(),1150 artifactId: v.string(),
@@ -1115,6 +1175,189 @@ export const listArtifactTools = internalQuery({ · 24 unmodified lines
1115 },1175 },
1116});1176});
1117 1177
1178export 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
1204export 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
1224export 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
1118export const listArtifactToolIndexSlice = internalQuery({1361export const listArtifactToolIndexSlice = internalQuery({
1119 args: {1362 args: {
1120 artifactId: v.string(),1363 artifactId: v.string(),
@@ -1292,6 +1535,10 @@ export const listWorkspaceNamespacesUncached = internalAction({ · 171 unmodified lines
1292 limit: v.optional(v.number()),1535 limit: v.optional(v.number()),
1293 },1536 },
1294 handler: async (ctx, args) => {1537 handler: async (ctx, args) => {
1538 await ctx.runAction(runtimeInternal.control_plane.tool_registry.ensureWorkspaceToolIndexCoverage, {
1539 workspaceId: args.workspaceId,
1540 });
1541
1295 const grouped = new Map<1542 const grouped = new Map<
1296 string,1543 string,
1297 {1544 {
apps/convex/convex/control_plane/tools.ts+9−0
@@ -303,6 +303,10 @@ export const listWorkspaceTools = action({ · 302 unmodified lines
303 workspaceId: v.string(),303 workspaceId: v.string(),
304 },304 },
305 handler: async (ctx, args): Promise<Array<SourceToolSummary>> => {305 handler: async (ctx, args): Promise<Array<SourceToolSummary>> => {
306 await ctx.runAction(runtimeInternal.control_plane.tool_registry.ensureWorkspaceToolIndexCoverage, {
307 workspaceId: args.workspaceId,
308 });
309
306 const summaries: Array<SourceToolSummary> = [];310 const summaries: Array<SourceToolSummary> = [];
307 let cursor: string | null = null;311 let cursor: string | null = null;
308 312
@@ -331,6 +335,11 @@ export const listSourceTools = action({ · 22 unmodified lines
331 sourceId: v.string(),335 sourceId: v.string(),
332 },336 },
333 handler: async (ctx, args): Promise<Array<SourceToolSummary>> => {337 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
334 const summaries: Array<SourceToolSummary> = [];343 const summaries: Array<SourceToolSummary> = [];
335 let cursor: string | null = null;344 let cursor: string | null = null;
336 345