feat(router): expose prompt-to-query via MCP - #3158
Conversation
|
Important Review skippedAuto reviews are disabled on base/target branches other than the default branch. Please check the settings in the CodeRabbit UI or the ⚙️ Run configurationConfiguration used: Organization UI Review profile: CHILL Plan: Pro Run ID: You can disable this status message by setting the Use the checkbox below for a quick retry:
Comment |
❌ Internal Query Planner CI checks failedThe Internal Query Planner CI checks failed in the celestial repository, and this is going to stop the merge of this PR. |
Codecov Report❌ Patch coverage is 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
🚀 New features to boost your workflow:
|
|
|
||
| // Invoke the `prompt to query` service | ||
| try { | ||
| const indexId = await this.ensureIndex(schemaVersion.sdl); |
There was a problem hiding this comment.
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
There was a problem hiding this comment.
Makes sense! PTAL and let me know if this is what you had in mind. Thanks!
There was a problem hiding this comment.
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.
There was a problem hiding this comment.
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
| 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)) | ||
| } |
|
|
||
| const defaultTimeout = 15 * time.Second | ||
|
|
||
| type Option func(*Client) |
There was a problem hiding this comment.
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
There was a problem hiding this comment.
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
| 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)) | ||
| } | ||
| } |
There was a problem hiding this comment.
Need to ensure that this does not retry on errors that should not be retryable
There was a problem hiding this comment.
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?
| Version: schemaVersionID, | ||
| Prompt: prompt, | ||
| }) | ||
| req.Header().Set("Authorization", "Bearer "+c.graphAPIToken) |
There was a problem hiding this comment.
I believe this is better done as part of a roundtripper on the client itself
There was a problem hiding this comment.
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
| "prompt": map[string]any{ | ||
| "type": "string", | ||
| "minLength": 1, | ||
| "description": "A natural-language description of the GraphQL operation to generate.", |
There was a problem hiding this comment.
I believe this may need a more descriptive instruction for how to format the prompt, tbd
94b0d9a to
561f7e2
Compare
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
There was a problem hiding this comment.
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.
Move bearer authentication and retry configuration into a shared internal transport so Control Plane clients can receive it from router setup.
Add an internal Control Plane client for AI query generation using the active schema version. Reuse the injected authenticated transport and bound requests with a private 15-second timeout.
Expose an optional OAuth scope for the generate_query built-in tool and enforce it through MCP authorization. Update the router configuration schema and fixtures accordingly.
Expose a natural-language GraphQL query-generation tool when an AI client is injected. Track the active router schema version across reloads and return generated operation details through MCP.
Create the Control Plane query-generation client only when MCP is enabled and the graph token includes prompt-to-query. Pass the gated client into the MCP server so the tool remains absent otherwise.
Document the generate_query MCP tool, its feature and token requirements, response format, and optional OAuth scope. Note that existing graph tokens must be rotated after organization enablement.
8c0591b to
330ce7d
Compare
3 participants
Summary
generate_querytool through the router MCP serverRouterConfig.versionto the control-planeGenerateQueryRPCEnsureIndexon demand, and use its returned opaque index IDStack
main.Verification
CGO_ENABLED=0 go test ./...inrouter