feat(router): expose prompt-to-query via MCP by fiam · Pull Request #3158 · wundergraph/cosmo

fiam · 2026-08-12T17:39:25Z

Summary

expose a generate_query tool through the router MCP server
send the active RouterConfig.version to the control-plane GenerateQuery RPC
resolve that version to the stored SDL in the control plane, call Yoko EnsureIndex on demand, and use its returned opaque index ID
add OAuth scope support and MCP documentation for the tool

Stack

Implements ROUTER-575
Depends on feat: implement ptq rpc and index schemas #3138
This PR is intentionally based on the head branch of feat: implement ptq rpc and index schemas #3138. Once feat: implement ptq rpc and index schemas #3138 merges, this PR can be retargeted to main.
Prompt-to-query service behavior remains covered by the base PR; this stack focuses on the router/MCP integration.

Verification

CGO_ENABLED=0 go test ./... in router
control-plane TypeScript build

coderabbitai · 2026-08-12T17:39:33Z

Important

Review skipped

Auto reviews are disabled on base/target branches other than the default branch.

Please check the settings in the CodeRabbit UI or the .coderabbit.yaml file in this repository. To trigger a single review, invoke the @coderabbitai review command.

⚙️ Run configuration

Configuration used: Organization UI

Review profile: CHILL

Plan: Pro

Run ID: bbbe0658-a9f1-4bc2-a882-387cca981824

You can disable this status message by setting the reviews.review_status to false in the CodeRabbit configuration file.

Use the checkbox below for a quick retry:

🔍 Trigger review

_{Comment @coderabbitai help to get the list of available commands.}

github-actions · 2026-08-12T17:41:45Z

❌ Internal Query Planner CI checks failed

The Internal Query Planner CI checks failed in the celestial repository, and this is going to stop the merge of this PR.
If you are part of the WunderGraph organization, you can see the PR with more details.

codecov · 2026-08-12T18:49:58Z

Codecov Report

❌ Patch coverage is 36.57143% with 111 lines in your changes missing coverage. Please review.
✅ Project coverage is 43.96%. Comparing base (c9d6d43) to head (064d7d2).

Files with missing lines	Patch %	Lines
router/pkg/controlplane/prompttoquery/client.go	0.00%	44 Missing ⚠️
...rolplane/src/core/services/PromptToQueryService.ts	10.00%	27 Missing ⚠️
router/pkg/mcpserver/server.go	71.08%	16 Missing and 8 partials ⚠️
router/core/supervisor_instance.go	0.00%	7 Missing ⚠️
router/core/router.go	0.00%	4 Missing and 1 partial ⚠️
router/gen/proto/wg/cosmo/node/v1/node.pb.go	0.00%	2 Missing ⚠️
controlplane/src/core/bufservices/NodeService.ts	0.00%	1 Missing ⚠️
router/core/graph_server.go	0.00%	0 Missing and 1 partial ⚠️

Additional details and impacted files

@@                                       Coverage Diff                                        @@
##           wilson/cosmo-335-cosmo-cloud-prompt-to-query-via-mcp-product    #3158      +/-   ##
================================================================================================
+ Coverage                                                         43.85%   43.96%   +0.10%     
================================================================================================
  Files                                                              1072     1078       +6     
  Lines                                                            140577   141078     +501     
  Branches                                                           7361     7446      +85     
================================================================================================
+ Hits                                                              61655    62020     +365     
- Misses                                                            77071    77190     +119     
- Partials                                                           1851     1868      +17

Files with missing lines	Coverage Δ
router/core/router_config.go	`93.97% <ø> (ø)`
router/pkg/config/config.go	`83.00% <ø> (ø)`
router/pkg/mcpserver/auth_middleware.go	`90.69% <100.00%> (+0.10%)`	⬆️
controlplane/src/core/bufservices/NodeService.ts	`19.17% <0.00%> (ø)`
router/core/graph_server.go	`85.76% <0.00%> (ø)`
router/gen/proto/wg/cosmo/node/v1/node.pb.go	`28.69% <0.00%> (+1.13%)`	⬆️
router/core/router.go	`70.76% <0.00%> (-0.23%)`	⬇️
router/core/supervisor_instance.go	`0.00% <0.00%> (ø)`
router/pkg/mcpserver/server.go	`74.54% <71.08%> (+3.98%)`	⬆️
...rolplane/src/core/services/PromptToQueryService.ts	`27.86% <10.00%> (-4.41%)`	⬇️
... and 1 more

