Skip to content

fix(mcp): forward upstream initialize instructions on cold gateway init#28231

Merged
mateo-berri merged 2 commits into
litellm_internal_stagingfrom
litellm_mcp_gateway_upstream_initialize_instructions
May 22, 2026
Merged

fix(mcp): forward upstream initialize instructions on cold gateway init#28231
mateo-berri merged 2 commits into
litellm_internal_stagingfrom
litellm_mcp_gateway_upstream_initialize_instructions

Conversation

@milan-berri

@milan-berri milan-berri commented May 19, 2026

Copy link
Copy Markdown
Collaborator

Report: MCP gateway does not forward upstream server instructions defined in MCP server code (e.g. McpServer / FastMCP instructions=). YAML/DB instructions work on first initialize; upstream code instructions do not. Calling list_tools / call_tool in the same MCP session also does not help—the client cannot re-initialize, and gateway instructions are fixed at first initialize. Upstream text can appear only on a new connection after list_tools has warmed the in-memory cache.

Pre-Submission checklist

Please complete all items before asking a LiteLLM maintainer to review your PR

  • I have Added testing in the tests/test_litellm/ directory, Adding at least 1 test is a hard requirement - see details
  • My PR passes all unit tests on make test-unit
  • My PR's scope is as isolated as possible, it only solves 1 specific problem
  • I have requested a Greptile review by commenting @greptileai and received a Confidence Score of at least 4/5 before requesting a maintainer review

Delays in PR merge?

If you're seeing a delay in your PR being merged, ping the LiteLLM Team on Slack (#pr-review).

CI (LiteLLM team)

CI status guideline:

  • 50-55 passing tests: main is stable with minor issues.
  • 45-49 passing tests: acceptable but needs attention
  • <= 40 passing tests: unstable; be careful with your merges and assess the risk.
  • Branch creation CI run
    Link:

  • CI run for the last commit
    Link:

  • Merge / cherry-pick CI run
    Links:

Screenshots / Proof of Fix

Problem: Gateway InitializeResult.instructions was only merged from YAML/DB at first initialize, or from an in-memory cache populated by a prior upstream session (list_tools, call_tool, health check). Upstream instructions set in MCP server code were not forwarded on cold initialize.

What did not work (before fix):

Step instructions on gateway initialize
Cold initialize (fresh gateway, no YAML) None
Same session: initializelist_toolscall_tool Still None (first init result unchanged; client cannot re-init)
New connection after list_tools warmed cache Upstream text (cache only)
YAML instructions in config Works on cold initialize

Live verify (after fix): FastMCP upstream with instructions= → LiteLLM gateway (no YAML instructions) → cold and subsequent initialize both return upstream text.

gateway COLD init: 'Example instructions from upstream MCP server code (FastMCP instructions=).'
gateway WARM init: 'Example instructions from upstream MCP server code (FastMCP instructions=).'

Setup: local FastMCP on :18801, gateway mcp_servers HTTP to upstream on :18803, Authorization: Bearer sk-123, x-mcp-servers: upstream_code_instructions (venv_berriai editable install of berriai/litellm).

Type

🐛 Bug Fix

Changes

Root cause

_merge_gateway_initialize_instructions only read YAML/DB overrides or _upstream_initialize_instructions_by_server_id at gateway initialize time. The cache was filled when the proxy opened a separate upstream session (list_tools, call_tool, health check)—not during the client's first gateway initialize. So list_tools in the same client session did not surface instructions; only a later connection could see cached upstream text.

Fix

  • MCPServerManager._ensure_upstream_initialize_instructions_cached(server) — if needed, one upstream run_with_session (noop) to capture InitializeResult.instructions into _upstream_initialize_instructions_by_server_id (same pattern as health check).
  • _gateway_initialize_instructions_request_scope — after resolving allowed servers, asyncio.gather(..., return_exceptions=True) the ensure call per server, then _merge_gateway_initialize_instructions as before. return_exceptions=True is required because anyio's task-group teardown on connection refused surfaces as CancelledError; without it, a single bad upstream would 500 the entire gateway initialize.

