From 9074f4db3f0a3a9324e53f962e0fefb5ec38eadb Mon Sep 17 00:00:00 2001 From: Yug Date: Fri, 1 May 2026 12:42:42 +0530 Subject: [PATCH] resolve --- .gitignore | 208 +++++++++--------- .../proxy/_experimental/mcp_server/server.py | 41 ++-- .../mcp_test_config.yaml | 13 -- .../custom_mcp_server.py | 0 .../test_config_sampling_elicitation.yaml | 19 ++ tests/mcp_tests/test_mcp_server.py | 1 + .../test_sampling_elicitation_live.py} | 31 ++- 7 files changed, 154 insertions(+), 159 deletions(-) delete mode 100644 tests/mcp_sampling_elicitation/mcp_test_config.yaml rename tests/{mcp_sampling_elicitation => mcp_tests}/custom_mcp_server.py (100%) create mode 100644 tests/mcp_tests/test_configs/test_config_sampling_elicitation.yaml rename tests/{mcp_sampling_elicitation/test_live_mcp.py => mcp_tests/test_sampling_elicitation_live.py} (68%) diff --git a/.gitignore b/.gitignore index 38bf9554b5b..7bc84ae2bf4 100644 --- a/.gitignore +++ b/.gitignore @@ -1,103 +1,105 @@ -.python-version -.venv -.venv_policy_test -.env -.claude -.newenv -newenv/* -litellm/proxy/myenv/* -litellm_uuid.txt -__pycache__/ -*.pyc -bun.lockb -**/.DS_Store -.aider* -litellm_results.jsonl -secrets.toml -.gitignore -litellm/proxy/litellm_secrets.toml -litellm/proxy/api_log.json -.idea/ -router_config.yaml -litellm_server/config.yaml -litellm/proxy/_secret_config.yaml -.aws-sam/ -litellm/tests/aiologs.log -litellm/tests/exception_data.txt -litellm/tests/config_*.yaml -litellm/tests/langfuse.log -langfuse.log -.langfuse.log -litellm/tests/test_custom_logger.py -litellm/tests/langfuse.log -litellm/tests/dynamo*.log -.vscode/settings.json -litellm/proxy/log.txt -proxy_server_config_@.yaml -.gitignore -proxy_server_config_2.yaml -litellm/proxy/secret_managers/credentials.json -hosted_config.yaml -litellm/proxy/tests/node_modules -litellm/proxy/tests/package.json -litellm/proxy/tests/package-lock.json -ui/litellm-dashboard/.next -ui/litellm-dashboard/node_modules -ui/litellm-dashboard/next-env.d.ts -ui/litellm-dashboard/package.json -ui/litellm-dashboard/package-lock.json -deploy/charts/litellm/*.tgz -deploy/charts/litellm/charts/* -deploy/charts/*.tgz -litellm/proxy/vertex_key.json -**/.vim/ -**/node_modules -kub.yaml -loadtest_kub.yaml -litellm/proxy/_new_secret_config.yaml -litellm/proxy/_new_secret_config.yaml -litellm/proxy/_super_secret_config.yaml -litellm/proxy/_super_secret_config.yaml -litellm/proxy/myenv/bin/activate -litellm/proxy/myenv/bin/Activate.ps1 -myenv/* -litellm/tests/log.txt -litellm/tests/langfuse.log -litellm/tests/langfuse.log -litellm/proxy/google-cloud-sdk/* -tests/llm_translation/log.txt -venv/ -tests/local_testing/log.txt - -.codegpt -litellm/proxy/_new_new_secret_config.yaml -litellm/proxy/custom_guardrail.py -**/.mypy_cache/ -litellm/proxy/application.log -tests/llm_translation/vertex_test_account.json -tests/llm_translation/test_vertex_key.json -litellm/proxy/migrations/0_init/migration.sql -litellm/proxy/db/migrations/0_init/migration.sql -litellm/proxy/db/migrations/* -litellm/proxy/migrations/*config.yaml -litellm/proxy/migrations/* -litellm/proxy/to_delete_loadtest_work/* -config.yaml -tests/litellm/litellm_core_utils/llm_cost_calc/log.txt -tests/test_custom_dir/* -test.py - -litellm_config.yaml -!.github/observatory/litellm_config.yaml -.cursor -.vscode/launch.json -litellm/proxy/to_delete_loadtest_work/* -update_model_cost_map.py -tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py -scripts/test_vertex_ai_search.py -LAZY_LOADING_IMPROVEMENTS.md -STABILIZATION_TODO.md -**/test-results -**/playwright-report -**/*.storageState.json -**/coverage \ No newline at end of file +.python-version +.venv +.venv_policy_test +.env +.claude +.newenv +newenv/* +litellm/proxy/myenv/* +litellm_uuid.txt +__pycache__/ +*.pyc +bun.lockb +**/.DS_Store +.aider* +litellm_results.jsonl +secrets.toml +.gitignore +litellm/proxy/litellm_secrets.toml +litellm/proxy/api_log.json +.idea/ +router_config.yaml +litellm_server/config.yaml +litellm/proxy/_secret_config.yaml +.aws-sam/ +litellm/tests/aiologs.log +litellm/tests/exception_data.txt +litellm/tests/config_*.yaml +litellm/tests/langfuse.log +langfuse.log +.langfuse.log +litellm/tests/test_custom_logger.py +litellm/tests/langfuse.log +litellm/tests/dynamo*.log +.vscode/settings.json +litellm/proxy/log.txt +proxy_server_config_@.yaml +.gitignore +proxy_server_config_2.yaml +litellm/proxy/secret_managers/credentials.json +hosted_config.yaml +litellm/proxy/tests/node_modules +litellm/proxy/tests/package.json +litellm/proxy/tests/package-lock.json +ui/litellm-dashboard/.next +ui/litellm-dashboard/node_modules +ui/litellm-dashboard/next-env.d.ts +ui/litellm-dashboard/package.json +ui/litellm-dashboard/package-lock.json +deploy/charts/litellm/*.tgz +deploy/charts/litellm/charts/* +deploy/charts/*.tgz +litellm/proxy/vertex_key.json +**/.vim/ +**/node_modules +kub.yaml +loadtest_kub.yaml +litellm/proxy/_new_secret_config.yaml +litellm/proxy/_new_secret_config.yaml +litellm/proxy/_super_secret_config.yaml +litellm/proxy/_super_secret_config.yaml +litellm/proxy/myenv/bin/activate +litellm/proxy/myenv/bin/Activate.ps1 +myenv/* +litellm/tests/log.txt +litellm/tests/langfuse.log +litellm/tests/langfuse.log +litellm/proxy/google-cloud-sdk/* +tests/llm_translation/log.txt +venv/ +tests/local_testing/log.txt + +.codegpt +litellm/proxy/_new_new_secret_config.yaml +litellm/proxy/custom_guardrail.py +**/.mypy_cache/ +litellm/proxy/application.log +tests/llm_translation/vertex_test_account.json +tests/llm_translation/test_vertex_key.json +litellm/proxy/migrations/0_init/migration.sql +litellm/proxy/db/migrations/0_init/migration.sql +litellm/proxy/db/migrations/* +litellm/proxy/migrations/*config.yaml +litellm/proxy/migrations/* +litellm/proxy/to_delete_loadtest_work/* +config.yaml +tests/litellm/litellm_core_utils/llm_cost_calc/log.txt +tests/test_custom_dir/* +test.py + +litellm_config.yaml +!.github/observatory/litellm_config.yaml +.cursor +.vscode/launch.json +litellm/proxy/to_delete_loadtest_work/* +update_model_cost_map.py +tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py +scripts/test_vertex_ai_search.py +LAZY_LOADING_IMPROVEMENTS.md +STABILIZATION_TODO.md +**/test-results +**/playwright-report +**/*.storageState.json +**/coverage +tests/mcp_sampling_elicitation/mcp_test_config.yaml +custom_server.log \ No newline at end of file diff --git a/litellm/proxy/_experimental/mcp_server/server.py b/litellm/proxy/_experimental/mcp_server/server.py index 2bcd50ab5ca..5248ad4e3bb 100644 --- a/litellm/proxy/_experimental/mcp_server/server.py +++ b/litellm/proxy/_experimental/mcp_server/server.py @@ -1333,6 +1333,7 @@ if MCP_AVAILABLE: extra_headers=extra_headers, add_prefix=True, # Always add server prefix raw_headers=raw_headers, + user_api_key_auth=user_api_key_auth, ) filtered_tools = filter_tools_by_allowed_tools(tools, server) filtered_tools = await filter_tools_by_key_team_permissions( @@ -2840,15 +2841,20 @@ if MCP_AVAILABLE: except ValueError: pass - set_auth_context( - user_api_key_auth=user_api_key_auth, - mcp_auth_header=mcp_auth_header, - mcp_servers=mcp_servers, - mcp_server_auth_headers=mcp_server_auth_headers, - oauth2_headers=oauth2_headers, - raw_headers=raw_headers, - client_ip=_sse_client_ip, - ) + # Reuse existing session auth if available to ensure identity stability for cleanup + if session_id_param and 'session_auth' in locals() and session_auth is not None: + auth_context_var.set(session_auth) + else: + set_auth_context( + user_api_key_auth=user_api_key_auth, + mcp_auth_header=mcp_auth_header, + mcp_servers=mcp_servers, + mcp_server_auth_headers=mcp_server_auth_headers, + oauth2_headers=oauth2_headers, + raw_headers=raw_headers, + client_ip=_sse_client_ip, + ) + await sse.handle_post_message(scope, receive, send) except Exception as e: await _render_mcp_error( @@ -3051,26 +3057,9 @@ if MCP_AVAILABLE: if existing_auth is not None: _session_obj_auth_storage[id(session)] = existing_auth else: - # Shouldn't normally happen — _session_auth_storage is - # populated before server.run(). Log for diagnostics. - verbose_logger.debug( - "get_or_extract_auth_context: no existing auth in " - "_session_auth_storage; creating standalone entry " - "(cleanup may not remove it via identity check)" - ) auth = auth_context_var.get() if auth and isinstance(auth, MCPAuthenticatedUser): _session_obj_auth_storage[id(session)] = auth - else: - _session_obj_auth_storage[id(session)] = MCPAuthenticatedUser( - user_api_key_auth=user_api_key_auth, - mcp_auth_header=mcp_auth_header, - mcp_servers=mcp_servers, - mcp_server_auth_headers=mcp_server_auth_headers, - oauth2_headers=oauth2_headers, - raw_headers=raw_headers, - client_ip=_client_ip, - ) else: # Fallback: try session-object-identity lookup first (robust) stored: Optional[MCPAuthenticatedUser] = None diff --git a/tests/mcp_sampling_elicitation/mcp_test_config.yaml b/tests/mcp_sampling_elicitation/mcp_test_config.yaml deleted file mode 100644 index d94a472ccdd..00000000000 --- a/tests/mcp_sampling_elicitation/mcp_test_config.yaml +++ /dev/null @@ -1,13 +0,0 @@ -model_list: - - model_name: groq-model - litellm_params: - model: groq/llama-3.3-70b-versatile - api_key: os.environ/GROQ_API_KEY -litellm_settings: - default_mcp_sampling_model: groq/llama-3.3-70b-versatile -mcp_servers: - test_server: - transport: stdio - command: "python" - args: ["tests/mcp_sampling_elicitation/custom_mcp_server.py"] - allow_all_keys: true diff --git a/tests/mcp_sampling_elicitation/custom_mcp_server.py b/tests/mcp_tests/custom_mcp_server.py similarity index 100% rename from tests/mcp_sampling_elicitation/custom_mcp_server.py rename to tests/mcp_tests/custom_mcp_server.py diff --git a/tests/mcp_tests/test_configs/test_config_sampling_elicitation.yaml b/tests/mcp_tests/test_configs/test_config_sampling_elicitation.yaml new file mode 100644 index 00000000000..8c04f0a2af1 --- /dev/null +++ b/tests/mcp_tests/test_configs/test_config_sampling_elicitation.yaml @@ -0,0 +1,19 @@ +general_settings: + master_key: sk-1234 + +litellm_settings: + drop_params: true + +model_list: + - model_name: gpt-4o-mini + litellm_params: + model: gpt-4o-mini + api_key: os.environ/OPENAI_API_KEY + +mcp_servers: + test_server: + transport: stdio + command: python + args: + - tests/mcp_tests/custom_mcp_server.py + allow_all_keys: true diff --git a/tests/mcp_tests/test_mcp_server.py b/tests/mcp_tests/test_mcp_server.py index 7b2730efabe..f93300ee533 100644 --- a/tests/mcp_tests/test_mcp_server.py +++ b/tests/mcp_tests/test_mcp_server.py @@ -938,6 +938,7 @@ async def test_get_tools_from_mcp_servers(): # noqa: PLR0915 extra_headers=None, add_prefix=False, raw_headers=None, + user_api_key_auth=None, ): if server.server_id == "server1_id": return [mock_tool_1] diff --git a/tests/mcp_sampling_elicitation/test_live_mcp.py b/tests/mcp_tests/test_sampling_elicitation_live.py similarity index 68% rename from tests/mcp_sampling_elicitation/test_live_mcp.py rename to tests/mcp_tests/test_sampling_elicitation_live.py index 71e8e358972..bd8e402b818 100644 --- a/tests/mcp_sampling_elicitation/test_live_mcp.py +++ b/tests/mcp_tests/test_sampling_elicitation_live.py @@ -1,6 +1,7 @@ import asyncio import os import logging +import pytest from mcp.client.sse import sse_client from mcp.client.session import ClientSession from mcp.types import ElicitResult @@ -9,10 +10,12 @@ logging.basicConfig(level=logging.INFO) logger = logging.getLogger(__name__) -async def main(): - if os.environ.get("RUN_LIVE_MCP_TEST") != "1": - logger.info("Skipping live integration test. Set RUN_LIVE_MCP_TEST=1 to run.") - return +@pytest.mark.asyncio +@pytest.mark.skipif( + os.environ.get("RUN_LIVE_MCP_TEST") != "1", + reason="Skipping live integration test. Set RUN_LIVE_MCP_TEST=1 to run.", +) +async def test_live_mcp_pipeline(): logger.info("Connecting to LiteLLM Proxy via SSE...") async def my_elicitation_callback(context, params): @@ -43,16 +46,10 @@ async def main(): logger.info("\n--- Testing Complex Pipeline (Elicitation + Sampling) ---") logger.info("Calling 'test_server-test_complex_pipeline'...") - try: - result = await session.call_tool( - "test_server-test_complex_pipeline", arguments={} - ) - logger.info("\nFINAL TOOL RESULT:") - logger.info("==================") - logger.info(result.content[0].text) - except Exception as e: - logger.info(f"Error calling test_complex_pipeline: {e}") - - -if __name__ == "__main__": - asyncio.run(main()) + result = await session.call_tool( + "test_server-test_complex_pipeline", arguments={} + ) + logger.info("\nFINAL TOOL RESULT:") + logger.info("==================") + logger.info(result.content[0].text) + assert "Pipeline Complete!" in result.content[0].text