fix: migrate MCP server to SDK v2 - #636
Conversation
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
There was a problem hiding this comment.
Your trial has ended. Reactivate Greptile to resume code reviews.
Bugbot couldn't run - usage limit reachedBugbot is counted against Cursor usage for this user or team, and this run hit a usage or spend limit. A user or team admin can review and increase usage limits in the Cursor dashboard. (requestId: serverGenReqId_b8324236-9b49-4591-9cbc-f78d28c47447) |
📝 WalkthroughWalkthroughThe MCP dependency moves to v2. The server adopts low-level v2 callbacks, typed requests, schema validation, snake_case fields, structured results, and metadata. The mock transport and MCP tests update to the new client, schema, annotation, and error APIs. ChangesMCP v2 server migration
Estimated code review effort: 4 (Complex) | ~45 minutes Possibly related PRs
Poem
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches 💡 1📝 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 |
![macroscopeapp[bot]](https://avatars.githubusercontent.com/in/900172?s=60&v=4){kind=small}
Bugbot couldn't run - usage limit reachedBugbot is counted against Cursor usage for this user or team, and this run hit a usage or spend limit. A user or team admin can review and increase usage limits in the Cursor dashboard. (requestId: serverGenReqId_07c6e0c7-6483-4730-9069-ca4b00750de0) |
There was a problem hiding this comment.
Your trial has ended. Reactivate Greptile to resume code reviews.
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 4
🤖 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 `@pyproject.toml`:
- Around line 19-20: Update the jsonschema dependency specification in
pyproject.toml to add an upper bound excluding the next major version, matching
the existing mcp dependency constraint while preserving the current minimum
version.
In `@src/brainlayer/mcp/__init__.py`:
- Around line 153-154: Update the exception handler around _error_result to log
the caught exception with the module’s existing logger before returning the
client-facing error result, preserving the current str(exc) response behavior.
- Around line 115-129: Update _normalize_call_tool_result to reject bare strings
and tuples whose length is not exactly two, rather than routing them through the
iterable fallback. Restrict the iterable branch to supported iterable content
forms, preserving valid two-item tuple, dict, and other accepted collection
handling, and raise TypeError for unsupported values.
- Around line 132-144: Cache the static tool definitions used by
_handle_call_tool_request instead of awaiting list_tools() for every call. Apply
a single-entry cache to _full_tool_definitions() or the corresponding
definition-site function, while preserving ToolPalette.expose()’s dynamic
session filtering and existing schema validation behavior.
🪄 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: Organization UI
Review profile: ASSERTIVE
Plan: Pro Plus
Run ID: aa9edbb1-d682-4c50-a524-228b741a18ab
📒 Files selected for processing (33)
pyproject.tomlscripts/smoke/repogolem_mcp_transport_smoke.pysrc/brainlayer/brainbar_hybrid_helper.pysrc/brainlayer/mcp/__init__.pysrc/brainlayer/mcp/_shared.pysrc/brainlayer/mcp/palette.pytests/mock_mcp/base.pytests/test_3tool_aliases.pytests/test_audit_search_quality.pytests/test_brainbar_hybrid_helper.pytests/test_brainstore.pytests/test_chunk_lifecycle.pytests/test_enrichment_controller.pytests/test_entity_type_sync.pytests/test_issue_type.pytests/test_mcp_digest_modes.pytests/test_mcp_input_schema_limits.pytests/test_mcp_labeled_field_output.pytests/test_mcp_palette.pytests/test_mcp_timeout.pytests/test_mcp_warm_route.pytests/test_phase3_digest.pytests/test_phase5.pytests/test_phase6_sentiment.pytests/test_precompact_chunk_origin.pytests/test_search_fanout.pytests/test_search_filter_params.pytests/test_search_validation.pytests/test_smart_search_entity_dedup.pytests/test_store_handler.pytests/test_think_recall_integration.pytests/test_tool_annotations.pytests/test_write_queue.py
📜 Review details
⏰ Context from checks skipped due to timeout. (3)
- GitHub Check: test (3.12)
- GitHub Check: test (3.13)
- GitHub Check: test (3.11)
🧰 Additional context used
📓 Path-based instructions (5)
**/*.py
📄 CodeRabbit inference engine (AGENTS.md)
**/*.py: Flag risky DB or concurrency changes explicitly and do not hand-wave lock behavior
Enforce one-write-at-a-time concurrency constraint; reads are safe but brain_digest is write-heavy and must not run in parallel with other MCP work
Run pytest before claiming behavior changed safely; current test suite has 929 tests
Files:
src/brainlayer/mcp/_shared.pytests/test_think_recall_integration.pyscripts/smoke/repogolem_mcp_transport_smoke.pytests/test_phase3_digest.pytests/test_brainbar_hybrid_helper.pytests/test_phase6_sentiment.pytests/test_mcp_warm_route.pytests/test_audit_search_quality.pytests/test_tool_annotations.pytests/test_phase5.pytests/test_store_handler.pytests/test_issue_type.pysrc/brainlayer/brainbar_hybrid_helper.pytests/test_mcp_labeled_field_output.pytests/test_search_validation.pytests/test_search_fanout.pytests/test_chunk_lifecycle.pytests/test_write_queue.pytests/test_precompact_chunk_origin.pysrc/brainlayer/mcp/palette.pytests/test_smart_search_entity_dedup.pytests/test_entity_type_sync.pytests/test_search_filter_params.pytests/test_enrichment_controller.pytests/test_mcp_palette.pytests/test_brainstore.pytests/test_mcp_digest_modes.pytests/mock_mcp/base.pytests/test_mcp_input_schema_limits.pytests/test_mcp_timeout.pytests/test_3tool_aliases.pysrc/brainlayer/mcp/__init__.py
src/**/*.py
📄 CodeRabbit inference engine (CLAUDE.md)
src/**/*.py: Usepaths.py:get_db_path()for all database path resolution; do not hard-code the BrainLayer database path.
Each worker must use its own database connection and retry onSQLITE_BUSY.
For bulk database operations, stop enrichment workers, checkpoint WAL before and after, drop and recreate FTS triggers around large deletes, batch deletes in 5–10K chunks, and checkpoint every three batches.
Use the canonical BrainLayer naming:BrainLayer (זיכרון)means “memory.”
Files:
src/brainlayer/mcp/_shared.pysrc/brainlayer/brainbar_hybrid_helper.pysrc/brainlayer/mcp/palette.pysrc/brainlayer/mcp/__init__.py
src/brainlayer/mcp/**/*.py
📄 CodeRabbit inference engine (CLAUDE.md)
src/brainlayer/mcp/**/*.py: Use the fixed-size readonly WALVectorStorepool; enforceBRAINLAYER_READ_POOL_SIZE,BRAINLAYER_READ_BUSY_TIMEOUT_MS, and the startup memory limit.
Add new MCP tools under themcp/package and expose them through thebrainlayer-mcpentrypoint.
Files:
src/brainlayer/mcp/_shared.pysrc/brainlayer/mcp/palette.pysrc/brainlayer/mcp/__init__.py
src/brainlayer/**/*.py
📄 CodeRabbit inference engine (CLAUDE.md)
src/brainlayer/**/*.py: Preserveai_code,stack_trace, anduser_messageverbatim; skip noise, summarize build logs, and retain only structure for directory listings.
Use AST-aware tree-sitter chunking, never split stack traces, and mask large tool output.
Keep Groq as the primary enrichment backend, Gemini as fallback, and Ollama as the offline last resort; honorBRAINLAYER_ENRICH_BACKENDandBRAINLAYER_ENRICH_RATE.
Default searches must exclude lifecycle-managed chunks;include_archived=Truemay expose history.brain_supersedemust apply its personal-data safety gate, andbrain_archivemust soft-delete with a timestamp.
Files:
src/brainlayer/mcp/_shared.pysrc/brainlayer/brainbar_hybrid_helper.pysrc/brainlayer/mcp/palette.pysrc/brainlayer/mcp/__init__.py
tests/**/*.py
📄 CodeRabbit inference engine (CLAUDE.md)
Tests must not refresh the production backup heartbeat log; use
BRAINLAYER_BACKUP_LOG_PATHand set provenance topytest.
Files:
tests/test_think_recall_integration.pytests/test_phase3_digest.pytests/test_brainbar_hybrid_helper.pytests/test_phase6_sentiment.pytests/test_mcp_warm_route.pytests/test_audit_search_quality.pytests/test_tool_annotations.pytests/test_phase5.pytests/test_store_handler.pytests/test_issue_type.pytests/test_mcp_labeled_field_output.pytests/test_search_validation.pytests/test_search_fanout.pytests/test_chunk_lifecycle.pytests/test_write_queue.pytests/test_precompact_chunk_origin.pytests/test_smart_search_entity_dedup.pytests/test_entity_type_sync.pytests/test_search_filter_params.pytests/test_enrichment_controller.pytests/test_mcp_palette.pytests/test_brainstore.pytests/test_mcp_digest_modes.pytests/mock_mcp/base.pytests/test_mcp_input_schema_limits.pytests/test_mcp_timeout.pytests/test_3tool_aliases.py
🪛 ast-grep (0.45.0)
tests/mock_mcp/base.py
[info] 64-64: use jsonify instead of json.dumps for JSON output
Context: json.dumps({"mock": True, "tool": params.name})
Note: [CWE-116] Improper Encoding or Escaping of Output.
(use-jsonify)
[info] 67-67: use jsonify instead of json.dumps for JSON output
Context: json.dumps(result)
Note: [CWE-116] Improper Encoding or Escaping of Output.
(use-jsonify)
src/brainlayer/mcp/__init__.py
[info] 121-121: use jsonify instead of json.dumps for JSON output
Context: json.dumps(result, indent=2)
Note: [CWE-116] Improper Encoding or Escaping of Output.
(use-jsonify)
🔇 Additional comments (38)
tests/mock_mcp/base.py (4)
17-19: LGTM!Also applies to: 35-36
52-73: LGTM!
100-100: LGTM!
145-148: LGTM!tests/test_mcp_input_schema_limits.py (3)
6-7: LGTM!Also applies to: 52-52
62-75: LGTM!
77-95: LGTM!tests/test_3tool_aliases.py (1)
151-151: LGTM!Also applies to: 197-197, 327-327, 336-336, 345-345, 453-453, 468-476, 492-504
tests/test_mcp_palette.py (1)
45-50: LGTM!Also applies to: 61-61
tests/test_mcp_timeout.py (1)
16-16: LGTM!tests/test_mcp_warm_route.py (1)
272-272: LGTM!tests/test_phase3_digest.py (1)
305-305: LGTM!Also applies to: 349-349
tests/test_phase5.py (1)
305-305: LGTM!Also applies to: 331-331
tests/test_phase6_sentiment.py (1)
195-195: LGTM!tests/test_precompact_chunk_origin.py (1)
946-946: LGTM!Also applies to: 986-986
tests/test_search_fanout.py (1)
271-271: LGTM!Also applies to: 299-299
tests/test_audit_search_quality.py (1)
116-116: LGTM!Also applies to: 129-129
tests/test_brainbar_hybrid_helper.py (1)
183-183: LGTM!tests/test_brainstore.py (1)
587-587: LGTM!Also applies to: 597-597
tests/test_chunk_lifecycle.py (1)
267-275: LGTM!Also applies to: 314-314
tests/test_enrichment_controller.py (1)
1590-1590: LGTM!Also applies to: 1606-1606
tests/test_think_recall_integration.py (1)
278-278: LGTM!tests/test_tool_annotations.py (1)
91-121: LGTM!Also applies to: 136-136, 146-146
tests/test_write_queue.py (1)
1481-1481: LGTM!Also applies to: 2474-2474
tests/test_entity_type_sync.py (1)
78-86: LGTM!Also applies to: 97-104, 119-119
tests/test_issue_type.py (1)
166-186: LGTM!tests/test_mcp_digest_modes.py (1)
77-77: LGTM!Also applies to: 86-86
tests/test_mcp_labeled_field_output.py (1)
272-273: LGTM!tests/test_search_filter_params.py (1)
532-532: LGTM!Also applies to: 810-810
tests/test_search_validation.py (1)
51-51: LGTM!Also applies to: 66-66, 81-81, 134-134, 150-150
tests/test_smart_search_entity_dedup.py (1)
172-175: LGTM!Also applies to: 231-231, 383-383, 393-393
tests/test_store_handler.py (1)
179-179: LGTM!src/brainlayer/mcp/__init__.py (2)
176-217: LGTM!Also applies to: 469-470, 652-653, 678-783, 812-813, 996-1040, 1110-1251, 1273-1273, 1326-1328, 1334-1354, 1394-1403, 1568-1628
145-151: 🎯 Functional CorrectnessOutput-schema validation already skips error results.
src/brainlayer/mcp/_shared.py (1)
325-327: LGTM!scripts/smoke/repogolem_mcp_transport_smoke.py (1)
116-118: LGTM!src/brainlayer/brainbar_hybrid_helper.py (1)
191-196: LGTM!src/brainlayer/mcp/palette.py (1)
46-52: LGTM!
| "mcp>=2.0.0,<3.0.0", | ||
| "jsonschema>=4.20.0", # MCP v2 low-level handlers require explicit schema validation |
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win
Consider pinning an upper bound for jsonschema.
The mcp dependency uses <3.0.0 to guard against a future breaking major release. jsonschema>=4.20.0 has no upper bound. Add a similar cap for consistency and to avoid an untested future major version breaking schema validation.
♻️ Proposed pin
- "jsonschema>=4.20.0", # MCP v2 low-level handlers require explicit schema validation
+ "jsonschema>=4.20.0,<5.0.0", # MCP v2 low-level handlers require explicit schema validation📝 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.
| "mcp>=2.0.0,<3.0.0", | |
| "jsonschema>=4.20.0", # MCP v2 low-level handlers require explicit schema validation | |
| "mcp>=2.0.0,<3.0.0", | |
| "jsonschema>=4.20.0,<5.0.0", # MCP v2 low-level handlers require explicit schema validation |
🤖 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 `@pyproject.toml` around lines 19 - 20, Update the jsonschema dependency
specification in pyproject.toml to add an upper bound excluding the next major
version, matching the existing mcp dependency constraint while preserving the
current minimum version.
| def _normalize_call_tool_result(result: Any) -> CallToolResult: | ||
| """Preserve the result forms accepted by the MCP v1 call-tool decorator.""" | ||
| if isinstance(result, CallToolResult): | ||
| return result | ||
| if isinstance(result, tuple) and len(result) == 2: | ||
| content, structured_content = result | ||
| elif isinstance(result, dict): | ||
| content = [TextContent(type="text", text=json.dumps(result, indent=2))] | ||
| structured_content = result | ||
| elif hasattr(result, "__iter__"): | ||
| content = result | ||
| structured_content = None | ||
| else: | ||
| raise TypeError(f"Unexpected return type from tool: {type(result).__name__}") | ||
| return CallToolResult(content=list(content), structured_content=structured_content, is_error=False) |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win
Tighten the fallback branch in _normalize_call_tool_result.
elif hasattr(result, "__iter__") also matches plain strings and tuples of any length other than 2. A bare string return gets converted to list(content), which splits it into one content item per character instead of raising TypeError. A tuple with 1 or 3+ elements is silently treated as a content list rather than being rejected. The docstring states this function should "preserve the result forms accepted by the MCP v1 call-tool decorator," which did not include bare strings or mismatched tuples.
🐛 Proposed fix
elif hasattr(result, "__iter__"):
+ elif isinstance(result, list):
content = result
structured_content = None🧰 Tools
🪛 ast-grep (0.45.0)
[info] 121-121: use jsonify instead of json.dumps for JSON output
Context: json.dumps(result, indent=2)
Note: [CWE-116] Improper Encoding or Escaping of Output.
(use-jsonify)
🤖 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 `@src/brainlayer/mcp/__init__.py` around lines 115 - 129, Update
_normalize_call_tool_result to reject bare strings and tuples whose length is
not exactly two, rather than routing them through the iterable fallback.
Restrict the iterable branch to supported iterable content forms, preserving
valid two-item tuple, dict, and other accepted collection handling, and raise
TypeError for unsupported values.
| async def _handle_call_tool_request(_ctx: Any, params: CallToolRequestParams) -> CallToolResult: | ||
| """Validate and route a typed MCP v2 tool call while preserving v1 behavior.""" | ||
| arguments = params.arguments or {} | ||
| try: | ||
| tool = next((candidate for candidate in await list_tools() if candidate.name == params.name), None) | ||
| if tool is not None: | ||
| try: | ||
| jsonschema.validate(instance=arguments, schema=tool.input_schema) | ||
| except jsonschema.ValidationError as exc: | ||
| return _error_result(f"Input validation error: {exc.message}") | ||
|
|
||
| result = _normalize_call_tool_result(await call_tool(params.name, arguments)) | ||
|
|
There was a problem hiding this comment.
🚀 Performance & Scalability | 🔵 Trivial | ⚡ Quick win
Cache the static tool definitions to avoid rebuilding them on every tool call.
await list_tools() here rebuilds the full 13-tool inventory (_full_tool_definitions()), including the large brain_search/brain_recall schemas, on every call_tool invocation solely to find one tool's schema for validation. Since _full_tool_definitions() has no external state dependency, cache it once instead of reconstructing it per request.
♻️ Proposed fix
+import functools
+
def _full_tool_definitions() -> list[Tool]:
"""Build the complete canonical Python MCP tool inventory."""
+
+
+_full_tool_definitions = functools.lru_cache(maxsize=1)(_full_tool_definitions)Or decorate the function directly with @functools.lru_cache(maxsize=1) at its definition site (line 461). ToolPalette.expose() still filters the cached list dynamically per session, so palette expansion state is unaffected.
🤖 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 `@src/brainlayer/mcp/__init__.py` around lines 132 - 144, Cache the static tool
definitions used by _handle_call_tool_request instead of awaiting list_tools()
for every call. Apply a single-entry cache to _full_tool_definitions() or the
corresponding definition-site function, while preserving ToolPalette.expose()’s
dynamic session filtering and existing schema validation behavior.
| except Exception as exc: | ||
| return _error_result(str(exc)) |
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win
Log the exception before converting it to an error result.
except Exception as exc: return _error_result(str(exc)) silently swallows the traceback. logger is already defined in this module (Line 12) but unused here. Without a log entry, diagnosing an unexpected tool failure in production requires reproducing it from the client-visible error string alone.
♻️ Proposed fix
except Exception as exc:
+ logger.exception("MCP call_tool failed for %s", params.name)
return _error_result(str(exc))📝 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.
| except Exception as exc: | |
| return _error_result(str(exc)) | |
| except Exception as exc: | |
| logger.exception("MCP call_tool failed for %s", params.name) | |
| return _error_result(str(exc)) |
🤖 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 `@src/brainlayer/mcp/__init__.py` around lines 153 - 154, Update the exception
handler around _error_result to log the caught exception with the module’s
existing logger before returning the client-facing error result, preserving the
current str(exc) response behavior.
1 participant
Summary
on_*callbacks(content, structured)tool resultsmcp>=2.0.0,<3.0.0and explicitjsonschemavalidation supportThis is option (a) from the task brief: the migration was bounded, so the repository now follows the current SDK instead of retaining the emergency
<2pin.Reproduction and collection evidence
Before, in an isolated environment resolving
mcp==2.0.0:After:
Test API rewrites
Two tests required API-surface rewrites without changing their behavioral expectations:
test_mcp_input_schema_limits.pynow calls the server through MCP v2's supported in-memoryClientinstead of the removedrequest_handlersimplementation detail; it also covers combined-content normalization through a real client.test_mcp_labeled_field_output.pychecks the proprietary result-size hint inTool._meta, because MCP v2 intentionally drops unknownToolAnnotationsfields.All other test edits are canonical v2 model attribute renames (
input_schema,is_error,structured_content, and annotation hint names).Verification
The local CodeRabbit CLI remained in
reviewingbeyond its bounded three-minute gate and was stopped; reviewer assignment is intentionally left to the lead.Note
Medium Risk
Touches the primary Claude Code integration surface and a major MCP dependency bump; new schema validation can reject inputs that previously slipped through, though behavior is largely preserved via normalization and broad test updates.
Overview
Migrates the BrainLayer MCP server from SDK v1 decorators to the v2 low-level
ServerAPI (on_list_tools,on_call_tool,on_completion) and bumpsmcpto>=2.0.0,<3.0.0with explicitjsonschemafor tool I/O checks.The v2 call path adds
_handle_call_tool_request: validate arguments against each tool’sinput_schema, route through existingcall_tool, normalize legacy return shapes via_normalize_call_tool_result(tuples, dicts, iterables →CallToolResult), and validatestructured_contentwhenoutput_schemais set. Errors useis_error=Trueinstead ofisError.Model/API renames are applied repo-wide:
input_schema,output_schema,structured_content,read_only_hint, etc. Read tools moveanthropic/maxResultSizeCharsfromToolAnnotationstoTool.meta.Tests and harness:
tests/mock_mcp/base.pyuses in-memoryClient(..., mode="legacy")instead ofcreate_connected_server_and_client_session; schema/validation tests call the real server through that client rather thanrequest_handlers.Reviewed by Cursor Bugbot for commit 7fbc9fe. Bugbot is set up for automated code reviews on this repo. Configure here.
Note
Migrate MCP server to SDK v2 with snake_case field names and low-level handler API
mcpdependency to>=2.0.0,<3.0.0in pyproject.toml and addsjsonschema>=4.20.0for input/output schema validation.on_list_tools,on_call_tool,on_completion) in src/brainlayer/mcp/init.py.inputSchema→input_schema,isError→is_error,structuredContent→structured_content,readOnlyHint→read_only_hint, etc._normalize_call_tool_resultto coerce varied tool return types into a consistentCallToolResult, and validates arguments against declaredinput_schemausing jsonschema before dispatch.Clientinstead ofClientSession.Macroscope summarized 7fbc9fe.
Summary by CodeRabbit
New Features
Bug Fixes
Tests