refactor(mcp): introduce McpServerBinding, unify ensure functions, fixethrowaway session

[key4ng]

Description

Problem

1. MCP server bindings use an opaque (String, String) tuple for (label, server_key) across ~20 call sites, making code hard to read and prone to field-swap bugs.
2. Three separate functions duplicate the same MCP server connection logic: ensure_request_mcp_client (OpenAI), ensure_connection (Anthropic), and ensure_mcp_connection (gRPC).
3. The Anthropic router creates a throwaway McpToolSession in ensure_connection solely to inject tools into the request, then discards it. The handler then creates a second session for actual tool execution.

Solution

1. Replace all (String, String) tuples with a named McpServerBinding { label, server_key } struct.
2. Introduce a unified ensure_mcp_servers() in mcp_utils.rs that handles both dynamic server connection and builtin tool routing. Existing callers delegate to it.
3. Split the Anthropic ensure_connection into connect (router level, via nsure_mcp_servers) and inject (handler level, using the real session). Delete the throwaway session pattern.

Changes

• Define McpServerBinding struct in mcp/src/core/session.rs, export from crate root
  • Migrate all Vec<(String, String)> fields, params, and returns to Vec across MCP crate, Anthropic, OpenAI, and gRPC routers (12 files)
  • Add ensure_mcp_servers(orchestrator, inputs, builtin_types) and extract_builtin_types() in mcp_utils.rs
  • Refactor ensure_request_mcp_client to delegate to ensure_mcp_servers
  • Delete ensure_connection from anthropic/mcp.rs; move tool injection to non_streaming.rs and streaming.rs where the real McpToolSession is created
  • Make inject_mcp_tools_into_request and collect_allowed_tools_from_toolsets pub(crate)

Test Plan

Checklist

• cargo +nightly fmt passes
• cargo clippy --all-targets --all-features -- -D warnings passes
• (Optional) Documentation updated

Summary by CodeRabbit

• New Features

  • Introduced an explicit server-binding model for MCP servers, plus utilities to manage and connect them.
  • Added a configurable default iteration limit to make MCP tool loops safer.

• Refactor

  • Replaced tuple-based server handling with the new structured server representation across routers and utilities.
  • Simplified and unified MCP tool injection and connection flows for more consistent behavior.

[coderabbitai]

📝 Walkthrough

Walkthrough

This PR introduces a new public struct McpServerBinding { label: String, server_key: String } and replaces tuple-based MCP server representations (String, String) with McpServerBinding across core session/orchestrator code and router plumbing, updating signatures, imports, and iteration logic accordingly.

Changes

Cohort / File(s)	Summary
MCP Core Library mcp/src/core/mod.rs, mcp/src/core/session.rs, mcp/src/core/orchestrator.rs, mcp/src/lib.rs	Adds and re-exports McpServerBinding; replaces Vec<(String,String)> with Vec<McpServerBinding> in session/orchestrator APIs and internal logic; updates access from tuple indexing to binding.label / binding.server_key.
MCP Utils model_gateway/src/routers/mcp_utils.rs	Introduces McpServerBinding, DEFAULT_MAX_ITERATIONS, extract_builtin_types, ensure_mcp_servers, and ensure_request_mcp_client; changes connect_mcp_servers/related returns to Vec<McpServerBinding> and refactors server-connect logic to construct bindings.
Anthropic Router model_gateway/src/routers/anthropic/context.rs, model_gateway/src/routers/anthropic/mcp.rs, model_gateway/src/routers/anthropic/non_streaming.rs, model_gateway/src/routers/anthropic/streaming.rs, model_gateway/src/routers/anthropic/router.rs	Changes RequestContext.mcp_servers to Option<Vec<McpServerBinding>>; exposes inject_mcp_tools_into_request and collect_allowed_tools_from_toolsets; removes some prior ensure_connection orchestration and uses ensure_mcp_servers flow; injects MCP tools before worker calls.
gRPC Routers model_gateway/src/routers/grpc/.../responses/utils.rs, model_gateway/src/routers/grpc/harmony/.../non_streaming.rs, model_gateway/src/routers/grpc/harmony/.../streaming.rs, model_gateway/src/routers/grpc/regular/.../non_streaming.rs, model_gateway/src/routers/grpc/regular/.../streaming.rs	Updates signatures and imports to accept Vec<McpServerBinding> instead of Vec<(String,String)>; adjusts iterations and lookups to use binding.label / binding.server_key.
OpenAI Router model_gateway/src/routers/openai/responses/mcp.rs, model_gateway/src/routers/openai/responses/streaming.rs	Replaces tuple-based MCP server handling with McpServerBinding iteration and uses binding.label/binding.server_key when constructing MCP list-tools metadata.

Sequence Diagram(s)

(Skipped — changes are refactors of data representation and signatures, not a new multi-component control-flow feature requiring a sequence diagram.)