... and 23 files with indirect coverage changes

🚀 New features to boost your workflow:

❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.
📦 JS Bundle Analysis: Save yourself from yourself by tracking and limiting bundle sizes in JS merges.

wilsonrivera · 2026-08-12T19:01:20Z

+
    // Invoke the `prompt to query` service
    try {
+      const indexId = await this.ensureIndex(schemaVersion.sdl);


This method is not intended to be called by the query generator. The reason for this is that the service could take a bit to parse and index the schema, use the crypto to generate the SHA256 for the schema and use the value directly.

If the index doesn't exist, let it fail

Makes sense! PTAL and let me know if this is what you had in mind. Thanks!

ensureIndex is implemented via a task queue so it should actually respond instantly with most fields, but importantly query generation will error if the index is not completed so you may want to poll here or fail fast.

In the context of Yoko that is correct, however, this is just a helper for ControlPlane to send the schema for indexation. If we want to check whether the index is already there, we need to make a new helper

endigma · 2026-08-14T09:33:52Z

+	if cfg.MCP.Enabled && cfg.Graph.Token != "" {
+		promptToQueryClient, err := prompttoquery.New(cfg.ControlplaneURL, cfg.Graph.Token,
+			prompttoquery.WithLogger(logger),
+		)
+		if err != nil {
+			return nil, fmt.Errorf("could not create prompt-to-query client: %w", err)
+		}
+		options = append(options, WithPromptToQueryClient(promptToQueryClient))
+	}


Unsure about this, see my PM

endigma · 2026-08-14T09:33:53Z

+
+const defaultTimeout = 15 * time.Second
+
+type Option func(*Client)


Is this worth it for just a logger and timeout vs named parameters to New(...)? This is internal-only, the default is either used or it isn't

I agree, probably a bit overkill. Hardcoded the timeout and made the logger a required argument. I've also moved the package to be internal so we don't have to care about a stable interface b7f1b0f - PTAL

endigma · 2026-08-14T09:33:53Z

+	retryClient := retryablehttp.NewClient()
+	retryClient.RetryWaitMax = 15 * time.Second
+	retryClient.RetryMax = 3
+	retryClient.Backoff = retryablehttp.DefaultBackoff
+	retryClient.Logger = nil
+	retryClient.RequestLogHook = func(_ retryablehttp.Logger, _ *http.Request, retry int) {
+		if retry > 0 {
+			c.logger.Info("Generate query through controlplane", zap.Int("retry", retry))
+		}
+	}


Need to ensure that this does not retry on errors that should not be retryable

The default policy retries on 429 and >= 500 (except 501). Looking at the CP side, it seems we only return 5xx for unhandled errors, all application errors return 200.

So the only thing that could generate a retry is either an intermediate proxy/infra or either an unhandled exception (which might or might not be worth retrying). However, some errors in k8s ingress or some service meshes also end up generating a 500 when the service is unavailable (e.g. some rollout strategy not properly configured).

Since prompt-to-query is supposed to be idempotent I think it might be safer to over-retry than under-retry (if that makes sense). WDYT?

endigma · 2026-08-14T09:33:53Z

+		Version: schemaVersionID,
+		Prompt:  prompt,
+	})
+	req.Header().Set("Authorization", "Bearer "+c.graphAPIToken)


I believe this is better done as part of a roundtripper on the client itself

I think that's a cool idea, but as an internal package, so eventually all communication to the control plane could use the same transport. PTAL 561f7e2

endigma · 2026-08-14T09:33:54Z