Skips upstream connect (no-op) when:

  • OpenAPI server (spec_path)
  • YAML/DB instructions already set
  • Cache already has text for server_id
  • requires_per_user_auth (same as health check — no user context in this path)
  • Auth type requires a static token but authentication_token is missing (except none / aws_sigv4)
  • A prior probe attempt for this server_id is within MCP_HEALTH_CHECK_TIMEOUT (default 10s, env-tunable via LITELLM_MCP_HEALTH_CHECK_TIMEOUT). The probe is a health-check-shaped operation and already uses this knob as its inner call timeout, so we reuse it as the cooldown rather than introducing a new env var. Applies to any prior attempt — success, empty upstream, or failure — so a server with no instructions configured never triggers a reconnect on every gateway initialize.

Behavior notes

  • YAML/DB instructions still override upstream.
  • Upstream prefetch runs once per server_id per cooldown window. Successful, non-empty caches never expire. Empty/failed probes are remembered for MCP_HEALTH_CHECK_TIMEOUT seconds (default 10).
  • Existing cache paths (list_tools, call_tool, health check) unchanged.
  • Per-user OAuth MCP servers are unchanged (still need YAML/DB or user-scoped list_tools / call_tool to populate cache).
  • Negative-cache map is cleared on load_servers_from_config and reload_servers_from_database so config/DB changes re-probe immediately.
  • Multi-pod: cache and cooldown are per-process. Each pod probes once on cold start, then honors the cooldown. No cross-pod coordination is needed for correctness; throughput-sensitive deployments can tune LITELLM_MCP_HEALTH_CHECK_TIMEOUT (also affects health-check call timeout).

Tests

  • TestEnsureUpstreamInitializeInstructionsCached — skip when YAML / cache / spec_path; runs session and caches when eligible; cooldown after empty upstream; cooldown after upstream failure; load_servers_from_config clears cooldown map.
  • TestMergeGatewayInitializeInstructions — existing merge behavior.
  • TestGatewayCreateInitializationOptions — ContextVar injection unchanged.
pytest tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py::TestEnsureUpstreamInitializeInstructionsCached \
       tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py::TestMergeGatewayInitializeInstructions \
       tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py::TestGatewayCreateInitializationOptions -q
# 20 passed

Live regression check

Live-tested both paths against the editable berriai/litellm checkout:

  • Happy path (upstream up): cold + warm gateway initialize both return upstream text; list_tools unchanged.
  • Failure path (upstream down): gateway initialize no longer 500s; 1st init ~0.2s (probe + swallow), 2nd/3rd ~0.01s (cooldown skip). Without the cooldown + return_exceptions=True, every init would re-open a refused upstream connection and crash with CancelledError.

@greptile-apps

greptile-apps Bot commented May 19, 2026

Copy link
Copy Markdown
Contributor

Greptile Summary

This PR fixes the MCP gateway not forwarding upstream InitializeResult.instructions on cold initialize — previously only YAML/DB overrides or a warm cache (populated by a prior list_tools/health-check session) were forwarded; upstream code-defined instructions were invisible to the first connection.

  • _ensure_upstream_initialize_instructions_cached opens one noop upstream session per server per cooldown window to capture InitializeResult.instructions, mirrors the existing health-check pattern, and stamps a _upstream_initialize_instructions_probed_at timestamp before the call so empty/failed probes never cause repeated reconnects.
  • _gateway_initialize_instructions_request_scope concurrently probes all allowed servers via asyncio.gather(..., return_exceptions=True) so a single bad upstream (anyio CancelledError on connection refused) cannot 500 the gateway initialize; both maps are cleared on load_servers_from_config / reload_servers_from_database so config changes re-probe immediately.

Confidence Score: 5/5

Safe to merge — the new probe path is narrowly scoped, mirrors the existing health-check code exactly, and cannot break the gateway initialize even when all upstreams are down.

