fix(mcp): forward upstream initialize instructions on cold gateway init#28231
Conversation
Greptile SummaryThis PR fixes the MCP gateway not forwarding upstream
Confidence Score: 5/5Safe 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.
|
| 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
Codecov Report❌ Patch coverage is
📢 Thoughts on this report? Let us know! |
PR overviewMCP initialize instruction prefetch addedThis 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
Risk: 2/10 |
117d769 to
10a4126
Compare
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]>
10a4126 to
1f3b967
Compare
|
@greptileai checka again |
|
@milan-berri Can you share a video of this working? |
|
@Sameerlite sharing screenshots |
…ream_initialize_instructions
…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]>



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
tests/test_litellm/directory, Adding at least 1 test is a hard requirement - see detailsmake test-unit@greptileaiand received a Confidence Score of at least 4/5 before requesting a maintainer reviewDelays 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)
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.instructionswas only merged from YAML/DB at firstinitialize, 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 coldinitialize.What did not work (before fix):
instructionson gatewayinitializeinitialize(fresh gateway, no YAML)Noneinitialize→list_tools→call_toolNone(first init result unchanged; client cannot re-init)list_toolswarmed cacheinstructionsin configinitializeLive verify (after fix): FastMCP upstream with
instructions=→ LiteLLM gateway (no YAMLinstructions) → cold and subsequentinitializeboth return upstream text.Setup: local FastMCP on
:18801, gatewaymcp_serversHTTP to upstream on:18803,Authorization: Bearer sk-123,x-mcp-servers: upstream_code_instructions(venv_berriaieditable install ofberriai/litellm).Type
🐛 Bug Fix
Changes
Root cause
_merge_gateway_initialize_instructionsonly read YAML/DB overrides or_upstream_initialize_instructions_by_server_idat gatewayinitializetime. The cache was filled when the proxy opened a separate upstream session (list_tools,call_tool, health check)—not during the client's first gatewayinitialize. Solist_toolsin 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 upstreamrun_with_session(noop) to captureInitializeResult.instructionsinto_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_instructionsas before.return_exceptions=Trueis required because anyio's task-group teardown on connection refused surfaces asCancelledError; without it, a single bad upstream would 500 the entire gateway initialize.Skips upstream connect (no-op) when:
spec_path)instructionsalready setserver_idrequires_per_user_auth(same as health check — no user context in this path)authentication_tokenis missing (exceptnone/aws_sigv4)server_idis withinMCP_HEALTH_CHECK_TIMEOUT(default 10s, env-tunable viaLITELLM_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
instructionsstill override upstream.server_idper cooldown window. Successful, non-empty caches never expire. Empty/failed probes are remembered forMCP_HEALTH_CHECK_TIMEOUTseconds (default 10).list_tools,call_tool, health check) unchanged.list_tools/call_toolto populate cache).load_servers_from_configandreload_servers_from_databaseso config/DB changes re-probe immediately.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_configclears 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 passedLive regression check
Live-tested both paths against the editable
berriai/litellmcheckout:initializeboth return upstream text;list_toolsunchanged.initializeno 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 withCancelledError.