+					"prompt": map[string]any{
+						"type":        "string",
+						"minLength":   1,
+						"description": "A natural-language description of the GraphQL operation to generate.",


I believe this may need a more descriptive instruction for how to format the prompt, tbd

Replace the functional options with an explicit logger constructor parameter and keep the fixed client timeout internal. Move the package under router/internal since it is only consumed by router core.

Move bearer authentication and retry configuration into a shared internal transport. Inject it into prompt-to-query from router setup.

Ensure the schema index exists before generating a query and poll Yoko until the index is ready. Propagate request cancellation through indexing and generation so polling stops when the caller disconnects or reaches its deadline.

claude

Claude Code Review

Claude Code Review is paused for this repository. To reconnect it, an admin of this repository's GitHub organization (or the account owner, for personal repositories) who can also manage your Claude organization's Code Review settings needs to re-link GitHub in Code Review settings. This is a one-time step.

_{Tip: disable this comment in your organization's Code Review settings.}

fiam added 4 commits August 12, 2026 17:37

wip: expose prompt-to-query through MCP

b3a786e

feat: resolve prompt-to-query indexes by schema version

7fe0fa1

refactor: align generate query version field

fd817ff

refactor: resolve prompt schema inside service

189412b

github-actions Bot added controlplane protocol connect router docs-website labels Aug 12, 2026

fiam added 4 commits August 12, 2026 19:44

style(router): clean up prompt-to-query client

25cbbee

test(router): remove prompt-to-query client test

4fb96f9

test(controlplane): remove stacked prompt-to-query tests

116737f

test(router): remove generate-query error cases

c310389

chore(router): remove local prompt-to-query script

155e2d9

wilsonrivera reviewed Aug 12, 2026

View reviewed changes

Comment thread controlplane/src/core/services/PromptToQueryService.ts Outdated

wilsonrivera reviewed Aug 12, 2026

View reviewed changes

wilsonrivera requested changes Aug 12, 2026

View reviewed changes

fiam added 2 commits August 12, 2026 20:13

refactor(controlplane): use zod UUID validation

4d2f733

fix(controlplane): avoid indexing schemas during query generation

064d7d2

fiam marked this pull request as ready for review August 13, 2026 08:56

fiam requested review from a team as code owners August 13, 2026 08:56

endigma requested changes Aug 14, 2026

View reviewed changes

fiam force-pushed the alberto/router-575-router-mcp-query-generation-support branch from ee3a30e to 665d0b9 Compare August 17, 2026 21:44

refactor(router): simplify prompt-to-query client

b7f1b0f

Replace the functional options with an explicit logger constructor parameter and keep the fixed client timeout internal. Move the package under router/internal since it is only consumed by router core.

fiam force-pushed the alberto/router-575-router-mcp-query-generation-support branch from 665d0b9 to b7f1b0f Compare August 17, 2026 21:45

refactor(router): centralize control plane transport

561f7e2

Move bearer authentication and retry configuration into a shared internal transport. Inject it into prompt-to-query from router setup.

fiam force-pushed the alberto/router-575-router-mcp-query-generation-support branch from 94b0d9a to 561f7e2 Compare August 17, 2026 22:32

fix(controlplane): wait for prompt-to-query index

8c0591b

Ensure the schema index exists before generating a query and poll Yoko until the index is ready. Propagate request cancellation through indexing and generation so polling stops when the caller disconnects or reaches its deadline.

claude Bot reviewed Aug 18, 2026

View reviewed changes


		const defaultTimeout = 15 * time.Second

		type Option func(*Client)

Uh oh!

Conversation

fiam commented Aug 12, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary

Stack

Verification

Uh oh!

coderabbitai Bot commented Aug 12, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Review skipped

Uh oh!

github-actions Bot commented Aug 12, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

❌ Internal Query Planner CI checks failed

Uh oh!

codecov Bot commented Aug 12, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Codecov Report

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

claude Bot left a comment

Choose a reason for hiding this comment

Claude Code Review

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

fiam commented Aug 12, 2026 •

edited

Loading

coderabbitai Bot commented Aug 12, 2026 •

edited

Loading

github-actions Bot commented Aug 12, 2026 •

edited

Loading

codecov Bot commented Aug 12, 2026 •

edited

Loading