Estimated code review effort

🎯 4 (Complex) | ⏱️ ~45 minutes

Possibly related PRs

• feat(message api): MCP tool use with streaming and non-streaming support #352: Touches the same MCP server representation and replaces tuple server pairs with a structured binding type across orchestrator and router call sites.
• refactor(mcp): remove requested_servers and introduce ResponsesCallContext #382: Refactors MCP server plumbing and per-request mcp_servers passing—overlaps with signature and call-site changes to accept McpServerBinding.
• fix(mcp): filter builtin servers from mcp_list_tools output #450: Modifies MCP server representation and McpToolSession usage; closely related to changes in session/server binding handling.

Suggested reviewers

• slin1237
• CatherineSue

Poem

🐰 I hopped through tuples, one by one,
Found label and key beneath the sun,
I stitched them cozy, bound in a struct,
Now servers shout their names—how plucky, how tucked! 🥕

🚥 Pre-merge checks | ✅ 3

✅ Passed checks (3 passed)

Check name	Status	Explanation
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.
Docstring Coverage	✅ Passed	Docstring coverage is 100.00% which is sufficient. The required threshold is 80.00%.
Title check	✅ Passed	The title directly addresses the main changes: introducing McpServerBinding, unifying ensure functions, and fixing the throwaway session pattern in Anthropic.

_{✏️ Tip: You can configure your own custom pre-merge checks in the settings.}

✨ Finishing Touches

• 📝 Generate docstrings (stacked PR)
• 📝 Generate docstrings (commit on current branch)

🧪 Generate unit tests (beta)

• Create PR with unit tests
• Post copyable unit tests in a comment
• Commit unit tests in branch keyang/mcp-session

Tip

Issue Planner is now in beta. Read the docs and try it out! Share your feedback on Discord.

---

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

[gemini-code-assist]

Summary of Changes

Hello @key4ng, I'm Gemini Code Assist¹! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!

This pull request significantly refactors the Multi-Cloud Proxy (MCP) integration by introducing a dedicated struct for server bindings, unifying the server connection logic across different routers, and optimizing the tool injection process within the Anthropic router. These changes aim to enhance code clarity, reduce redundancy, and improve the overall robustness of MCP server interactions.

Highlights

• Introduced McpServerBinding Struct: Replaced opaque (String, String) tuples representing MCP server labels and keys with a dedicated McpServerBinding struct, improving code readability and preventing field-swap bugs across the MCP core and various router implementations.
• Unified MCP Server Connection Logic: Consolidated disparate MCP server connection functions (ensure_request_mcp_client, ensure_connection, ensure_mcp_connection) into a single, unified ensure_mcp_servers function within mcp_utils.rs, reducing code duplication and simplifying server management.
• Refactored Anthropic Tool Injection: Eliminated the pattern of creating and discarding a 'throwaway' McpToolSession in the Anthropic router. Tool injection now occurs at the handler level using the actual session, streamlining the tool execution flow.

Changelog

• mcp/src/core/mod.rs
  • Exported the new McpServerBinding struct.
• mcp/src/core/orchestrator.rs
  • Imported McpServerBinding for use within the orchestrator.
  • Updated execute_tool_calls to accept &[McpServerBinding] instead of &[(String, String)].
• mcp/src/core/session.rs
  • Defined the McpServerBinding struct to encapsulate server label and key.
  • Updated McpToolSession fields and constructor to use Vec<McpServerBinding>.
  • Modified various methods within McpToolSession to access label and server_key from McpServerBinding.
  • Updated test cases to instantiate McpServerBinding structs directly.
• mcp/src/lib.rs
  • Re-exported McpServerBinding from the mcp crate root.
• model_gateway/src/routers/anthropic/context.rs
  • Imported McpServerBinding.
  • Updated the mcp_servers field in RequestContext to use Option<Vec<McpServerBinding>>.
• model_gateway/src/routers/anthropic/mcp.rs
  • Removed the ensure_connection function, delegating its responsibilities to mcp_utils.
  • Changed visibility of inject_mcp_tools_into_request and collect_allowed_tools_from_toolsets to pub(crate).
• model_gateway/src/routers/anthropic/non_streaming.rs
  • Added calls to mcp::collect_allowed_tools_from_toolsets and mcp::inject_mcp_tools_into_request to handle tool injection for non-streaming requests.
• model_gateway/src/routers/anthropic/router.rs
  • Replaced the call to the removed mcp::ensure_connection with a direct call to mcp_utils::ensure_mcp_servers for connecting MCP servers.
  • Imported mcp_utils and HashMap.
• model_gateway/src/routers/anthropic/streaming.rs
  • Added calls to mcp::collect_allowed_tools_from_toolsets and mcp::inject_mcp_tools_into_request for tool injection in streaming requests.
• model_gateway/src/routers/grpc/common/responses/utils.rs
  • Imported McpServerBinding.
  • Updated ensure_mcp_connection to return Result<(bool, Vec<McpServerBinding>), Response>.
