fix: refresh MCP source tool manifests (#862)

Jad committed 5 months ago

Commit 5ae2ce81 parent bd61b6e2 files · +145 −88Browse files

Unified · Split

2 files

apps/cloud/src/services/sources-api.node.test.ts +48 −0
packages/plugins/mcp/src/sdk/plugin.ts +97 −88

apps/cloud/src/services/sources-api.node.test.ts+48−0

@@ -457,6 +457,54 @@ describe("sources api (HTTP)", () => { · 456 unmodified lines
457	}),	457	}),
458	);	458	);
459		459
		460	it.effect("generic source refresh updates MCP source tool rows", () =>
		461	Effect.gen(function* () {
		462	let toolName = "before_refresh";
		463	const server = yield* serveMcpServer(() =>
		464	makeGreetingMcpServer({
		465	name: "cloud-refresh-mcp",
		466	toolName,
		467	text: "refresh-ok",
		468	}),
		469	);
		470	const org = `org_${crypto.randomUUID()}`;
		471	const namespace = `mcp_${crypto.randomUUID().replace(/-/g, "_")}`;
		472	const scopeId = ScopeId.make(org);
		473
		474	const added = yield* asOrg(org, (client) =>
		475	client.mcp.addSource({
		476	params: { scopeId },
		477	payload: {
		478	transport: "remote",
		479	name: "Cloud Refresh MCP",
		480	endpoint: server.endpoint,
		481	remoteTransport: "streamable-http",
		482	namespace,
		483	},
		484	}),
		485	);
		486	expect(added).toEqual({ namespace, toolCount: 1 });
		487
		488	const beforeTools = yield* asOrg(org, (client) =>
		489	client.sources.tools({ params: { scopeId, sourceId: namespace } }),
		490	);
		491	expect(beforeTools.map((tool) => tool.id)).toContain(`${namespace}.before_refresh`);
		492	expect(beforeTools.map((tool) => tool.id)).not.toContain(`${namespace}.after_refresh`);
		493
		494	toolName = "after_refresh";
		495	const refreshResult = yield* asOrg(org, (client) =>
		496	client.sources.refresh({ params: { scopeId, sourceId: namespace } }),
		497	);
		498	expect(refreshResult.refreshed).toBe(true);
		499
		500	const afterTools = yield* asOrg(org, (client) =>
		501	client.sources.tools({ params: { scopeId, sourceId: namespace } }),
		502	);
		503	expect(afterTools.map((tool) => tool.id)).not.toContain(`${namespace}.before_refresh`);
		504	expect(afterTools.map((tool) => tool.id)).toContain(`${namespace}.after_refresh`);
		505	}),
		506	);
		507
