Improve Aurora MCP tool discovery & steering#466
Open
isiddharthsingh wants to merge 10 commits into
Open
Conversation
…1, drop phantom github route, add CI/CD+Sentry+Grafana entries
…urora toward direct lookups
…mpts toward direct reads
…overage, and per-client setup
… zeroes out real token matches
Contributor
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Path: .coderabbit.yaml Review profile: ASSERTIVE Plan: Pro Run ID: 📒 Files selected for processing (2)
WalkthroughPromotes several read-only tools to Tier-1, replaces substring dispatch search with tokenized ranking plus fallback, removes promoted entries from the Tier-3 allowlist, updates prompts/docs to prefer direct tools over chat_with_aurora for factual lookups, and adds tests covering search and registration behavior. ChangesMCP Discovery and Steering Refactor
Estimated code review effort🎯 4 (Complex) | ⏱️ ~45 minutes Possibly related PRs
Suggested reviewers
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches📝 Generate docstrings
🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Actionable comments posted: 3
🤖 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 `@server/aurora_mcp/dispatch.py`:
- Around line 129-136: The hint text in server/aurora_mcp/dispatch.py currently
promises "traces" support even though no trace tool or dispatch exists; update
the "hint" string(s) (the dict entries with the "hint" key in dispatch.py) to
remove any mention of "traces" (and related phrasing about searching for trace
tools) so the MCP does not advertise an unsupported capability; repeat the same
removal for the other similar hint block in the file (the second "hint" entry
referenced in the review) to keep messaging consistent with registry.py's actual
tools.
In `@server/aurora_mcp/registry.py`:
- Around line 975-981: _tokenize_query currently splits only on whitespace so
punctuation-attached tokens like "deployments?", "mttr/dora", or "rca,tools"
remain joined; update _tokenize_query to use a regex-based tokenization (e.g.
re.findall for alphanumeric runs) instead of q.split(), filter out tokens
shorter than 3, and add the required import (re) so tokens delimited by
punctuation are separated correctly.
In `@server/tests/mcp/test_mcp_registry.py`:
- Line 143: The monkeypatch of registry._check_skill_connected currently uses a
lambda with an unused parameter named u which triggers Ruff ARG005; update the
lambda used in the monkeypatch (the registry._check_skill_connected replacement)
to rename the unused parameter to a leading-underscore placeholder (e.g., _u or
simply _) so the function signature remains the same but the linter no longer
flags the unused argument.
🪄 Autofix (Beta)
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: Path: .coderabbit.yaml
Review profile: ASSERTIVE
Plan: Pro
Run ID: 46f4f880-1218-436b-85d8-ac3b239c39e1
📒 Files selected for processing (8)
docs/design/mcp-agent-steering.mdserver/aurora_mcp/dispatch.pyserver/aurora_mcp/prompts.pyserver/aurora_mcp/registry.pyserver/aurora_mcp/tools_always_on.pyserver/mcp_server.pyserver/tests/mcp/test_mcp_registry.pywebsite/docs/integrations/mcp.md
…tuation, fix test lint
Contributor
There was a problem hiding this comment.
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
server/aurora_mcp/dispatch.py (1)
115-123:⚠️ Potential issue | 🟠 Major | ⚡ Quick winFetch the full ranked match set before reordering by
callable_now.This truncates before visibility sorting. For broad queries, disconnected entries can fill the first
limit * 2slots, so connected tools that rank slightly lower never get surfaced even though this function is trying to put callable tools first. Pull the full ranked set, then split visible/non-visible, and only applylimitat the end.Proposed fix
from .registry import ( + DISPATCH_ALLOWLIST, _get_cached_connector_status, dispatch_entry_visible, find_dispatch_entry, parse_and_cache_connector_status, search_dispatch_entries, @@ all_matches = search_dispatch_entries( query=query, category=category, connector=connector, - user_id=None, limit=limit * 2, + user_id=None, limit=len(DISPATCH_ALLOWLIST), )🤖 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 `@server/aurora_mcp/dispatch.py` around lines 115 - 123, The code currently truncates matches before reordering by visibility (all_matches = search_dispatch_entries(..., limit=limit * 2)), which can exclude callable entries; change the flow in dispatch.py so you fetch the full ranked set (call search_dispatch_entries without the premature limit or with a sufficiently large/unbounded limit), then build annotated = [(e, dispatch_entry_visible(e, user_id)) for e in all_matches], then form ordered by visible-first ([(e, True)...] + [(e, False)...]) and only apply the final [:limit] slice to ordered; adjust variable usage around all_matches, annotated, and ordered accordingly.
🤖 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.
Outside diff comments:
In `@server/aurora_mcp/dispatch.py`:
- Around line 115-123: The code currently truncates matches before reordering by
visibility (all_matches = search_dispatch_entries(..., limit=limit * 2)), which
can exclude callable entries; change the flow in dispatch.py so you fetch the
full ranked set (call search_dispatch_entries without the premature limit or
with a sufficiently large/unbounded limit), then build annotated = [(e,
dispatch_entry_visible(e, user_id)) for e in all_matches], then form ordered by
visible-first ([(e, True)...] + [(e, False)...]) and only apply the final
[:limit] slice to ordered; adjust variable usage around all_matches, annotated,
and ordered accordingly.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: ASSERTIVE
Plan: Pro
Run ID: a5efd3f1-6192-4196-a967-db40205d35a8
📒 Files selected for processing (3)
server/aurora_mcp/dispatch.pyserver/aurora_mcp/registry.pyserver/tests/mcp/test_mcp_registry.py
…ools aren't dropped
|
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Why
External agents (Cursor, Claude Desktop/Code, Codex, Windsurf) connecting to Aurora's self-hosted MCP server consistently over-used
chat_with_aurora/ask_incidentfor questions a cheap direct read could answer, and failed to discover obvious read-only tools (infrastructure context, service graph, RCA findings). The steering text was chat-maximalist in three layers, the high-value reads were buried insearch_tools-only Tier 3, and the per-sub-agent RCA findings weren't on the MCP surface at all. Result: wrong tool choices and ~40s latency where a ~1s direct read would do.The tiered architecture is correct (keeps the upfront surface under client tool-count caps). This keeps it and fixes the problem with (a) re-steering the prompt text toward direct reads, (b) promoting a curated set of reads to the always-visible surface, and (c) fixing the
search_toolsmatcher so the long tail is findable.What changed
1. Re-steered the steering text (chat-first → direct-read-first)
mcp_server.pyinstructions— routing guide naming direct-read categories; critical sentence front-loaded; phantomgithub_rcascrubbed.chat_with_auroradocstring — dropped "Default tool for any question", added a "NOT a structured-data lookup" scope note, trimmed under the 2KB cap (was being truncated).ask_incidentpoints at the direct readers first.query_logs/query_metricsdescriptions softened;search_toolsdocstring + hint made loud about the long tail.prompts.pyreordered to direct reads first;blast_radius_analysisrepointed to the newservice_impacttool.2. Promoted 6 curated reads to first-class Tier-1 tools
get_infrastructure_context,list_services,service_impact,incident_findings,incident_finding_detail,incident_list_alerts. The 3 graph entries were removed fromDISPATCH_ALLOWLIST(no double-exposure); findings are new.service_impactURL-encodes the service name itself.3. Fixed the
search_toolsmatcher (two bugs)"rca tools steps"and forced a chat fallback); ranks by match count.category="metrics"for an LLM-usage question whose tools live undercategory="usage"previously got zero results and gave up).4. Expanded connector coverage — added skill-gated read entries for Jenkins/CloudBees/Spinnaker (CI/CD), Sentry, and Grafana. Removed the phantom
github_list_reposentry (/github/repos404s).Docs — updated
website/docs/integrations/mcp.md(tier tables, routing note, per-client setup caveats) and addeddocs/design/mcp-agent-steering.md.Testing
docker exec aurora-server python -m pytest tests/mcp/ -q).search_tools/call_tooldiscovery, connector-gating, and the allowlist boundary all behave correctly; the wrong-category fallback resolves the reported empty-result case.instructionsand descriptions-only.Note for reviewers
The CLAUDE.md New Connector Checklist and the architectural RBAC test do not apply here — these are MCP tools wrapping already-RBAC-decorated REST routes, not new connectors/routes.
Summary by CodeRabbit
New Features
Documentation
Tests