• model_gateway/src/routers/grpc/harmony/responses/non_streaming.rs
  • Imported McpServerBinding.
  • Updated execute_with_mcp_loop to accept Vec<McpServerBinding>.
• model_gateway/src/routers/grpc/harmony/responses/streaming.rs
  • Imported McpServerBinding.
  • Updated execute_mcp_tool_loop_streaming to accept Vec<McpServerBinding>.
  • Modified tool listing logic to use McpServerBinding fields.
• model_gateway/src/routers/grpc/regular/responses/non_streaming.rs
  • Imported McpServerBinding.
  • Updated execute_tool_loop to accept Vec<McpServerBinding>.
• model_gateway/src/routers/grpc/regular/responses/streaming.rs
  • Imported McpServerBinding.
  • Updated execute_tool_loop_streaming and execute_tool_loop_streaming_internal to accept Vec<McpServerBinding>.
  • Modified tool listing logic to use McpServerBinding fields.
• model_gateway/src/routers/mcp_utils.rs
  • Imported McpServerBinding.
  • Updated connect_mcp_servers to return Vec<McpServerBinding> and construct McpServerBinding instances.
  • Introduced extract_builtin_types to parse built-in tool types from ResponseTool arrays.
  • Added ensure_mcp_servers as a unified function for connecting MCP servers, handling both dynamic/static and built-in types.
  • Refactored ensure_request_mcp_client to delegate to the new ensure_mcp_servers function.
  • Updated test cases to use McpServerBinding fields for assertions.
• model_gateway/src/routers/openai/responses/mcp.rs
  • Updated inject_mcp_metadata_streaming and build_incomplete_response to iterate over McpServerBinding and use its label and server_key fields.
• model_gateway/src/routers/openai/responses/streaming.rs
  • Imported McpServerBinding.
  • Updated handle_streaming_with_tool_interception to accept Vec<McpServerBinding>.
  • Modified send_mcp_list_tools_events to use McpServerBinding fields.

Using Gemini Code Assist

The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.

Invoking Gemini

You can request assistance from Gemini at any point by creating a comment using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands on the current page.

Feature	Command	Description
Code Review	/gemini review	Performs a code review for the current pull request in its current state.
Pull Request Summary	/gemini summary	Provides a summary of the current pull request in its current state.
Comment	@gemini-code-assist	Responds in comments when explicitly tagged, both in pull request comments and review comments.
Help	/gemini help	Displays a list of available commands.

Customization

To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.

Limitations & Feedback

Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here.

You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension.

Footnotes

1. Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩

[gemini-code-assist]

Code Review

This pull request is a significant and well-executed refactoring. It introduces the McpServerBinding struct to replace opaque (String, String) tuples, which greatly improves code readability and type safety. The unification of the various ensure_* connection functions into a single ensure_mcp_servers is a great architectural improvement that reduces code duplication and centralizes connection logic. Additionally, removing the throwaway McpToolSession in the Anthropic router simplifies the control flow.

I've made a couple of minor suggestions to use collect for more idiomatic and concise code. Overall, this is a high-quality change that improves the maintainability of the MCP codebase.

[coderabbitai]

Actionable comments posted: 1

🤖 Prompt for all review comments with AI agents

Verify each finding against the current code and only fix it if needed.

Inline comments:
In `@mcp/src/core/session.rs`:
- Around line 27-37: Add `PartialEq` and `Eq` derives to the McpServerBinding
struct so it can be directly compared; update the derive attribute on the
McpServerBinding declaration (currently #[derive(Debug, Clone)]) to include
PartialEq and Eq (e.g., #[derive(Debug, Clone, PartialEq, Eq)]) so tests and
duplicate-checking can use direct equality and collection helpers like contains.

[coderabbitai]

Actionable comments posted: 1

🤖 Prompt for all review comments with AI agents

Verify each finding against the current code and only fix it if needed.

Inline comments:
In `@model_gateway/src/routers/mcp_utils.rs`:
- Around line 188-197: extract_builtin_types currently collects BuiltinToolType
entries directly and can return duplicates when multiple ResponseTool items
share the same builtin type; change the implementation to deduplicate before
returning (e.g., map ResponseTool -> BuiltinToolType, insert into a
HashSet<BuiltinToolType> or use iterator.collect::<HashSet<_>>(), then convert
to Vec), keeping the same function signature and matching arms for
ResponseToolType::WebSearchPreview and ::CodeInterpreter so callers (and
ensure_mcp_servers) receive only unique builtin types.

---

Duplicate comments:
In `@mcp/src/core/session.rs`:
- Around line 27-37: No changes required—McpServerBinding is correctly defined
and documented with derives PartialEq and Eq; keep the struct as-is
(McpServerBinding with fields label and server_key) and proceed with the
approved change.