feat(router): add MCP schema discovery for search and query generation [DO NOT MERGE] - #3156
Conversation
|
Preview deployment for your docs. Learn more about Mintlify Previews.
💡 Tip: Enable Workflows to automatically generate PRs for you. |
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Organization UI Review profile: CHILL Plan: Pro Run ID: ⛔ Files ignored due to path filters (2)
📒 Files selected for processing (1)
WalkthroughThis change adds schema discovery to the router MCP server. It defines the Yoko protobuf API, indexes schemas through an external service, exposes three MCP tools, adds configuration and startup integration, and documents setup, workflows, and tool behavior. ChangesSchema discovery
Estimated code review effort: 4 (Complex) | ~60 minutes Possibly related PRs
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
Comment |
This comment has been minimized.
This comment has been minimized.
Router-nonroot image scan passed✅ No security vulnerabilities found in image: |
There was a problem hiding this comment.
Actionable comments posted: 7
🧹 Nitpick comments (5)
router/pkg/querygen/service.go (1)
234-260: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueSoften the guarantee in the doc comment.
The comment states that the text never names the token or any other secret. The default branch at line 245 appends
err.Error()from an arbitrary transport error, and the branch at line 258 appendsconnectErr.Message()from the remote service. Neither string is under the control of this package. State that the router never adds the token to the message.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@router/pkg/querygen/service.go` around lines 234 - 260, Update the doc comment for UserMessage to state that the router never adds the token to the returned message, rather than guaranteeing that no secret can appear. Leave the existing error-message branches and behavior unchanged.router/pkg/querygen/config.go (1)
43-48: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winValidate the URL scheme in
Validate.The doc comment on
URLstates that the value must contain a scheme, butValidateonly rejects an empty string. A value such aslocalhost:3400passes validation. The Connect client then fails on the first call with an unclear transport error instead of at router startup.♻️ Proposed check
func (c *Config) Validate() error { if c.URL == "" { return errors.New("schema discovery is enabled but no url is configured") } + u, err := url.Parse(c.URL) + if err != nil { + return fmt.Errorf("schema discovery url %q is not a valid url: %w", c.URL, err) + } + if u.Scheme != "http" && u.Scheme != "https" { + return fmt.Errorf("schema discovery url %q must use the http or https scheme", c.URL) + } return nil }Add
fmtandnet/urlto the imports.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@router/pkg/querygen/config.go` around lines 43 - 48, Update Config.Validate to parse c.URL with net/url and reject values without a scheme, while retaining the existing empty-URL error. Add the required fmt and net/url imports and return a clear validation error that identifies the invalid URL and requires a scheme.router/pkg/querygen/querygen_test.go (1)
359-418: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAdd coverage for the
Resyncpath and the address-mismatch path.The suite covers indexing, polling, timeout, failure, search, and generation. Two behaviours stay untested:
Service.handlecallsindexer.Resynconconnect.CodeNotFound. No test asserts that a secondEnsureIndexcall follows an expired index.buildreassignswantwhen the service returns a differentindex_id. That path leavesi.pendingstale, as flagged onrouter/pkg/querygen/indexer.go. A test would pin the fixed behaviour.I can generate both tests if you want.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@router/pkg/querygen/querygen_test.go` around lines 359 - 418, Extend the querygen test suite with coverage for Service.handle’s connect.CodeNotFound flow, asserting indexer.Resync triggers a subsequent EnsureIndex call. Add an address-mismatch test for build that returns a different index_id and verifies i.pending is cleared or refreshed rather than left stale. Use existing service/indexer fakes and assert the behavior through observable calls and state.router/pkg/querygen/indexer.go (1)
92-96: 🩺 Stability & Availability | 🔵 Trivial | ⚡ Quick winAvoid storing
context.ContextinbaseCtx.
Reloadpasses the server-lifetimes.ctx, which is canceled only byStop. The reload-scoped cancellation scenario does not apply. Use a long-lived indexer context forResyncto avoid thecontainedctxfinding.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@router/pkg/querygen/indexer.go` around lines 92 - 96, Update the initialization around the indexer’s poll context so `baseCtx` stores a long-lived indexer-owned context rather than the incoming `ctx` from `Reload`; preserve `pollCtx` and its cancellation for reload polling, and ensure `Resync` uses the long-lived context without storing a `context.Context` field sourced from the server lifetime context.router/pkg/config/config.schema.json (1)
2764-2808: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winRequire
urlin the schema whenenabledis true.The router already fails at startup through
querygen.Config.Validate. Themcpblock enforces the equivalent rule for OAuth with anif/thenat Lines 3014-3040. Add the same conditional here so an editor and a CI schema check report the missingurlbefore the router runs.♻️ Proposed conditional
} - } + }, + "if": { + "properties": { + "enabled": { "const": true } + }, + "required": ["enabled"] + }, + "then": { + "required": ["url"], + "properties": { + "url": { "minLength": 1 } + } + } },🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@router/pkg/config/config.schema.json` around lines 2764 - 2808, Add an if/then conditional to the schema_discovery object so that when enabled is true, url is required. Preserve the existing property definitions and validation behavior, placing the conditional alongside the object’s other schema keywords.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@docs-website/router/mcp/schema-discovery/overview.mdx`:
- Line 20: Update the schema upload lifecycle description near the router
discovery overview to state that the router sends and indexes each schema when
its hash is new, and skips the discovery call only when the schema hash is
unchanged.
In `@docs-website/router/mcp/schema-discovery/tools.mdx`:
- Line 154: Update the schema discovery error-handling table entry in the
“Schema discovery is not configured correctly” row so a missing or empty token
is not classified as invalid configuration. Describe the failure as rejected
credentials, while preserving the existing guidance not to retry and to inform
the operator.
- Around line 162-164: Update the schema-exposure guidance to include
get_symbols alongside search_schema and generate_query, and state that all three
tools must be gated with the tools_call scope because they lack per-tool
`@requiresScopes` directives.
In `@router/go.mod`:
- Around line 56-60: Move the
buf.build/gen/go/bufbuild/protovalidate/protocolbuffers/go dependency from the
indirect require block into the direct require block in router/go.mod,
preserving its version and leaving it marked as a direct dependency because
yoko.pb.go imports it directly.
In `@router/pkg/config/config.schema.json`:
- Around line 2774-2778: Update the schema discovery service `url` property to
validate HTTP(S)-only URLs, replacing the generic `uri` format with the existing
`http-url` format used elsewhere in the schema (or an equivalent `^https?://`
pattern). Preserve the existing string type and description.
In `@router/pkg/querygen/indexer.go`:
- Around line 129-160: Update the address reconciliation in Indexer.build so
i.pending is synchronized with the service-returned index ID whenever want
changes from the locally computed address. Ensure subsequent fail handling
compares against the updated address, allowing pending to clear and future Sync
or Resync calls to retry normally.
In `@router/pkg/querygen/querygen_test.go`:
- Around line 47-59: Update fakeService methods EnsureIndex, SearchSchema, and
GenerateQuery to guard reads of mutable fields with f.mu, matching GetIndex’s
locking pattern. Lock access to ensureErr, ensureStatus, stale, searchErr,
searchHits, and generateFn, while preserving the existing method behavior and
responses.
---
Nitpick comments:
In `@router/pkg/config/config.schema.json`:
- Around line 2764-2808: Add an if/then conditional to the schema_discovery
object so that when enabled is true, url is required. Preserve the existing
property definitions and validation behavior, placing the conditional alongside
the object’s other schema keywords.
In `@router/pkg/querygen/config.go`:
- Around line 43-48: Update Config.Validate to parse c.URL with net/url and
reject values without a scheme, while retaining the existing empty-URL error.
Add the required fmt and net/url imports and return a clear validation error
that identifies the invalid URL and requires a scheme.
In `@router/pkg/querygen/indexer.go`:
- Around line 92-96: Update the initialization around the indexer’s poll context
so `baseCtx` stores a long-lived indexer-owned context rather than the incoming
`ctx` from `Reload`; preserve `pollCtx` and its cancellation for reload polling,
and ensure `Resync` uses the long-lived context without storing a
`context.Context` field sourced from the server lifetime context.
In `@router/pkg/querygen/querygen_test.go`:
- Around line 359-418: Extend the querygen test suite with coverage for
Service.handle’s connect.CodeNotFound flow, asserting indexer.Resync triggers a
subsequent EnsureIndex call. Add an address-mismatch test for build that returns
a different index_id and verifies i.pending is cleared or refreshed rather than
left stale. Use existing service/indexer fakes and assert the behavior through
observable calls and state.
In `@router/pkg/querygen/service.go`:
- Around line 234-260: Update the doc comment for UserMessage to state that the
router never adds the token to the returned message, rather than guaranteeing
that no secret can appear. Leave the existing error-message branches and
behavior unchanged.
🪄 Autofix
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro
Run ID: d5c8732e-4246-47b1-80e5-95cc858b9d1c
⛔ Files ignored due to path filters (4)
buf.lockis excluded by!**/*.lockrouter/gen/proto/yoko/v1/yoko.pb.gois excluded by!**/*.pb.go,!**/gen/**router/gen/proto/yoko/v1/yokov1connect/yoko.connect.gois excluded by!**/gen/**router/go.sumis excluded by!**/*.sum
📒 Files selected for processing (28)
Makefilebuf.router.go.gen.yamlbuf.yamldocs-website/docs.jsondocs-website/router/mcp.mdxdocs-website/router/mcp/configuration.mdxdocs-website/router/mcp/schema-discovery/configuration.mdxdocs-website/router/mcp/schema-discovery/guides.mdxdocs-website/router/mcp/schema-discovery/overview.mdxdocs-website/router/mcp/schema-discovery/quickstart.mdxdocs-website/router/mcp/schema-discovery/tools.mdxdocs-website/router/mcp/tools.mdxproto/yoko/v1/yoko.protorouter/core/router.gorouter/go.modrouter/mcp.schema-discovery.config.yamlrouter/pkg/config/config.gorouter/pkg/config/config.schema.jsonrouter/pkg/config/testdata/config_defaults.jsonrouter/pkg/config/testdata/config_full.jsonrouter/pkg/mcpserver/schema_discovery_tools.gorouter/pkg/mcpserver/schema_discovery_tools_test.gorouter/pkg/mcpserver/server.gorouter/pkg/querygen/client.gorouter/pkg/querygen/config.gorouter/pkg/querygen/indexer.gorouter/pkg/querygen/querygen_test.gorouter/pkg/querygen/service.go
|
|
||
| ## What schema discovery does | ||
|
|
||
| The router sends the client schema to the discovery service one time. The service builds an index. An agent then searches that index instead of reading the schema. |
There was a problem hiding this comment.
🔒 Security & Privacy | 🟡 Minor | ⚡ Quick win
Correct the schema upload lifecycle.
Line 20 says that the router sends the schema one time. The router also indexes each changed schema. State that it sends the schema for each new schema hash and skips the call only when the hash is unchanged.
Proposed fix
-The router sends the client schema to the discovery service one time. The service builds an index. An agent then searches that index instead of reading the schema.
+The router sends the client schema to the discovery service for each new schema hash. The service builds an index. An agent then searches that index instead of reading the schema.📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| The router sends the client schema to the discovery service one time. The service builds an index. An agent then searches that index instead of reading the schema. | |
| The router sends the client schema to the discovery service for each new schema hash. The service builds an index. An agent then searches that index instead of reading the schema. |
🧰 Tools
🪛 LanguageTool
[style] ~20-~20: Consider using “once” instead of ‘one time’.
Context: ... client schema to the discovery service one time. The service builds an index. An agent ...
(ONCE_TWICE_PREMIUM)
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@docs-website/router/mcp/schema-discovery/overview.mdx` at line 20, Update the
schema upload lifecycle description near the router discovery overview to state
that the router sends and indexes each schema when its hash is new, and skips
the discovery call only when the schema hash is unchanged.
| google.golang.org/protobuf v1.36.12 | ||
| ) | ||
|
|
||
| require ( | ||
| buf.build/gen/go/bufbuild/protovalidate/protocolbuffers/go v1.36.12-20260709200747-435963d16310.1 |
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
# Description: Check where the protovalidate module is required and whether router code imports it directly.
rg -n -B2 -A2 'protovalidate' router/go.mod
rg -n 'buf/validate' router/gen/proto/yoko/v1 | head -20Repository: wundergraph/cosmo
Length of output: 642
🏁 Script executed:
#!/bin/bash
set -e
printf '%s\n' '--- router/go.mod ---'
sed -n '1,90p' router/go.mod
printf '%s\n' '--- protovalidate imports under router ---'
rg -n -g '*.go' 'buf\.build/gen/go/bufbuild/protovalidate|buf/validate' router
printf '%s\n' '--- module and generated package context ---'
sed -n '1,35p' router/gen/proto/yoko/v1/yoko.pb.go
rg -n '^(module|go) ' router/go.mod
printf '%s\n' '--- all require directives containing protovalidate ---'
rg -n -B3 -A3 'protovalidate' --glob 'go.mod' .Repository: wundergraph/cosmo
Length of output: 5724
🏁 Script executed:
#!/bin/bash
set -e
python3 - <<'PY'
from pathlib import Path
import re
mod = Path("router/go.mod").read_text().splitlines()
module = next(line.split()[1] for line in mod if line.startswith("module "))
target = "buf.build/gen/go/bufbuild/protovalidate/protocolbuffers/go"
in_require = False
entries = []
for number, line in enumerate(mod, 1):
stripped = line.strip()
if stripped == "require (":
in_require = True
continue
if in_require and stripped == ")":
in_require = False
continue
if in_require:
match = re.match(r"\s*(\S+)\s+(\S+)(?:\s+//\s*(.*))?$", line)
if match:
path, version, comment = match.groups()
entries.append((number, path, version, comment or ""))
target_entry = [entry for entry in entries if entry[1] == target]
print(f"module: {module}")
print(f"target require entry: {target_entry}")
imports = []
for path in Path("router").rglob("*.go"):
for number, line in enumerate(path.read_text(errors="replace").splitlines(), 1):
if target + "/" in line:
imports.append((str(path), number, line.strip()))
print("target imports:")
for item in imports:
print(item)
print("classification:")
if target_entry and imports and not target_entry[0][3].strip().startswith("indirect"):
print("The target is imported by router Go source and has no // indirect marker.")
else:
print("The target is not classified as a direct source dependency by these checks.")
PY
if command -v go >/dev/null 2>&1; then
go version
go help mod tidy | sed -n '1,80p'
else
echo "go tool unavailable"
fiRepository: wundergraph/cosmo
Length of output: 2194
🌐 Web query:
Go modules reference require blocks // indirect marker block placement go mod tidy official
💡 Result:
In Go modules, the go.mod file organizes dependencies using require directives. When a dependency is not directly imported by any package in the main module, the go tool marks it with a // indirect comment [1][2]. Regarding the placement and management of these markers: 1. Block Separation: Because there can be a significant number of indirect dependencies, the go command records them in a separate require block within the go.mod file [3][4]. This separation helps distinguish direct dependencies from indirect ones [5]. 2. Role of go mod tidy: The go mod tidy command is the official tool responsible for maintaining the accuracy of these requirements [6]. It automatically scans your code's imports, adds necessary modules, removes unused ones, and correctly categorizes dependencies by applying or removing the // indirect comment [6]. 3. Automated Formatting: The go.mod file is machine-managed [5]. When you run go mod tidy, the go command automatically handles the formatting, including the placement of dependencies into the appropriate (direct or indirect) require blocks and the addition of the required // indirect markers [3][6]. You generally do not need to manually manage these blocks or markers; running go mod tidy ensures your go.mod file correctly reflects your project's dependency graph [3][6].
Citations:
- 1: https://go.dev/ref/mod
- 2: https://tip.golang.org/ref/mod
- 3: https://go.dev/doc/modules/gomod-ref
- 4: https://tip.golang.org/doc/modules/gomod-ref
- 5: cmd/go: separate direct and indirect dependencies in go.mod golang/go#27887
- 6: https://go.dev/doc/modules/managing-dependencies
Move the protovalidate module into the direct require block. yoko.pb.go imports buf/validate directly, but router/go.mod:60 lists the module in the second, indirect dependency block. go mod tidy -diff can rewrite this placement.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@router/go.mod` around lines 56 - 60, Move the
buf.build/gen/go/bufbuild/protovalidate/protocolbuffers/go dependency from the
indirect require block into the direct require block in router/go.mod,
preserving its version and leaving it marked as a direct dependency because
yoko.pb.go imports it directly.
| func (i *Indexer) build(ctx context.Context, sdl, want string) { | ||
| log := i.logger.With(zap.String("index_id", want)) | ||
|
|
||
| res, err := i.client.EnsureIndex(ctx, connect.NewRequest(&yokov1.EnsureIndexRequest{Sdl: sdl})) | ||
| if err != nil { | ||
| if ctx.Err() != nil { | ||
| return // a newer Sync replaced this build | ||
| } | ||
| i.fail(want, fmt.Errorf("failed to send the schema: %w", err)) | ||
| log.Error("failed to send the schema to the discovery service", zap.Error(err)) | ||
| return | ||
| } | ||
|
|
||
| idx := res.Msg.GetIndex() | ||
| if got := idx.GetIndexId(); got != want { | ||
| // The service disagrees about the address. Trust the service, because | ||
| // it holds the index, but record the difference. | ||
| log.Warn("the discovery service returned a different address", | ||
| zap.String("returned_index_id", got)) | ||
| want = got | ||
| } | ||
|
|
||
| if idx.GetStatus() == yokov1.IndexStatus_INDEX_STATUS_READY { | ||
| i.adopt(want, idx) | ||
| log.Info("schema index is ready", | ||
| zap.Int64("symbol_count", idx.GetSymbolCount())) | ||
| return | ||
| } | ||
|
|
||
| log.Info("schema index is building") | ||
| i.poll(ctx, want, log) | ||
| } |
There was a problem hiding this comment.
🩺 Stability & Availability | 🟠 Major | ⚡ Quick win
Fix the stuck pending state when the service returns a different address.
Line 148 reassigns want to the address the service returned. i.pending still holds the locally computed address. If the build then fails, fail compares i.pending == address at line 228, the comparison is false, and i.pending keeps the stale value forever.
After that:
Syncwith the same SDL returns early at line 83, becausei.pending == want.Resyncreturns early at line 109, becausei.pending != "".
The indexer never retries and CurrentAddress returns ErrIndexNotReady for the lifetime of the process.
🐛 Proposed fix
idx := res.Msg.GetIndex()
if got := idx.GetIndexId(); got != want {
// The service disagrees about the address. Trust the service, because
// it holds the index, but record the difference.
log.Warn("the discovery service returned a different address",
zap.String("returned_index_id", got))
+ i.mu.Lock()
+ if i.pending == want {
+ i.pending = got
+ }
+ i.mu.Unlock()
want = got
}📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| func (i *Indexer) build(ctx context.Context, sdl, want string) { | |
| log := i.logger.With(zap.String("index_id", want)) | |
| res, err := i.client.EnsureIndex(ctx, connect.NewRequest(&yokov1.EnsureIndexRequest{Sdl: sdl})) | |
| if err != nil { | |
| if ctx.Err() != nil { | |
| return // a newer Sync replaced this build | |
| } | |
| i.fail(want, fmt.Errorf("failed to send the schema: %w", err)) | |
| log.Error("failed to send the schema to the discovery service", zap.Error(err)) | |
| return | |
| } | |
| idx := res.Msg.GetIndex() | |
| if got := idx.GetIndexId(); got != want { | |
| // The service disagrees about the address. Trust the service, because | |
| // it holds the index, but record the difference. | |
| log.Warn("the discovery service returned a different address", | |
| zap.String("returned_index_id", got)) | |
| want = got | |
| } | |
| if idx.GetStatus() == yokov1.IndexStatus_INDEX_STATUS_READY { | |
| i.adopt(want, idx) | |
| log.Info("schema index is ready", | |
| zap.Int64("symbol_count", idx.GetSymbolCount())) | |
| return | |
| } | |
| log.Info("schema index is building") | |
| i.poll(ctx, want, log) | |
| } | |
| func (i *Indexer) build(ctx context.Context, sdl, want string) { | |
| log := i.logger.With(zap.String("index_id", want)) | |
| res, err := i.client.EnsureIndex(ctx, connect.NewRequest(&yokov1.EnsureIndexRequest{Sdl: sdl})) | |
| if err != nil { | |
| if ctx.Err() != nil { | |
| return // a newer Sync replaced this build | |
| } | |
| i.fail(want, fmt.Errorf("failed to send the schema: %w", err)) | |
| log.Error("failed to send the schema to the discovery service", zap.Error(err)) | |
| return | |
| } | |
| idx := res.Msg.GetIndex() | |
| if got := idx.GetIndexId(); got != want { | |
| // The service disagrees about the address. Trust the service, because | |
| // it holds the index, but record the difference. | |
| log.Warn("the discovery service returned a different address", | |
| zap.String("returned_index_id", got)) | |
| i.mu.Lock() | |
| if i.pending == want { | |
| i.pending = got | |
| } | |
| i.mu.Unlock() | |
| want = got | |
| } | |
| if idx.GetStatus() == yokov1.IndexStatus_INDEX_STATUS_READY { | |
| i.adopt(want, idx) | |
| log.Info("schema index is ready", | |
| zap.Int64("symbol_count", idx.GetSymbolCount())) | |
| return | |
| } | |
| log.Info("schema index is building") | |
| i.poll(ctx, want, log) | |
| } |
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@router/pkg/querygen/indexer.go` around lines 129 - 160, Update the address
reconciliation in Indexer.build so i.pending is synchronized with the
service-returned index ID whenever want changes from the locally computed
address. Ensure subsequent fail handling compares against the updated address,
allowing pending to clear and future Sync or Resync calls to retry normally.
| func (f *fakeService) EnsureIndex(_ context.Context, req *connect.Request[yokov1.EnsureIndexRequest]) (*connect.Response[yokov1.EnsureIndexResponse], error) { | ||
| f.ensureCalls.Add(1) | ||
| if f.ensureErr != nil { | ||
| return nil, f.ensureErr | ||
| } | ||
| return connect.NewResponse(&yokov1.EnsureIndexResponse{ | ||
| Index: &yokov1.Index{ | ||
| IndexId: Address(req.Msg.GetSdl()), | ||
| Status: f.ensureStatus, | ||
| Stale: f.stale, | ||
| }, | ||
| }), nil | ||
| } |
There was a problem hiding this comment.
🩺 Stability & Availability | 🟠 Major | ⚡ Quick win
Guard the fakeService fields with f.mu in EnsureIndex.
GetIndex reads the mutable fields under f.mu, but EnsureIndex reads f.ensureErr, f.ensureStatus, and f.stale without the lock. TestSync_SchemaChangeKeepsServingTheOldIndex writes f.ensureStatus at line 246 under f.mu while the build goroutine can serve a request. That is a data race, and go test -race fails.
SearchSchema reads f.searchErr and f.searchHits without the lock as well. GenerateQuery reads f.generateFn without the lock.
🐛 Proposed fix for `EnsureIndex`
func (f *fakeService) EnsureIndex(_ context.Context, req *connect.Request[yokov1.EnsureIndexRequest]) (*connect.Response[yokov1.EnsureIndexResponse], error) {
f.ensureCalls.Add(1)
+
+ f.mu.Lock()
+ defer f.mu.Unlock()
+
if f.ensureErr != nil {
return nil, f.ensureErr
}
return connect.NewResponse(&yokov1.EnsureIndexResponse{
Index: &yokov1.Index{
IndexId: Address(req.Msg.GetSdl()),
Status: f.ensureStatus,
Stale: f.stale,
},
}), nil
}Apply the same lock to SearchSchema and GenerateQuery.
📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| func (f *fakeService) EnsureIndex(_ context.Context, req *connect.Request[yokov1.EnsureIndexRequest]) (*connect.Response[yokov1.EnsureIndexResponse], error) { | |
| f.ensureCalls.Add(1) | |
| if f.ensureErr != nil { | |
| return nil, f.ensureErr | |
| } | |
| return connect.NewResponse(&yokov1.EnsureIndexResponse{ | |
| Index: &yokov1.Index{ | |
| IndexId: Address(req.Msg.GetSdl()), | |
| Status: f.ensureStatus, | |
| Stale: f.stale, | |
| }, | |
| }), nil | |
| } | |
| func (f *fakeService) EnsureIndex(_ context.Context, req *connect.Request[yokov1.EnsureIndexRequest]) (*connect.Response[yokov1.EnsureIndexResponse], error) { | |
| f.ensureCalls.Add(1) | |
| f.mu.Lock() | |
| defer f.mu.Unlock() | |
| if f.ensureErr != nil { | |
| return nil, f.ensureErr | |
| } | |
| return connect.NewResponse(&yokov1.EnsureIndexResponse{ | |
| Index: &yokov1.Index{ | |
| IndexId: Address(req.Msg.GetSdl()), | |
| Status: f.ensureStatus, | |
| Stale: f.stale, | |
| }, | |
| }), nil | |
| } |
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@router/pkg/querygen/querygen_test.go` around lines 47 - 59, Update
fakeService methods EnsureIndex, SearchSchema, and GenerateQuery to guard reads
of mutable fields with f.mu, matching GetIndex’s locking pattern. Lock access to
ensureErr, ensureStatus, stale, searchErr, searchHits, and generateFn, while
preserving the existing method behavior and responses.
This comment has been minimized.
This comment has been minimized.
connect-go - uncommitted changes detectedSeems like you forgot to commit some code. Possible causes:
Dirty files
|
230681a to
a9b0b5d
Compare
Codecov Report❌ Patch coverage is ❌ Your patch check has failed because the patch coverage (60.25%) is below the target coverage (90.00%). You can increase the patch coverage or adjust the target coverage. Additional details and impacted files@@ Coverage Diff @@
## main #3156 +/- ##
==========================================
- Coverage 62.40% 62.39% -0.01%
==========================================
Files 263 270 +7
Lines 31048 32195 +1147
==========================================
+ Hits 19375 20089 +714
- Misses 10163 10568 +405
- Partials 1510 1538 +28
🚀 New features to boost your workflow:
|
806b21a to
a60a313
Compare
Caution
DO NOT MERGE. This is a demo branch. It exists to show the capability and to start a discussion.
It is not a merge candidate. Nobody agreed that this ships. Nobody agreed on this shape. There is no hosted discovery service for users today. The vendored proto also publishes a service API that nobody cleared for release.
Linear: ROUTER-628
The demo
An agent works against a federated graph that is too large to fit in its context.
Today the MCP server has two ways to show the graph.
get_schemareturns the whole SDL and fills the context. Persisted operations become tools, but somebody must author each one first. Neither helps an agent that knows what it wants and does not know the schema.This branch adds three MCP tools:
search_schemaget_symbolsgenerate_queryThe Router indexes its client schema at startup. No operator action.
Run it
Point any MCP client at
http://localhost:5035/mcp. The Router logsschema index is readywhen the tools can serve.The demo needs a schema discovery service on port 3400.
What it looks like
Ask for data. Get a valid operation.
Prompt: list all employees with their id, first name, last name and current mood
The service checks the operation against the schema before it answers. Every field in it exists.
Values become variables, not literals.
Prompt: find one employee by id and show their hobbies
{ "findemployeesby_criteria": { "properties": { "department": { "enum": ["ENGINEERING", "MARKETING", "OPERATIONS", null] }, "id": { "type": ["integer", "null"] }, "title": { "type": ["string", "null"] } } } }The variables schema carries the allowed values from the graph. A model cannot invent a department. Generate the operation one time, then call it many times.
Ask for something that does not exist.
Prompt: list the invoices for a customer with their billing address
{ "unsatisfied": [ "The indexed schema exposes products, employees, locations, and work reviews, but no invoice entity, billing address, or payment status." ] }The result has no queries and one reason. This answers the question "does this capability already exist?". A developer asks before they build. The answer tells them to build or to reuse.
Run the operation. Set
enable_arbitrary_operationsto true. The agent then runs the document throughexecute_graphql. The subgraphs return the data.Measured
Against the
router-testsexecution config:Configuration
Off by default. One new block:
router/mcp.schema-discovery.config.yamlreproduces the demo.How it works
The index is content addressed. Its address is the SHA-256 hash of the schema bytes. The Router therefore computes the address locally, and holds no state in the service.
Syncruns on every config reload. It compares the hash and returns. Every service call happens in a goroutine, so a slow service cannot delay a reload. An unchanged schema makes no network call.The Router adopts a new address only when that address is ready. The previous index keeps serving during a rebuild.
The Router indexes the client schema, not the supergraph. No federation internals leave the Router.
The service never runs an operation. Authentication, rate limits, and audit logs stay in the Router.
Implementation notes
The MCP HTTP server used a write timeout of 30 seconds. Query generation takes 10 to 30 seconds. The server thus closed the response before the handler returned.
writeTimeout()now returns the larger of 30 seconds andrequest_timeoutplus 10 seconds. Behaviour is unchanged when schema discovery is off.The upstream service repository is private, so the Router cannot import it.
proto/yoko/v1/yoko.protois a byte identical copy, and stubs generate intorouter/gen/proto/yoko/v1through the existinggenerate-gotarget.Tests cover the index lifecycle, the tool registration, and the config. They run against a Connect test server, so no test needs a live service.
Documentation sits under
router/mcp/schema-discovery/, split by Diataxis.Not done
search_schemaandgenerate_queryshow the schema shape to any caller that reaches the MCP server. This matches whatget_schemaalready gives.Summary by CodeRabbit
New Features
Documentation
Bug Fixes