feat(router): expose prompt-to-query via MCP

controlplane/src/core/bufservices/NodeService.ts

@@ -82,7 +82,7 @@ export default function (opts: RouterOptions): Partial<ServiceImpl<typeof NodeSe
           opts.billingDefaultPlanId,
         );
 
-        return ptqService.generateQuery(req.schemaHash, req.prompt);
+        return ptqService.generateQuery(authContext.federatedGraphId, req.version, req.prompt);
       });
     },
   };

controlplane/src/core/services/PromptToQueryService.ts

@@ -10,15 +10,21 @@ import {
 import { create } from '@bufbuild/protobuf';
 import { EnumStatusCode } from '@wundergraph/cosmo-connect/dist/common/common_pb';
 import * as z from 'zod';
+import { validate as validateUUID } from 'uuid';
 import { traced } from '../tracing.js';
 import * as schema from '../../db/schema.js';
+import { FederatedGraphRepository } from '../repositories/FederatedGraphRepository.js';
 import { OrganizationRepository } from '../repositories/OrganizationRepository.js';
 
 const validationSchema = z.object({
-  schemaSha: z.string().regex(/^sha256:[\da-f]{64}$/i),
+  version: z.string().refine(validateUUID),
   prompt: z.string().trim().min(1),
 });
 
+const ensureIndexResponseSchema = z.object({
+  indexId: z.string().trim().min(1),
+});
+
 const ptqQuerySchema = z.object({
   description: z.string().optional(),
   document: z.string().min(1),
@@ -65,7 +71,7 @@ export class PromptToQueryService {
     });
   }
 
-  async generateQuery(schemaSha: string, prompt: string): Promise<GenerateQueryResponse> {
+  async generateQuery(federatedGraphId: string, version: string, prompt: string): Promise<GenerateQueryResponse> {
     if (!this.serviceAddress) {
       // The feature doesn't seem to be configured correctly
       return create(GenerateQueryResponseSchema, {
@@ -77,7 +83,7 @@ export class PromptToQueryService {
     }
 
     // Ensure that the provided parameters are valid
-    const parsed = validationSchema.safeParse({ schemaSha, prompt });
+    const parsed = validationSchema.safeParse({ version, prompt });
     if (!parsed.success) {
       return create(GenerateQueryResponseSchema, {
         response: {
@@ -98,13 +104,38 @@ export class PromptToQueryService {
       });
     }
 
+    const fedRepo = new FederatedGraphRepository(this.logger, this.db, this.organizationId);
+    const federatedGraph = await fedRepo.byId(federatedGraphId);
+    if (!federatedGraph) {
+      return create(GenerateQueryResponseSchema, {
+        response: {
+          code: EnumStatusCode.ERR_NOT_FOUND,
+          details: 'Federated graph not found',
+        },
+      });
+    }
+
+    const schemaVersion = await fedRepo.getSdlBasedOnSchemaVersion({
+      targetId: federatedGraph.targetId,
+      schemaVersionId: parsed.data.version,
+    });
+    if (!schemaVersion?.sdl) {
+      return create(GenerateQueryResponseSchema, {
+        response: {
+          code: EnumStatusCode.ERR_NOT_FOUND,
+          details: 'Schema version not found for this federated graph',
+        },
+      });
+    }
+
     // Invoke the `prompt to query` service
     try {
+      const indexId = await this.ensureIndex(schemaVersion.sdl);
       const response = await this.#httpClient('/yoko.v1.YokoService/GenerateQuery', {
         method: 'POST',
         headers: { 'Content-Type': 'application/json' },
         data: JSON.stringify({
-          indexId: parsed.data.schemaSha,
+          indexId,
           prompt: parsed.data.prompt,
         }),
       });
@@ -145,11 +176,21 @@ export class PromptToQueryService {
     }
 
     // Fire and forget the schema indexation
-    this.#httpClient('/yoko.v1.YokoService/EnsureIndex', {
+    this.ensureIndex(schema).catch((e) => this.logger.error(e, 'Failed to index schema due an unexpected error'));
+  }
+
+  private async ensureIndex(schemaSDL: string): Promise<string> {
+    const response = await this.#httpClient('/yoko.v1.YokoService/EnsureIndex', {
       method: 'POST',
       headers: { 'Content-Type': 'application/json' },
-      data: JSON.stringify({ sdl: schema }),
-    }).catch((e) => this.logger.error(e, 'Failed to index schema due an unexpected error'));
+      data: JSON.stringify({ sdl: schemaSDL }),
+    });
+    const parsed = ensureIndexResponseSchema.safeParse(response.data);
+    if (!parsed.success) {
+      throw new Error('It was not possible to parse the response returned by the Prompt to Query index service');
+    }
+
+    return parsed.data.indexId;
   }
 
   private static getOperationType(type: z.infer<typeof ptqQuerySchema>['operationType']): SatisfiedOperationType {

controlplane/test/prompt-to-query-service.test.ts

@@ -0,0 +1,90 @@
+import { PostgresJsDatabase } from 'drizzle-orm/postgres-js';
+import { pino } from 'pino';
+import { http, HttpResponse } from 'msw';
+import { setupServer } from 'msw/node';
+import { afterAll, afterEach, beforeAll, describe, expect, test, vi } from 'vitest';
+import { EnumStatusCode } from '@wundergraph/cosmo-connect/dist/common/common_pb';
+import { FederatedGraphRepository } from '../src/core/repositories/FederatedGraphRepository.js';
+import { OrganizationRepository } from '../src/core/repositories/OrganizationRepository.js';
+import { PromptToQueryService } from '../src/core/services/PromptToQueryService.js';
+import * as schema from '../src/db/schema.js';
+import type { FederatedGraphDTO } from '../src/types/index.js';
+
+describe('PromptToQueryService', () => {
+  const yokoURL = 'http://yoko.test';
+  const mockServer = setupServer();
+
+  beforeAll(() => mockServer.listen({ onUnhandledRequest: 'error' }));
+  afterEach(() => {
+    mockServer.resetHandlers();
+    vi.restoreAllMocks();
+  });
+  afterAll(() => mockServer.close());
+
+  test('uses the index ID returned by EnsureIndex to generate the query', async () => {
+    const schemaSDL = 'type Query { employees: [String!]! }';
+    let ensureIndexRequest: unknown;
+    let generateQueryRequest: unknown;
+
+    mockServer.use(
+      http.post(`${yokoURL}/yoko.v1.YokoService/EnsureIndex`, async ({ request }) => {
+        ensureIndexRequest = await request.json();
+        return HttpResponse.json({ indexId: 'opaque-yoko-index-id' });
+      }),
+      http.post(`${yokoURL}/yoko.v1.YokoService/GenerateQuery`, async ({ request }) => {
+        generateQueryRequest = await request.json();
+        return HttpResponse.json({
+          resolution: {
+            queries: [
+              {
+                description: 'Lists employees',
+                document: 'query ListEmployees { employees }',
+                operationName: 'ListEmployees',
+                operationType: 'query',
+                variablesSchema: '{}',
+              },
+            ],
+            unsatisfied: [],
+          },
+        });
+      }),
+    );
+    vi.spyOn(OrganizationRepository.prototype, 'getFeature').mockResolvedValue({
+      id: 'prompt-to-query',
+      enabled: true,
+    });
+    const byId = vi
+      .spyOn(FederatedGraphRepository.prototype, 'byId')
+      .mockResolvedValue({ targetId: 'target-id' } as FederatedGraphDTO);
+    const getSdl = vi.spyOn(FederatedGraphRepository.prototype, 'getSdlBasedOnSchemaVersion').mockResolvedValue({
+      sdl: schemaSDL,
+      clientSchema: schemaSDL,
+    });
+
+    const service = new PromptToQueryService(
+      {} as PostgresJsDatabase<typeof schema>,
+      pino(),
+      yokoURL,
+      'organization-id',
+      undefined,
+    );
+    const response = await service.generateQuery(
+      'graph-id',
+      '14a1d197-7e3a-48df-88d7-a663de90527e',
+      'List all employees',
+    );
+
+    expect(ensureIndexRequest).toEqual({ sdl: schemaSDL });
+    expect(byId).toHaveBeenCalledWith('graph-id');
+    expect(getSdl).toHaveBeenCalledWith({
+      targetId: 'target-id',
+      schemaVersionId: '14a1d197-7e3a-48df-88d7-a663de90527e',
+    });
+    expect(generateQueryRequest).toEqual({
+      indexId: 'opaque-yoko-index-id',
+      prompt: 'List all employees',
+    });
+    expect(response.response?.code).toBe(EnumStatusCode.OK);
+    expect(response.query?.operationName).toBe('ListEmployees');
+  });
+});

controlplane/test/prompt-to-query.test.ts

@@ -0,0 +1,101 @@
+import { http, HttpResponse } from 'msw';
+import { setupServer } from 'msw/node';
+import { afterAll, afterEach, beforeAll, describe, expect, test } from 'vitest';
+import { EnumStatusCode } from '@wundergraph/cosmo-connect/dist/common/common_pb';
+import { joinLabel, routerConfigFromJsonString } from '@wundergraph/cosmo-shared';
+import { afterAllSetup, beforeAllSetup, genID, genUniqueLabel } from '../src/core/test-util.js';
+import {
+  createFederatedGraph,
+  createThenPublishSubgraph,
+  DEFAULT_NAMESPACE,
+  DEFAULT_ROUTER_URL,
+  DEFAULT_SUBGRAPH_URL_ONE,
+  SetupTest,
+} from './test-util.js';
+
+let dbname = '';
+
+describe('Prompt to Query RPC', () => {
+  const yokoURL = 'http://yoko.test';
+  const mockServer = setupServer();
+
+  beforeAll(async () => {
+    mockServer.listen({ onUnhandledRequest: 'bypass' });
+    dbname = await beforeAllSetup();
+  });
+  afterEach(() => mockServer.resetHandlers());
+  afterAll(async () => {
+    mockServer.close();
+    await afterAllSetup(dbname);
+  });
+
+  test("resolves the router schema version and generates a query with Yoko's index ID", async (testContext) => {
+    const ensureIndexRequests: unknown[] = [];
+    let generateQueryRequest: unknown;
+    mockServer.use(
+      http.post(`${yokoURL}/yoko.v1.YokoService/EnsureIndex`, async ({ request }) => {
+        ensureIndexRequests.push(await request.json());
+        return HttpResponse.json({ indexId: 'opaque-yoko-index-id' });
+      }),
+      http.post(`${yokoURL}/yoko.v1.YokoService/GenerateQuery`, async ({ request }) => {
+        generateQueryRequest = await request.json();
+        return HttpResponse.json({
+          resolution: {
+            queries: [
+              {
+                description: 'Returns hello',
+                document: 'query GetHello { hello }',
+                operationName: 'GetHello',
+                operationType: 'query',
+                variablesSchema: '{}',
+              },
+            ],
+            unsatisfied: [],
+          },
+        });
+      }),
+    );
+
+    const { client, nodeClient, server, users, blobStorage } = await SetupTest({
+      dbname,
+      enabledFeatures: ['prompt-to-query'],
+      promptToQueryServiceAddress: yokoURL,
+    });
+    testContext.onTestFinished(() => server.close());
+
+    const subgraphName = genID('subgraph');
+    const graphName = genID('graph');
+    const label = genUniqueLabel();
+    await createThenPublishSubgraph(
+      client,
+      subgraphName,
+      DEFAULT_NAMESPACE,
+      'type Query { hello: String! }',
+      [label],
+      DEFAULT_SUBGRAPH_URL_ONE,
+    );
+    await createFederatedGraph(client, graphName, DEFAULT_NAMESPACE, [joinLabel(label)], DEFAULT_ROUTER_URL);
+
+    const graphResponse = await client.getFederatedGraphByName({ name: graphName, namespace: DEFAULT_NAMESPACE });
+    expect(graphResponse.response?.code).toBe(EnumStatusCode.OK);
+    const tokenResponse = await client.generateRouterToken({ fedGraphName: graphName, namespace: DEFAULT_NAMESPACE });
+    expect(tokenResponse.response?.code).toBe(EnumStatusCode.OK);
+
+    const configBlob = await blobStorage.getObject({
+      key: `${users.adminAliceCompanyA.organizationId}/${graphResponse.graph?.id}/routerconfigs/latest.json`,
+    });
+    const config = routerConfigFromJsonString(await new Response(configBlob.stream).text());
+    const response = await nodeClient.generateQuery(
+      {
+        version: config?.version,
+        prompt: 'Return hello',
+      },
+      { headers: { Authorization: `Bearer ${tokenResponse.token}` } },
+    );
+
+    expect(response.response?.code).toBe(EnumStatusCode.OK);
+    expect(response.query?.operationName).toBe('GetHello');
+    expect(ensureIndexRequests.at(-1)).toEqual({ sdl: config?.engineConfig?.graphqlSchema });
+    expect(generateQueryRequest).toEqual({ indexId: 'opaque-yoko-index-id', prompt: 'Return hello' });
+  });
+});

controlplane/test/test-util.ts

@@ -86,6 +86,7 @@ export const SetupTest = async function ({
   createScimKey,
   setupBilling,
   organizationId,
+  promptToQueryServiceAddress,
 }: {
   dbname: string;
   chClient?: ClickHouseClient;
@@ -96,6 +97,7 @@ export const SetupTest = async function ({
     plan: 'developer@1' | 'launch@1' | 'scale@1' | 'enterprise';
   };
   organizationId?: UUID;
+  promptToQueryServiceAddress?: string;
 }) {
   const log = pino();
   const databaseConnectionUrl = `postgresql://postgres:changeme@localhost:5432/${dbname}`;
@@ -173,6 +175,7 @@ export const SetupTest = async function ({
         clientSecret: 'test',
       },
       cdnBaseUrl: 'http://localhost:11000',
+      promptToQueryServiceAddress,
       admissionWebhookJWTSecret: 'secret',
       keycloakApiUrl: apiUrl,
       blobStorage,

docs-website/router/mcp/oauth/configuration.mdx

@@ -18,6 +18,7 @@ icon: 'sliders-up'
 | `oauth.scopes.execute_graphql`               | Scopes required to call the `execute_graphql` built-in tool. Additive to `tools_call`. Only relevant when `enable_arbitrary_operations` is `true`.                                                                            | `[]`    |
 | `oauth.scopes.get_operation_info`            | Scopes required to call the `get_operation_info` built-in tool. Additive to `tools_call`.                                                                                                                                     | `[]`    |
 | `oauth.scopes.get_schema`                    | Scopes required to call the `get_schema` built-in tool. Additive to `tools_call`. Only relevant when `expose_schema` is `true`.                                                                                               | `[]`    |
+| `oauth.scopes.generate_query`                | Scopes required to call the `generate_query` built-in tool. Additive to `tools_call`. Only relevant when a graph-token-backed control-plane client is configured.                                                            | `[]`    |
 | `oauth.jwks`                                 | List of JWKS providers for JWT verification. Supports remote JWKS URLs or symmetric secrets.                                                                                                                                  | `[]`    |
 
 ## JWKS Configuration
@@ -67,6 +68,7 @@ oauth:
 | `MCP_OAUTH_AUTHORIZATION_SERVER_URL`             | `mcp.oauth.authorization_server_url`             |
 | `MCP_OAUTH_SCOPE_CHALLENGE_INCLUDE_TOKEN_SCOPES` | `mcp.oauth.scope_challenge_include_token_scopes` |
 | `MCP_OAUTH_MAX_SCOPE_COMBINATIONS`               | `mcp.oauth.max_scope_combinations`               |
+| `MCP_OAUTH_SCOPES_GENERATE_QUERY`                 | `mcp.oauth.scopes.generate_query`                 |
 
 ## HTTP Error Responses