460	it.effect("sources.remove deletes the source and it drops off sources.list", () =>	508	it.effect("sources.remove deletes the source and it drops off sources.list", () =>
461	Effect.gen(function* () {	509	Effect.gen(function* () {
462	const org = `org_${crypto.randomUUID()}`;	510	const org = `org_${crypto.randomUUID()}`;

packages/plugins/mcp/src/sdk/plugin.ts+97−88

   return entry;
 };
+const refreshSourceInternal = (
+  ctx: PluginCtx<McpBindingStore>,
+  sourceId: string,
+  scope: string,
+  allowStdio: boolean,
+) =>
+  Effect.gen(function* () {
+    const sd = yield* ctx.storage.getSourceConfig(sourceId, scope).pipe(
+      Effect.withSpan("mcp.plugin.load_source_config", {
+        attributes: { "mcp.source.namespace": sourceId },
+      }),
+    );
+    if (!sd) {
+      return yield* new McpConnectionError({
+        transport: "remote",
+        message: `No stored config for MCP source "${sourceId}"`,
+      });
+    }
+    const ci = yield* resolveConnectorInput(sourceId, scope, sd, ctx, allowStdio).pipe(
+      Effect.catchTag("McpAuthRequiredError", ({ message }) =>
+        Effect.fail(new McpConnectionError({ transport: sd.transport, message })),
+      ),
+      Effect.withSpan("mcp.plugin.resolve_connector", {
+        attributes: {
+          "mcp.source.namespace": sourceId,
+          "mcp.source.transport": sd.transport,
+        },
+      }),
+    );
+    const manifest = yield* discoverTools(createMcpConnector(ci)).pipe(
+      Effect.mapError(
+        ({ message }) =>
+          new McpToolDiscoveryError({
+            stage: "list_tools",
+            message: `MCP refresh failed: ${message}`,
+          }),
+      ),
+      Effect.withSpan("mcp.plugin.discover_tools", {
+        attributes: { "mcp.source.namespace": sourceId },
+      }),
+    );
+    const existing = yield* ctx.storage.getSource(sourceId, scope);
+    const sourceName = manifest.server?.name ?? existing?.name ?? sourceId;
+    yield* ctx
+      .transaction(
+        Effect.gen(function* () {
+          yield* ctx.storage.removeBindingsByNamespace(sourceId, scope);
+          yield* ctx.core.sources.unregister({ id: sourceId, targetScope: scope });
+          yield* ctx.storage.putBindings(
+            sourceId,
+            scope,
+            manifest.tools.map((e) => ({
+              toolId: `${sourceId}.${e.toolId}`,
+              binding: toBinding(e),
+            })),
+          );
+          yield* ctx.core.sources.register({
+            id: sourceId,
+            scope,
+            kind: "mcp",
+            name: sourceName,
+            url: sd.transport === "remote" ? sd.endpoint : undefined,
+            canRemove: true,
+            canRefresh: true,
+            canEdit: sd.transport === "remote",
+            tools: manifest.tools.map((e) => ({
+              name: e.toolId,
+              description: e.description ?? `MCP tool: ${e.toolName}`,
+              inputSchema: e.inputSchema,
+              outputSchema: mcpCallToolResultOutputSchema(e.outputSchema),
+            })),
+          });
+        }),
+      )
+      .pipe(
+        Effect.withSpan("mcp.plugin.persist_source", {
+          attributes: {
+            "mcp.source.namespace": sourceId,
+            "mcp.source.tool_count": manifest.tools.length,
+          },
+        }),
+      );
+    return { toolCount: manifest.tools.length };
+  }).pipe(
+    Effect.withSpan("mcp.plugin.refresh_source", {
+      attributes: { "mcp.source.namespace": sourceId },
+    }),
+  );
 export const mcpPlugin = definePlugin((options?: McpPluginOptions) => {
   const allowStdio = options?.dangerouslyAllowStdioMCP ?? false;
   // Per-plugin-instance runtime holder. Captured by closures in
         );
       const refreshSource = (namespace: string, scope: string) =>
-        Effect.gen(function* () {
+        refreshSourceInternal(ctx, namespace, scope, allowStdio);
-          const sd = yield* ctx.storage.getSourceConfig(namespace, scope).pipe(
-            Effect.withSpan("mcp.plugin.load_source_config", {
-              attributes: { "mcp.source.namespace": namespace },
-            }),
-          );
-          if (!sd) {
-            return yield* new McpConnectionError({
-              transport: "remote",
-              message: `No stored config for MCP source "${namespace}"`,
-            });
-          }
-          const ci = yield* resolveConnectorInput(namespace, scope, sd, ctx, allowStdio).pipe(
-            Effect.catchTag("McpAuthRequiredError", ({ message }) =>
-              Effect.fail(new McpConnectionError({ transport: sd.transport, message })),
-            ),
-            Effect.withSpan("mcp.plugin.resolve_connector", {
-              attributes: {
-                "mcp.source.namespace": namespace,
-                "mcp.source.transport": sd.transport,
-              },
-            }),
-          );
-          const manifest = yield* discoverTools(createMcpConnector(ci)).pipe(
-            Effect.mapError(
-              ({ message }) =>
-                new McpToolDiscoveryError({
-                  stage: "list_tools",
-                  message: `MCP refresh failed: ${message}`,
-                }),
-            ),
-            Effect.withSpan("mcp.plugin.discover_tools", {
-              attributes: { "mcp.source.namespace": namespace },
-            }),
-          );
-          const existing = yield* ctx.storage.getSource(namespace, scope);
-          const sourceName = manifest.server?.name ?? existing?.name ?? namespace;
-          yield* ctx
-            .transaction(
-              Effect.gen(function* () {
-                yield* ctx.storage.removeBindingsByNamespace(namespace, scope);
-                yield* ctx.core.sources.unregister({ id: namespace, targetScope: scope });
-                yield* ctx.storage.putBindings(
-                  namespace,
-                  scope,
-                  manifest.tools.map((e) => ({
-                    toolId: `${namespace}.${e.toolId}`,
-                    binding: toBinding(e),
-                  })),
-                );
-                yield* ctx.core.sources.register({
-                  id: namespace,
-                  scope,
-                  kind: "mcp",
-                  name: sourceName,
-                  url: sd.transport === "remote" ? sd.endpoint : undefined,
-                  canRemove: true,
-                  canRefresh: true,
-                  canEdit: sd.transport === "remote",
-                  tools: manifest.tools.map((e) => ({
-                    name: e.toolId,
-                    description: e.description ?? `MCP tool: ${e.toolName}`,
-                    inputSchema: e.inputSchema,
-                    outputSchema: mcpCallToolResultOutputSchema(e.outputSchema),
-                  })),
-                });
-              }),
-            )
-            .pipe(
-              Effect.withSpan("mcp.plugin.persist_source", {
-                attributes: {
-                  "mcp.source.namespace": namespace,
-                  "mcp.source.tool_count": manifest.tools.length,
-                },
-              }),
-            );
-          return { toolCount: manifest.tools.length };
-        }).pipe(
-          Effect.withSpan("mcp.plugin.refresh_source", {
-            attributes: { "mcp.source.namespace": namespace },
-          }),
-        );
       const getSource = (namespace: string, scope: string) =>
         ctx.storage.getSource(namespace, scope).pipe(
     usagesForConnection: () => Effect.succeed([]),
-    refreshSource: () => Effect.void,
+    refreshSource: ({ ctx, sourceId, scope }) =>
+      refreshSourceInternal(ctx, sourceId, scope, allowStdio).pipe(Effect.asVoid),
     // Connection refresh for oauth2-minted sources is owned by the
     // canonical `"oauth2"` ConnectionProvider that core registers via