fix(wren): align MCP handlers with ServeContext and query limits by ttw225 · Pull Request #2500 · Canner/WrenAI

ttw225 · 2026-07-14T06:41:18Z

Summary

#2438 introduced wren serve mcp and centralized the selected project,
engine, and capability flags in ServeContext.

This follow-up builds on that foundation by using the startup context
consistently across memory tools and workflow guidance, refining cube query
limit handling, and adding dedicated MCP test coverage.

Changes

Use ServeContext.project as the source of truth for all MCP memory paths,
keeping memory access aligned with the project selected through --project.
Generate the wren_workflow prompt from the active capabilities.
Embed the cube N+1 truncation probe in generated SQL and execute it without
applying a second connector limit:
- preserve the caller's limit in SQL-only output
- keep the truncation probe bounded in generated SQL regardless of how
  individual connectors implement their limit argument
- preserve valid LIMIT/OFFSET ordering for MySQL and Doris
Reject negative limits consistently for run_sql and query_cube.
Add 14 handler-level MCP tests and a dedicated test-mcp CI job with the
optional mcp dependency.

Workflow capabilities

Server flags	Query guidance	`store_query`
Default	`dry_run` → `run_sql`, with `query_cube` for metrics	Omitted
`--allow-write`	`dry_run` → `run_sql`, with `query_cube` for metrics	Included
`--no-connect`	`dry_plan`	Omitted
`--no-connect --allow-write`	`dry_plan`	Included

The generated workflow only references tools registered for the active
server configuration, with contiguous step numbering in every mode.

Test plan

uv lock --check
uv run --no-sync ruff format --check src/wren/mcp_server.py tests/unit/test_mcp_server.py
uv run --no-sync ruff check src/wren/mcp_server.py tests/unit/test_mcp_server.py
uv run --no-sync pytest tests/unit/test_mcp_server.py -q
- 14 passed
uv run --no-sync pytest tests/unit -q
- 1007 passed, 43 skipped

Summary by CodeRabbit

Bug Fixes
- Improved MCP memory storage path handling to correctly use the project’s .wren/memory directory.
- Added validation to reject negative limit values for SQL and cube queries.
- Refined cube query truncation behavior (including correct LIMIT/OFFSET handling) and consistent truncated/row_count reporting, including SQL-only output.
- Updated MCP workflow prompts to dynamically match connection and write capabilities.
Tests
- Expanded unit tests covering memory paths, workflow prompt generation, negative-limit validation, and cube truncation logic.
Chores
- Split MCP server tests into a dedicated CI job for clearer, targeted execution.

Use ServeContext as the source of truth for MCP memory paths and workflow capabilities. This keeps explicit project selection and registered tools aligned throughout the server. Embed the cube N+1 probe in generated SQL so truncation detection works without applying a second connector limit. Preserve the requested limit in SQL-only output and reject negative limits for MCP query tools. Add handler-level MCP tests and a dedicated CI job that installs the mcp extra.

coderabbitai · 2026-07-14T06:41:40Z

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: Repository UI

Review profile: CHILL

Plan: Pro

Run ID: ed3a62a9-221d-4c76-9cc5-cfc3b5f1e526

📥 Commits

Reviewing files that changed from the base of the PR and between 31156b6 and 525b0f3.

📒 Files selected for processing (1)

.github/workflows/wren-ci.yml

🚧 Files skipped from review as they are similar to previous changes (1)

.github/workflows/wren-ci.yml

Walkthrough

The MCP server now uses project-scoped memory paths, context-dependent workflow prompts, validated query limits, and SQL-level cube-query probing. Dedicated CI execution covers the MCP server tests.

Changes

MCP server behavior

Layer / File(s)	Summary
Cube query limit probing `core/wren/src/wren/mcp_server.py`, `core/wren/tests/unit/test_mcp_server.py`	Cube queries reject negative limits, preserve offsets, embed probe limits in SQL, and report truncation metadata.
Context-aware memory and workflow `core/wren/src/wren/mcp_server.py`, `core/wren/tests/unit/test_mcp_server.py`	Memory tools use `ServeContext.project`, and workflow prompts vary with connection and write permissions.
Dedicated MCP test execution `.github/workflows/wren-ci.yml`, `core/wren/tests/unit/test_mcp_server.py`	MCP tests are excluded from the general unit job and run in a dedicated MCP-enabled job.