The fix is self-contained: a new optional prefetch that runs concurrently and swallows all per-server errors. The cooldown timestamp is set before the first await, eliminating TOCTOU races in the single-threaded asyncio model. The previous review concern about infinite reconnects for empty upstreams is addressed by the probed-at map. Test coverage includes skip conditions, happy-path caching, empty-upstream cooldown, failure cooldown, and reload clearing the negative-cache — all exercised without real network calls.

No files require special attention.

Important Files Changed

Filename Overview
litellm/proxy/_experimental/mcp_server/mcp_server_manager.py Adds _ensure_upstream_initialize_instructions_cached with a _upstream_initialize_instructions_probed_at cooldown map to prefetch upstream InitializeResult.instructions once per server per cooldown window; clears both maps on config/DB reload.
litellm/proxy/_experimental/mcp_server/server.py Adds asyncio.gather with return_exceptions=True in _gateway_initialize_instructions_request_scope to run per-server prefetch probes concurrently before merging instructions; correctly swallows per-server CancelledError without cancelling the gateway initialize.
tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py Adds TestEnsureUpstreamInitializeInstructionsCached with 7 cases covering all skip conditions, happy path caching, cooldown after empty upstream, cooldown after failure, and reload clearing the negative-cache map.

Reviews (2): Last reviewed commit: "fix(mcp): forward upstream initialize in..." | Re-trigger Greptile

Comment thread litellm/proxy/_experimental/mcp_server/mcp_server_manager.py
@codecov

codecov Bot commented May 19, 2026

Copy link
Copy Markdown

Codecov Report

❌ Patch coverage is 86.66667% with 4 lines in your changes missing coverage. Please review.

Files with missing lines Patch % Lines
...oxy/_experimental/mcp_server/mcp_server_manager.py 89.28% 3 Missing ⚠️
litellm/proxy/_experimental/mcp_server/server.py 50.00% 1 Missing ⚠️

📢 Thoughts on this report? Let us know!

@veria-ai

veria-ai Bot commented May 19, 2026

Copy link
Copy Markdown
Contributor

PR overview

MCP initialize instruction prefetch added

This PR caches upstream MCP initialize instructions during gateway initialization and throttles failed or empty probes per server. I checked the new prefetch path against the existing allowed-server filtering, per-user auth skips, static auth handling, and MCP request auth flow and did not find an attacker-reachable security issue.

Security review

  • No new security issues were flagged in the latest review.
  • No review issues remain open on this pull request.

Risk: 2/10

Comment thread litellm/proxy/_experimental/mcp_server/mcp_server_manager.py
@milan-berri milan-berri force-pushed the litellm_mcp_gateway_upstream_initialize_instructions branch from 117d769 to 10a4126 Compare May 19, 2026 09:14
Prefetch upstream InitializeResult.instructions before merging gateway
initialize options when YAML/DB do not set instructions, so clients receive
upstream server text on the first MCP initialize without list_tools.

Co-authored-by: Cursor <[email protected]>
@milan-berri milan-berri force-pushed the litellm_mcp_gateway_upstream_initialize_instructions branch from 10a4126 to 1f3b967 Compare May 19, 2026 09:25
@milan-berri

Copy link
Copy Markdown
Collaborator Author

@greptileai checka again

@Sameerlite

Copy link
Copy Markdown
Collaborator

@milan-berri Can you share a video of this working?

@milan-berri

Copy link
Copy Markdown
Collaborator Author

@Sameerlite sharing screenshots
image
image
image

@mateo-berri mateo-berri left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

LGTM; thanks!

@mateo-berri mateo-berri merged commit 7270f72 into litellm_internal_staging May 22, 2026
117 checks passed
fzowl pushed a commit to fzowl/litellm that referenced this pull request Jun 24, 2026
…it (BerriAI#28231)

Prefetch upstream InitializeResult.instructions before merging gateway
initialize options when YAML/DB do not set instructions, so clients receive
upstream server text on the first MCP initialize without list_tools.

Co-authored-by: Cursor <[email protected]>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants