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, ctx.signal);
       });
     },
   };

controlplane/src/core/services/PromptToQueryService.ts

@@ -1,3 +1,4 @@
+import { setTimeout } from 'node:timers/promises';
 import { PostgresJsDatabase } from 'drizzle-orm/postgres-js';
 import { FastifyBaseLogger } from 'fastify';
 import { type AxiosInstance, create as createHttpClient } from 'axios';
@@ -12,13 +13,24 @@ import { EnumStatusCode } from '@wundergraph/cosmo-connect/dist/common/common_pb
 import * as z from 'zod';
 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().uuid(),
   prompt: z.string().trim().min(1),
 });
 
+const indexPollInterval = 1000;
+
+const indexResponseSchema = z.object({
+  index: z.object({
+    indexId: z.string().trim().min(1),
+    status: z.enum(['INDEX_STATUS_INDEXING', 'INDEX_STATUS_READY', 'INDEX_STATUS_FAILED']),
+    error: z.string().optional(),
+  }),
+});
+
 const ptqQuerySchema = z.object({
   description: z.string().optional(),
   document: z.string().min(1),
@@ -65,7 +77,12 @@ export class PromptToQueryService {
     });
   }
 
-  async generateQuery(schemaSha: string, prompt: string): Promise<GenerateQueryResponse> {
+  async generateQuery(
+    federatedGraphId: string,
+    version: string,
+    prompt: string,
+    signal?: AbortSignal,
+  ): Promise<GenerateQueryResponse> {
     if (!this.serviceAddress) {
       // The feature doesn't seem to be configured correctly
       return create(GenerateQueryResponseSchema, {
@@ -77,7 +94,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,15 +115,41 @@ 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, signal);
       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,
         }),
+        signal,
       });
 
       const parsedResponse = ptqResponseSchema.safeParse(response.data);
@@ -152,6 +195,43 @@ export class PromptToQueryService {
     }).catch((e) => this.logger.error(e, 'Failed to index schema due an unexpected error'));
   }
 
+  private async ensureIndex(schemaSDL: string, signal?: AbortSignal): Promise<string> {
+    const response = await this.#httpClient('/yoko.v1.YokoService/EnsureIndex', {
+      method: 'POST',
+      headers: { 'Content-Type': 'application/json' },
+      data: JSON.stringify({ sdl: schemaSDL }),
+      signal,
+    });
+
+    let index = PromptToQueryService.parseIndexResponse(response.data);
+    while (index.status === 'INDEX_STATUS_INDEXING') {
+      await setTimeout(indexPollInterval, undefined, { signal });
+
+      const response = await this.#httpClient('/yoko.v1.YokoService/GetIndex', {
+        method: 'POST',
+        headers: { 'Content-Type': 'application/json' },
+        data: JSON.stringify({ indexId: index.indexId }),
+        signal,
+      });
+      index = PromptToQueryService.parseIndexResponse(response.data);
+    }
+
+    if (index.status === 'INDEX_STATUS_FAILED') {
+      throw new Error(`Prompt to Query index generation failed${index.error ? `: ${index.error}` : ''}`);
+    }
+
+    return index.indexId;
+  }
+
+  private static parseIndexResponse(response: unknown): z.infer<typeof indexResponseSchema>['index'] {
+    const parsed = indexResponseSchema.safeParse(response);
+    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.index;
+  }
+
   private static getOperationType(type: z.infer<typeof ptqQuerySchema>['operationType']): SatisfiedOperationType {
     switch (type) {
       case 'query': {

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
 

docs-website/router/mcp/tools.mdx

@@ -16,6 +16,7 @@ The MCP server gives AI models a set of tools they can discover and execute. It
 | `get_operation_info` | Returns instructions for executing the operation behind one of your tools directly via HTTP and integrating it into an application. |
 | `get_schema`         | Returns the full GraphQL schema of the API, helping AI models understand the entire API structure.                                  |
 | `execute_graphql`    | Executes an arbitrary GraphQL query or mutation, letting AI models craft operations beyond the tools you have created.              |
+| `generate_query`     | Uses Cosmo Cloud Prompt to Query to generate a GraphQL operation from a natural-language prompt.                                    |
 
 <Warning>
   `get_schema` and `execute_graphql` are disabled by default because they expose your full API surface to AI models.
@@ -24,6 +25,35 @@ The MCP server gives AI models a set of tools they can discover and execute. It
   intended. Prefer creating focused tools.
 </Warning>
 
+### Generate a GraphQL Operation with Cosmo Cloud
+
+The `generate_query` tool is available when the router has a graph token and router registration is enabled. It sends the
+prompt and a hash of the router's active schema to Cosmo Cloud's Prompt to Query service. The generated operation is
+returned to the MCP client but is not executed.
+
+The tool accepts one required field:
+
+```json
+{
+  "prompt": "List the names and email addresses of the five most recently created users"
+}
+```
+
+A successful response contains the GraphQL document and the metadata needed to use it:
+
+```json
+{
+  "description": "Lists the five most recently created users",
+  "document": "query RecentUsers { users(first: 5, orderBy: CREATED_AT_DESC) { name email } }",
+  "operationName": "RecentUsers",
+  "operationType": "query",
+  "variablesSchema": "{\"type\":\"object\"}"
+}
+```
+
+Availability also depends on the Prompt to Query feature being enabled for the Cosmo Cloud organization associated with
+the graph token. Control-plane or entitlement errors are returned as MCP tool errors.
+
 ## Creating Tools
 
 Create a directory for your tools (as specified in your [storage provider configuration](/router/mcp/configuration#storage-providers)) and add `.graphql` or `.gql` files containing GraphQL operations.

proto/wg/cosmo/node/v1/node.proto

@@ -61,7 +61,8 @@ message SelfRegisterResponse {
 }
 
 message GenerateQueryRequest {
-  string schema_hash = 1;
+  // Copied from RouterConfig.version in the active router configuration.
+  string version = 1;
   string prompt = 2;
 }
 

router/core/graph_server.go

@@ -1611,7 +1611,7 @@ func (s *graphServer) buildGraphMux(
 
 	// We support the MCP only on the base graph. Feature flags are not supported yet.
 	if opts.IsBaseGraph() && s.mcpServer != nil {
-		if mErr := s.mcpServer.Reload(executor.ClientSchema, opts.EngineConfig.FieldConfigurations); mErr != nil {
+		if mErr := s.mcpServer.Reload(executor.ClientSchema, opts.EngineConfig.FieldConfigurations, opts.RouterConfigVersion); mErr != nil {
 			return nil, fmt.Errorf("failed to reload MCP server: %w", mErr)
 		}
 	}

router/core/router.go

@@ -1218,6 +1218,9 @@ func (r *Router) startMCPServer(ctx context.Context) error {
 		mcpserver.WithServerTitle(r.mcp.Server.Title),
 		mcpserver.WithServerDescription(r.mcp.Server.Description),
 	}
+	if r.promptToQueryClient != nil {
+		mcpOpts = append(mcpOpts, mcpserver.WithPromptToQueryClient(r.promptToQueryClient))
+	}
 
 	if r.corsOptions != nil {
 		mcpOpts = append(mcpOpts, mcpserver.WithCORS(*r.corsOptions))
@@ -2089,6 +2092,13 @@ func WithSelfRegistration(sr selfregister.SelfRegister) Option {
 	}
 }
 
+// WithPromptToQueryClient sets the control-plane client used by the MCP generate_query tool.
+func WithPromptToQueryClient(client mcpserver.PromptToQueryClient) Option {
+	return func(r *Router) {
+		r.promptToQueryClient = client
+	}
+}
+
 // WithGracePeriod sets the grace period for the router to shutdown.
 func WithGracePeriod(timeout time.Duration) Option {
 	return func(r *Router) {

router/core/router_config.go

@@ -143,6 +143,7 @@ type Config struct {
 	// Poller
 	configPoller                 configpoller.ConfigPoller
 	selfRegister                 selfregister.SelfRegister
+	promptToQueryClient          mcpserver.PromptToQueryClient
 	registrationInfo             *nodev1.RegistrationInfo
 	securityConfiguration        config.SecurityConfiguration
 	customModules                []Module

router/core/supervisor_instance.go

@@ -9,6 +9,8 @@ import (
 
 	"github.com/KimMachineGun/automemlimit/memlimit"
 	"github.com/dustin/go-humanize"
+	"github.com/wundergraph/cosmo/router/internal/controlplane"
+	"github.com/wundergraph/cosmo/router/internal/prompttoquery"
 	"github.com/wundergraph/cosmo/router/pkg/authentication"
 	"github.com/wundergraph/cosmo/router/pkg/config"
 	"github.com/wundergraph/cosmo/router/pkg/controlplane/selfregister"
@@ -159,6 +161,19 @@ func newRouter(ctx context.Context, params RouterResources, additionalOptions ..
 		options = append(options, WithSelfRegistration(selfRegister))
 	}
 
+	if cfg.MCP.Enabled && cfg.Graph.Token != "" {
+		controlplaneTransport, err := controlplane.NewTransport(cfg.Graph.Token, logger)
+		if err != nil {
+			return nil, fmt.Errorf("could not create controlplane transport: %w", err)
+		}
+
+		promptToQueryClient, err := prompttoquery.New(cfg.ControlplaneURL, controlplaneTransport)
+		if err != nil {
+			return nil, fmt.Errorf("could not create prompt-to-query client: %w", err)
+		}
+		options = append(options, WithPromptToQueryClient(promptToQueryClient))
+	}
+
 	if opt := optionFromExecutionConfig(&cfg.ExecutionConfig, cfg.RouterConfigPath); opt != nil {
 		options = append(options, opt)
 	} else {