Estimated code review effort: 4 (Complex) | ~45 minutes

Sequence Diagram(s)

sequenceDiagram
  participant MCPClient
  participant query_cube
  participant SQLBuilder
  participant Connector
  MCPClient->>query_cube: Submit cube query with limit and offset
  query_cube->>SQLBuilder: Build SQL with probe limit
  query_cube->>Connector: Execute SQL with limit=None
  Connector-->>query_cube: Return probe rows
  query_cube-->>MCPClient: Return rows and truncation metadata

Possibly related PRs

Canner/WrenAI#2438: Modifies the same MCP server tooling and behavior.
Canner/WrenAI#2439: Updates related MCP workflow prompt generation.

Suggested labels: dependencies

Suggested reviewers: yichieh-lu

Poem

I’m a rabbit with queries tucked neatly away,
Project paths guide where memories stay.
Probes count rows; prompts change their tune,
MCP tests hop through CI beneath the moon.

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.
Title check	✅ Passed	The title clearly summarizes the main changes: MCP handlers now align with ServeContext and query-limit behavior.
Docstring Coverage	✅ Passed	No functions found in the changed files to evaluate docstring coverage. Skipping docstring coverage check.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.

✨ Finishing Touches

🧪 Generate unit tests (beta)

Create PR with unit tests

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.

❤️ Share

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

coderabbitai

🧹 Nitpick comments (1)

.github/workflows/wren-ci.yml (1)
221-221: 🔒 Security & Privacy | 🔵 Trivial | 💤 Low value

Consider disabling credential persistence in checkout.

To prevent the GITHUB_TOKEN from being persisted in the local git config, it's a security best practice to set persist-credentials: false when checking out the repository, especially if the .git directory could ever be uploaded as an artifact or accessed by third-party tools.
🛠️ Proposed fix
-      - uses: actions/checkout@v4
+      - uses: actions/checkout@v4
+        with:
+          persist-credentials: false
🤖 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 @.github/workflows/wren-ci.yml at line 221, Update the actions/checkout@v4
step to disable credential persistence by setting persist-credentials to false,
ensuring the checkout does not retain the GITHUB_TOKEN in local git
configuration.
Source: Linters/SAST tools

🤖 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.

Nitpick comments:
In @.github/workflows/wren-ci.yml:
- Line 221: Update the actions/checkout@v4 step to disable credential
persistence by setting persist-credentials to false, ensuring the checkout does
not retain the GITHUB_TOKEN in local git configuration.

ℹ️ Review info

⚙️ Run configuration

Configuration used: Repository UI

Review profile: CHILL

Plan: Pro

Run ID: c8904e73-4679-470b-990d-5d68a04efac2

📥 Commits

Reviewing files that changed from the base of the PR and between e4f0139 and 31156b6.

📒 Files selected for processing (3)

.github/workflows/wren-ci.yml
core/wren/src/wren/mcp_server.py
core/wren/tests/unit/test_mcp_server.py

github-actions Bot added python Pull requests that update Python code core ci labels Jul 14, 2026

coderabbitai Bot reviewed Jul 14, 2026

View reviewed changes

ci(wren): disable credential persistence in MCP tests

525b0f3

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

fix(wren): align MCP handlers with ServeContext and query limits#2500

fix(wren): align MCP handlers with ServeContext and query limits#2500
ttw225 wants to merge 2 commits into
Canner:mainfrom
ttw225:fix/mcp-project-memory-path

ttw225 commented Jul 14, 2026 •

edited by coderabbitai Bot

Loading

Uh oh!

coderabbitai Bot commented Jul 14, 2026 •

edited

Loading

Uh oh!

coderabbitai Bot left a comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant

Uh oh!

Conversation

ttw225 commented Jul 14, 2026 • edited by coderabbitai Bot Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary

Changes

Workflow capabilities

Test plan

Summary by CodeRabbit

Uh oh!

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

Uh oh!

Walkthrough

Changes

Sequence Diagram(s)

Uh oh!

coderabbitai Bot left a comment

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant

ttw225 commented Jul 14, 2026 •

edited by coderabbitai Bot

Loading

coderabbitai Bot commented Jul 14, 2026 •

edited

Loading