From 322483525d8bd81de10db465377a65d66a550b9e Mon Sep 17 00:00:00 2001 From: mubashir1osmani Date: Tue, 11 Aug 2026 15:00:14 -0700 Subject: [PATCH] test(e2e): pin MCP toolset naming, key-grant survival, pricing-override clear, and responses-mode contracts Covers the writable rows of the customer bug audit: - MCP toolset rows store bare tool names and serve the prefixed wire name to toolset-scoped keys (LIT-3419, regressed twice) - /key/update with only max_budget preserves the key's MCP grants (PR #34452 regression) and stale deleted-server ids are dropped instead of 403ing the edit (LIT-3278) - an explicit-null PATCH clears a deployment pricing override and Reload Price Data does not resurrect it (ticket #6844) - a /model/new deployment declaring model_info.mode=responses bridges /chat/completions to the Responses path with no static cost-map entry, and /health/test_connection mode=responses completes without the duplicate acompletion TypeError - /v1/responses streaming announces every text item before its deltas on multi-step tool conversations (#27671) --- .../coverage_registry/llm_conversational.yaml | 2 + tests/e2e/coverage_registry/mcp.yaml | 32 ++++ tests/e2e/coverage_registry/mgmt.yaml | 2 + tests/e2e/coverage_registry/other.yaml | 1 + .../test_responses_routing_e2e.py | 136 +++++++++++++++++ .../test_responses_stream_text_start_e2e.py | 140 +++++++++++++++++ .../test_model_pricing_override_e2e.py | 144 ++++++++++++++++++ tests/e2e/mcp/test_mcp_key_permissions_e2e.py | 115 ++++++++++++++ tests/e2e/mcp/test_mcp_toolset_e2e.py | 137 +++++++++++++++++ tests/e2e/models.py | 8 +- 10 files changed, 715 insertions(+), 2 deletions(-) create mode 100644 tests/e2e/llm_translation/test_responses_routing_e2e.py create mode 100644 tests/e2e/llm_translation/test_responses_stream_text_start_e2e.py create mode 100644 tests/e2e/management/test_model_pricing_override_e2e.py create mode 100644 tests/e2e/mcp/test_mcp_key_permissions_e2e.py create mode 100644 tests/e2e/mcp/test_mcp_toolset_e2e.py diff --git a/tests/e2e/coverage_registry/llm_conversational.yaml b/tests/e2e/coverage_registry/llm_conversational.yaml index 8163866abd1..13358b77165 100644 --- a/tests/e2e/coverage_registry/llm_conversational.yaml +++ b/tests/e2e/coverage_registry/llm_conversational.yaml @@ -69,3 +69,5 @@ - {id: llm.responses.vertex.tool_use.nonstream.works, module: llm, tier: P1, subject_endpoint: responses, route: vertex, capability: tool_use, streaming: nonstream, assertions: [works], source: "model_prices json", rationale: "Responses tool calls w/ Vertex"} - {id: llm.responses.azure_openai.basic.nonstream.works, module: llm, tier: P1, subject_endpoint: responses, route: azure_openai, capability: basic, streaming: nonstream, assertions: [works], source: "response_api_endpoints/endpoints.py:26", rationale: "Responses w/ Azure OpenAI (smoke)"} - {id: llm.responses.azure_openai.tool_use.nonstream.works, module: llm, tier: P1, subject_endpoint: responses, route: azure_openai, capability: tool_use, streaming: nonstream, assertions: [works], source: "model_prices json", rationale: "Responses tool calls w/ Azure OpenAI"} +- {id: llm.chat_completions.gemini.basic.nonstream.bridges_to_responses, module: llm, tier: P0, subject_endpoint: chat_completions, route: gemini, capability: basic, streaming: nonstream, assertions: [works], source: "main.py:5434 responses_api_bridge_check + router.py:8493 _register_deployment_in_model_cost", rationale: "A /model/new deployment declaring model_info.mode=responses must bridge /chat/completions to the Responses path even with no static cost-map entry; the 1.92 cost-map routing move broke every responses-only model whose entry had not shipped ('does not support the /v1/chat/completions API')"} +- {id: llm.responses.gemini.tool_use.stream.text_start_precedes_deltas, module: llm, tier: P0, subject_endpoint: responses, route: gemini, capability: tool_use, streaming: stream, assertions: [works], source: "responses/litellm_completion_transformation/streaming_iterator.py:779 content_part.added before deltas", rationale: "Post-tool-round text streamed output_text.delta events whose item was never announced, so OpenAI-SDK parsers died with 'text part ... not found' (#27671); every delta must follow an output_item.added/content_part.added carrying its item id"} diff --git a/tests/e2e/coverage_registry/mcp.yaml b/tests/e2e/coverage_registry/mcp.yaml index ab644118a47..1809c03d8be 100644 --- a/tests/e2e/coverage_registry/mcp.yaml +++ b/tests/e2e/coverage_registry/mcp.yaml @@ -119,3 +119,35 @@ assertions: [succeeds] source: "server.py:1089" rationale: Smoke; rarely used; same auth model as tools +- id: mcp.toolset.api_key.persists_bare_tool_names + module: mcp + tier: P0 + operation: toolset + auth_family: api_key + assertions: [persists_bare_tool_names] + source: "toolset_db.py:24" + rationale: "Toolset rows store bare {server_id, tool_name} pairs; a pre-prefixed name double-prefixes at serve time and tools stop matching (LIT-3419, regressed twice)" +- id: mcp.toolset.api_key.serves_prefixed_tool_names + module: mcp + tier: P0 + operation: toolset + auth_family: api_key + assertions: [serves_prefixed_tool_names] + source: "utils.py:249 add_server_prefix_to_name" + rationale: "A toolset-scoped key must be served the documented {server_prefix}{separator}{tool} wire name, or clients calling the documented name get tool-not-found (LIT-3419)" +- id: mcp.key_update.api_key.budget_change_preserves_grants + module: mcp + tier: P0 + operation: key_update + auth_family: api_key + assertions: [budget_change_preserves_grants] + source: "key_management_endpoints.py:2047 object_permission gating" + rationale: "A max_budget-only /key/update silently dropped the key's MCP toolset/servers (fixed PR #34452); partial updates must not null unrelated fields" +- id: mcp.key_update.api_key.stale_server_ids_dropped + module: mcp + tier: P1 + operation: key_update + auth_family: api_key + assertions: [stale_server_ids_dropped] + source: "object_permission_utils.py:570 stale_identifiers" + rationale: "A key holding deleted MCP server ids was un-editable via the UI (403 'not allowed by team'), forcing hand-edited DB rows (LIT-3278); stale ids must be dropped, not fatal" diff --git a/tests/e2e/coverage_registry/mgmt.yaml b/tests/e2e/coverage_registry/mgmt.yaml index 2a0fc5c9f29..4432c90efe0 100644 --- a/tests/e2e/coverage_registry/mgmt.yaml +++ b/tests/e2e/coverage_registry/mgmt.yaml @@ -71,3 +71,5 @@ - {id: mgmt.workflow.list.happy_path, module: mgmt, tier: P2, surface: api, assertions: [happy_path], source: "workflow_management_endpoints.py", rationale: "Workflow tracking (smoke)"} - {id: mgmt.credential_migration.check.happy_path, module: mgmt, tier: P2, surface: api, assertions: [happy_path], source: "key_management_endpoints.py:4252", rationale: "Encryption migration (smoke)"} - {id: mgmt.credential.new.serves_request, module: mgmt, tier: P1, surface: api, assertions: [serves_request], source: "credential_endpoints/endpoints.py:42", rationale: "Stored credential resolves into a deployment and serves a live /messages request"} +- {id: mgmt.model.update.clears_pricing_override, module: mgmt, tier: P0, surface: api, assertions: [clears_pricing_override], source: "model_management_endpoints.py:391 explicit-null clears", rationale: "UI-pinned deployment prices could never be cleared, permanently shadowing the cost map (ticket #6844); an explicit-null PATCH must drop the override"} +- {id: mgmt.model.cost_map_reload.keeps_cleared_override, module: mgmt, tier: P1, surface: api, assertions: [keeps_cleared_override], source: "proxy_server.py:16397 /reload/model_cost_map", rationale: "Reload Price Data must refresh cost-map rates without resurrecting a cleared deployment override (ticket #6844)"} diff --git a/tests/e2e/coverage_registry/other.yaml b/tests/e2e/coverage_registry/other.yaml index ace4f8bcdc9..8a51d37c99a 100644 --- a/tests/e2e/coverage_registry/other.yaml +++ b/tests/e2e/coverage_registry/other.yaml @@ -37,3 +37,4 @@ - {id: other.a2a.message_send.bridge_invokes, module: other, tier: P1, area: a2a, assertions: [bridge_invokes], source: "a2a_protocol/litellm_completion_bridge/handler.py", rationale: "A2A message/send routes through the completion bridge to a real provider and logs an asend_message spend row"} - {id: other.a2a.version.serves_pinned_0_3, module: other, tier: P1, area: a2a, assertions: [serves_pinned_0_3], source: "agent_endpoints/a2a_endpoints.py _served_version", rationale: "An agent pinning 0.3 returns the flat 0.3 message shape (parts on the result)"} - {id: other.a2a.version.serves_pinned_1_0, module: other, tier: P1, area: a2a, assertions: [serves_pinned_1_0], source: "agent_endpoints/a2a_endpoints.py _served_version", rationale: "An agent pinning 1.0 returns the nested 1.0 message shape (result.message with ROLE_AGENT)"} +- {id: other.health.test_connection.responses_mode_succeeds, module: other, tier: P0, area: health, assertions: [responses_mode_succeeds], source: "main.py:627 kwargs.pop('acompletion') + health_check_helpers.py responses handler", rationale: "The UI Test Connection button on a responses-mode model crashed with functools.partial got multiple values for 'acompletion' (fixed v1.91.0-rc.1); the responses health handler must complete"} diff --git a/tests/e2e/llm_translation/test_responses_routing_e2e.py b/tests/e2e/llm_translation/test_responses_routing_e2e.py new file mode 100644 index 00000000000..bb2f4bc6203 --- /dev/null +++ b/tests/e2e/llm_translation/test_responses_routing_e2e.py @@ -0,0 +1,136 @@ +"""Live e2e: deployment-level mode=responses drives routing and health checks. + +Two regressions around the Responses API bridge, both hit through models +registered via /model/new rather than shipped in the static cost map: + +- routing: /chat/completions on a responses-only deployment used to work by + model-name matching, then 1.92 moved the decision to a cost-map lookup and + every model whose entry had not shipped yet started failing with "does not + support the /v1/chat/completions API". The router now injects the deployment's + model_info (including mode) into the cost map at registration, so a deployment + declaring mode=responses must bridge regardless of what the static map says. + The spend log's call_type distinguishes the two paths: "responses" only when + the bridge translated the call +- health: POST /health/test_connection with mode=responses used to crash with + "functools.partial got multiple values for 'acompletion'" (the bridge handed + litellm.acompletion a second acompletion kwarg). The UI's Test Connection + button on any responses-mode model hit this as a 500 +""" + +from __future__ import annotations + +import pytest +from pydantic import BaseModel + +from e2e_config import unique_marker +from e2e_http import unwrap +from lifecycle import ResourceManager +from models import ChatBody, ChatMessage, LiteLLMParamsBody +from proxy_client import ProxyClient + +pytestmark = pytest.mark.e2e + +BACKEND_MODEL = "gemini/gemini-2.5-flash" +GEMINI_API_KEY = "os.environ/GEMINI_API_KEY" + + +class HealthCheckParams(BaseModel): + model: str + + +class HealthCheckModelInfo(BaseModel): + id: str + + +class HealthCheckBody(BaseModel): + litellm_params: HealthCheckParams + model_info: HealthCheckModelInfo + mode: str + + +class HealthCheckResult(BaseModel): + error: str | None = None + + +class HealthCheckResponse(BaseModel): + status: str + result: HealthCheckResult = HealthCheckResult() + + +def _register_responses_mode_model(proxy: ProxyClient, resources: ResourceManager) -> tuple[str, str]: + """Register a gemini deployment declaring mode=responses in its model_info, + the shape a responses-only model (bedrock_mantle gpt-5.x, openai codex) + takes when added through the UI. Returns (model_name, model_id).""" + model_name = f"e2e-responses-mode-{unique_marker()}" + model_id = proxy.create_model(model_name, LiteLLMParamsBody(model=BACKEND_MODEL, api_key=GEMINI_API_KEY)) + resources.defer(lambda: proxy.delete_model(model_id)) + return model_name, model_id + + +class TestResponsesModeRouting: + @pytest.mark.covers( + "llm.chat_completions.gemini.basic.nonstream.bridges_to_responses", + exercised_on=["chat_completions"], + ) + def test_chat_completion_bridges_when_deployment_declares_responses_mode( + self, proxy: ProxyClient, resources: ResourceManager, scoped_key: str + ) -> None: + model_name = f"e2e-responses-mode-{unique_marker()}" + model_id = proxy.create_model( + model_name, + LiteLLMParamsBody(model=BACKEND_MODEL, api_key=GEMINI_API_KEY), + mode="responses", + ) + resources.defer(lambda: proxy.delete_model(model_id)) + + chat = unwrap( + proxy.chat( + scoped_key, + ChatBody( + model=model_name, + messages=[ChatMessage(role="user", content=f"Reply with only the word ok. {unique_marker()}")], + max_tokens=256, + ), + ) + ) + assert chat.id and chat.id.startswith("chatcmpl-"), ( + f"the bridged call must still answer in chat-completion shape: id={chat.id!r}" + ) + assert chat.choices and chat.choices[0].message is not None, ( + f"bridged call returned no choices: {chat}" + ) + + rows = proxy.poll_logs_for_key(scoped_key) + call_types = {row.call_type for row in rows} + assert "responses" in call_types, ( + f"a deployment with model_info.mode=responses must be served through the " + f"Responses bridge (spend log call_type='responses'); the call went down the " + f"chat path instead: {call_types}" + ) + + @pytest.mark.covers("other.health.test_connection.responses_mode_succeeds") + def test_health_test_connection_in_responses_mode_succeeds( + self, proxy: ProxyClient, resources: ResourceManager + ) -> None: + _, model_id = _register_responses_mode_model(proxy, resources) + + # The UI's Test Connection button sends the provider model plus the + # deployment id; mode=responses selects the aresponses health handler. + health = unwrap( + proxy.transport.post( + "/health/test_connection", + headers=proxy.transport.master, + json=HealthCheckBody( + litellm_params=HealthCheckParams(model=BACKEND_MODEL), + model_info=HealthCheckModelInfo(id=model_id), + mode="responses", + ), + response_type=HealthCheckResponse, + ) + ) + error = health.result.error or "" + assert "acompletion" not in error and "functools.partial" not in error, ( + f"health check crashed on the duplicate-acompletion TypeError the responses " + f"bridge used to raise: {error}" + ) + assert health.status == "success", f"responses-mode health check failed: {error}" diff --git a/tests/e2e/llm_translation/test_responses_stream_text_start_e2e.py b/tests/e2e/llm_translation/test_responses_stream_text_start_e2e.py new file mode 100644 index 00000000000..9387eb1c845 --- /dev/null +++ b/tests/e2e/llm_translation/test_responses_stream_text_start_e2e.py @@ -0,0 +1,140 @@ +"""Live e2e: /v1/responses streaming emits a text-start before every text delta. + +github issue #27671: on multi-step tool conversations, the text the model streams +after a tool round arrived as `response.output_text.delta` events whose item had +never been announced - no `response.output_item.added` / `response.content_part.added` +carried that item id first - so OpenAI-SDK parsers died with "text part +chatcmpl- not found". The customer hit it through OpenCode against bedrock +anthropic; the defect lives in the shared completion-to-responses streaming +iterator every non-native-responses provider goes through, which is what this +test drives (gemini here, same translation layer). + +The request replays a completed tool round (function_call + function_call_output +in the input) so the streamed text is exactly the post-tool "step 2" text the +issue describes. +""" + +from __future__ import annotations + +from typing import Literal + +import pytest +from pydantic import BaseModel + +from e2e_config import unique_marker +from e2e_http import StreamingResponse +from endpoints_client import FunctionParameterProperty, FunctionParameters, ResponsesFunctionTool +from lifecycle import ResourceManager +from models import LiteLLMParamsBody +from proxy_client import ProxyClient + +pytestmark = pytest.mark.e2e + +BACKEND_MODEL = "gemini/gemini-2.5-flash" + +_WEATHER_TOOL = ResponsesFunctionTool( + name="get_weather", + description="Get the current weather for a location", + parameters=FunctionParameters( + properties={"location": FunctionParameterProperty(type="string")}, + required=["location"], + ), +) + + +class _UserMessage(BaseModel): + role: Literal["user"] = "user" + content: str + + +class _FunctionCall(BaseModel): + type: Literal["function_call"] = "function_call" + call_id: str + name: str + arguments: str + + +class _FunctionCallOutput(BaseModel): + type: Literal["function_call_output"] = "function_call_output" + call_id: str + output: str + + +class _ToolRoundRequest(BaseModel): + model: str + stream: bool = True + tools: list[ResponsesFunctionTool] + input: list[_UserMessage | _FunctionCall | _FunctionCallOutput] + + +class _EventItem(BaseModel): + id: str | None = None + + +class _StreamEvent(BaseModel): + type: str + item_id: str | None = None + item: _EventItem | None = None + + def announced_item_id(self) -> str | None: + """The item id this event announces or targets, wherever the event + carries it (delta events use item_id; added/done events nest it).""" + if self.item_id is not None: + return self.item_id + return self.item.id if self.item else None + + +def _stream_post_tool_text(proxy: ProxyClient, key: str, model: str) -> StreamingResponse: + call_id = f"call_{unique_marker()}" + return proxy.transport.stream( + "/v1/responses", + headers=proxy.transport.bearer(key), + json=_ToolRoundRequest( + model=model, + tools=[_WEATHER_TOOL], + input=[ + _UserMessage(content="What is the weather in Paris?"), + _FunctionCall(call_id=call_id, name="get_weather", arguments='{"location": "Paris"}'), + _FunctionCallOutput(call_id=call_id, output="22C and sunny"), + ], + ), + ) + + +class TestResponsesStreamTextStart: + @pytest.mark.covers( + "llm.responses.gemini.tool_use.stream.text_start_precedes_deltas", + exercised_on=["responses"], + ) + def test_post_tool_text_deltas_follow_a_text_start_with_same_item_id( + self, proxy: ProxyClient, resources: ResourceManager, scoped_key: str + ) -> None: + model = f"e2e-text-start-{unique_marker()}" + model_id = proxy.create_model( + model, LiteLLMParamsBody(model=BACKEND_MODEL, api_key="os.environ/GEMINI_API_KEY") + ) + resources.defer(lambda: proxy.delete_model(model_id)) + + result = _stream_post_tool_text(proxy, scoped_key, model) + assert result.ok and result.is_streaming, f"responses stream was not established: {result}" + assert result.stream_error is None, f"stream carried an error event: {result.stream_error}" + + events = [_StreamEvent.model_validate_json(payload) for payload in result.stream_events] + deltas = [event for event in events if event.type == "response.output_text.delta"] + assert deltas, f"the post-tool round streamed no text deltas: {[e.type for e in events]}" + + announced: set[str] = set() + for event in events: + item_id = event.announced_item_id() + if event.type in ("response.output_item.added", "response.content_part.added") and item_id: + announced.add(item_id) + if event.type == "response.output_text.delta": + assert item_id in announced, ( + f"text delta for item {item_id!r} arrived before any output_item.added/" + f"content_part.added announced it; OpenAI-SDK parsers fail with " + f"'text part ... not found' (#27671). Events so far: " + f"{[e.type for e in events[: events.index(event) + 1]]}" + ) + assert any(event.type == "response.completed" for event in events), ( + f"stream never completed: {[e.type for e in events][-5:]}" + ) diff --git a/tests/e2e/management/test_model_pricing_override_e2e.py b/tests/e2e/management/test_model_pricing_override_e2e.py new file mode 100644 index 00000000000..5cb9d036d38 --- /dev/null +++ b/tests/e2e/management/test_model_pricing_override_e2e.py @@ -0,0 +1,144 @@ +"""Live e2e: a deployment's custom pricing override can be cleared, and Reload +Price Data restores cost-map rates (ticket #6844). + +Models created through the UI carry the displayed price as a deployment-level +override, which then shadows the shared cost map forever: Reload Price Data +becomes a no-op for that deployment and the admin cannot get back to canonical +pricing. The escape hatch is PATCH /model/{id}/update with the pricing field set +to an explicit null, which drops the override from both stored blobs +(SPECIAL_MODEL_INFO_PARAMS clears in update_db_model). These tests pin that +escape hatch and the reload endpoint the UI's Reload Price Data button calls. + +Known gap called out in the PR: the older POST /model/update route ignores +explicit nulls (its merge keeps the stored value for any null field), so only +the PATCH route can clear an override. +""" + +from __future__ import annotations + +import time + +import pytest +from pydantic import BaseModel + +from e2e_config import unique_marker +from e2e_http import NoBody, unwrap +from lifecycle import ResourceManager +from models import LiteLLMParamsBody, ModelInfoEntry +from proxy_client import ProxyClient + +pytestmark = pytest.mark.e2e + +BACKEND_MODEL = "gemini/gemini-2.5-flash" +# Deliberately ~100x the canonical gemini-2.5-flash input rate (3e-7) so an +# override that fails to clear is unmistakable from the cost-map value. +OVERRIDE_INPUT_RATE = 5e-05 + + +class PricingPatchBody(BaseModel): + """PATCH /model/{id}/update body carrying an explicit-null pricing clear. + + The pricing fields must serialize as literal nulls (the handler only clears + fields present in the payload), so they cannot ride LiteLLMParamsBody, whose + None fields are dropped by the transport's exclude_none serialization. + """ + + litellm_params: dict[str, float | None] + + +def _register_overridden_model(proxy: ProxyClient, resources: ResourceManager) -> tuple[str, str]: + model_name = f"e2e-pricing-override-{unique_marker()}" + model_id = proxy.create_model( + model_name, + LiteLLMParamsBody( + model=BACKEND_MODEL, + api_key="os.environ/GEMINI_API_KEY", + input_cost_per_token=OVERRIDE_INPUT_RATE, + ), + ) + resources.defer(lambda: proxy.delete_model(model_id)) + return model_name, model_id + + +def _model_entry(proxy: ProxyClient, model_name: str) -> ModelInfoEntry: + for entry in proxy.model_info(): + if entry.model_name == model_name: + return entry + pytest.fail(f"{model_name} absent from /model/info") + + +def _poll_override_cleared(proxy: ProxyClient, model_name: str) -> ModelInfoEntry: + """Poll /model/info until the deployment's input override reads as cleared + (the router picks the change up on its periodic DB reload).""" + deadline = time.monotonic() + proxy.poll_timeout + while True: + entry = _model_entry(proxy, model_name) + if entry.litellm_params.input_cost_per_token is None: + return entry + if time.monotonic() >= deadline: + return entry + time.sleep(proxy.poll_interval) + + +class TestPricingOverrideClearAndReload: + @pytest.mark.covers("mgmt.model.update.clears_pricing_override") + def test_explicit_null_patch_clears_pricing_override( + self, proxy: ProxyClient, resources: ResourceManager + ) -> None: + model_name, model_id = _register_overridden_model(proxy, resources) + assert _model_entry(proxy, model_name).litellm_params.input_cost_per_token == OVERRIDE_INPUT_RATE + + unwrap( + proxy.transport.patch( + f"/model/{model_id}/update", + headers=proxy.transport.master, + json=PricingPatchBody(litellm_params={"input_cost_per_token": None}), + response_type=NoBody, + ) + ) + + entry = _poll_override_cleared(proxy, model_name) + assert entry.litellm_params.input_cost_per_token is None, ( + f"PATCH with an explicit null must drop the stored pricing override; the " + f"deployment still bills at {entry.litellm_params.input_cost_per_token} " + f"(ticket #6844: stale UI-pinned prices cannot be cleared)" + ) + resolved_rate = entry.model_info.input_cost_per_token + assert resolved_rate is not None and resolved_rate < OVERRIDE_INPUT_RATE, ( + f"after the clear the deployment must fall back to the cost-map rate, " + f"not keep the {OVERRIDE_INPUT_RATE} override: resolved {resolved_rate}" + ) + + @pytest.mark.covers("mgmt.model.cost_map_reload.keeps_cleared_override") + def test_reload_price_data_does_not_overwrite_cleared_deployment( + self, proxy: ProxyClient, resources: ResourceManager + ) -> None: + model_name, model_id = _register_overridden_model(proxy, resources) + unwrap( + proxy.transport.patch( + f"/model/{model_id}/update", + headers=proxy.transport.master, + json=PricingPatchBody(litellm_params={"input_cost_per_token": None}), + response_type=NoBody, + ) + ) + _poll_override_cleared(proxy, model_name) + + unwrap( + proxy.transport.post( + "/reload/model_cost_map", + headers=proxy.transport.master, + json=NoBody(), + response_type=NoBody, + ) + ) + + entry = _poll_override_cleared(proxy, model_name) + assert entry.litellm_params.input_cost_per_token is None, ( + f"Reload Price Data resurrected the cleared deployment override: " + f"{entry.litellm_params.input_cost_per_token}" + ) + resolved_rate = entry.model_info.input_cost_per_token + assert resolved_rate is not None and resolved_rate < OVERRIDE_INPUT_RATE, ( + f"after the reload the deployment must price at the cost-map rate: {resolved_rate}" + ) diff --git a/tests/e2e/mcp/test_mcp_key_permissions_e2e.py b/tests/e2e/mcp/test_mcp_key_permissions_e2e.py new file mode 100644 index 00000000000..ee3e9114491 --- /dev/null +++ b/tests/e2e/mcp/test_mcp_key_permissions_e2e.py @@ -0,0 +1,115 @@ +"""Live e2e: a virtual key's MCP grants survive unrelated /key/update writes. + +Two regressions customers hit through the key-edit UI, both driven here exactly as +the UI drives them (POST /key/update with a partial body): + +- a budget-only update silently dropped the key's attached MCP servers (fixed in + PR #34452): /key/update must leave object_permission untouched when the body + does not carry one +- a key holding ids of since-deleted MCP servers could not be edited at all: the + update was rejected with "Key requests MCP servers not allowed by team", forcing + admins to fix DB rows by hand (LIT-3278). Stale ids must be dropped, not fatal +""" + +from __future__ import annotations + +import pytest + +from datadog_mcp import register_datadog_mcp +from e2e_config import unique_marker +from e2e_http import NoBody, unwrap +from lifecycle import ResourceManager +from mcp_client import McpClient +from models import KeyGenerateBody, KeyInfoParams, KeyInfoResponse, KeyUpdateBody, ObjectPermission + +pytestmark = pytest.mark.e2e + + +def _mcp_key(client: McpClient, resources: ResourceManager, server_ids: list[str]) -> str: + key = client.proxy.generate_key( + KeyGenerateBody( + user_id=f"e2e-mcp-key-perms-{unique_marker()}", + object_permission=ObjectPermission(mcp_servers=server_ids), + ) + ) + resources.defer(lambda: client.proxy.delete_key(key)) + return key + + +def _key_mcp_servers(client: McpClient, key: str) -> list[str]: + info = unwrap( + client.proxy.transport.get( + "/key/info", + headers=client.proxy.transport.master, + params=KeyInfoParams(key=key), + response_type=KeyInfoResponse, + ) + ).info + if info.object_permission is None: + return [] + return info.object_permission.mcp_servers or [] + + +class TestKeyUpdatePreservesMcpGrants: + @pytest.mark.covers("mcp.key_update.api_key.budget_change_preserves_grants") + def test_budget_only_update_preserves_mcp_servers( + self, client: McpClient, resources: ResourceManager + ) -> None: + server_id = register_datadog_mcp(client, resources) + client.await_registered(server_id) + key = _mcp_key(client, resources, [server_id]) + assert _key_mcp_servers(client, key) == [server_id] + + unwrap( + client.proxy.transport.post( + "/key/update", + headers=client.proxy.transport.master, + json=KeyUpdateBody(key=key, max_budget=42.0), + response_type=NoBody, + ) + ) + + info = unwrap( + client.proxy.transport.get( + "/key/info", + headers=client.proxy.transport.master, + params=KeyInfoParams(key=key), + response_type=KeyInfoResponse, + ) + ).info + assert info.max_budget == 42.0, f"the budget update itself did not land: {info}" + servers = (info.object_permission.mcp_servers or []) if info.object_permission else [] + assert servers == [server_id], ( + f"a max_budget-only /key/update dropped the key's MCP servers " + f"(PR #34452 regression): {info.object_permission}" + ) + + @pytest.mark.covers("mcp.key_update.api_key.stale_server_ids_dropped") + def test_update_with_deleted_server_ids_succeeds_and_drops_them( + self, client: McpClient, resources: ResourceManager + ) -> None: + surviving = register_datadog_mcp(client, resources) + doomed = register_datadog_mcp(client, resources) + client.await_registered(surviving) + client.await_registered(doomed) + key = _mcp_key(client, resources, [surviving, doomed]) + + client.delete_server(doomed) + + # The UI resends the key's full server list, deleted ids included. + unwrap( + client.proxy.transport.post( + "/key/update", + headers=client.proxy.transport.master, + json=KeyUpdateBody( + key=key, + object_permission=ObjectPermission(mcp_servers=[surviving, doomed]), + ), + response_type=NoBody, + ) + ) + + assert _key_mcp_servers(client, key) == [surviving], ( + f"/key/update must drop the deleted server id {doomed} and keep {surviving}; " + f"rejecting or persisting stale ids locks the key out of UI edits (LIT-3278)" + ) diff --git a/tests/e2e/mcp/test_mcp_toolset_e2e.py b/tests/e2e/mcp/test_mcp_toolset_e2e.py new file mode 100644 index 00000000000..4e6247064a4 --- /dev/null +++ b/tests/e2e/mcp/test_mcp_toolset_e2e.py @@ -0,0 +1,137 @@ +"""Live e2e: MCP toolset persistence and the tool-name prefix contract (LIT-3419). + +A toolset row stores bare `{server_id, tool_name}` pairs; the gateway applies the +`{server_prefix}{separator}{tool}` wire prefix only at serve time. A regression in +either direction broke customers twice: the UI once persisted pre-prefixed names +(so serve-time prefixing double-prefixed them and tools/list stopped matching), and +once the served name lost the prefix entirely (so clients calling the documented +`{server}-{tool}` name got "tool not found"). These tests pin both halves against +the real Datadog remote MCP server: + +- the stored toolset row keeps exactly the bare tool name it was created with +- /v1/mcp/tools serves that same tool under the prefixed wire name for a key + scoped to the toolset via object_permission.mcp_toolsets +""" + +from __future__ import annotations + +import pytest +from pydantic import BaseModel + +from datadog_mcp import SEARCH_LOGS_TOOL, register_datadog_mcp +from e2e_config import unique_marker +from e2e_http import NoBody, unwrap +from lifecycle import ResourceManager +from mcp_client import McpClient +from models import KeyGenerateBody, ObjectPermission + +pytestmark = pytest.mark.e2e + + +class ToolsetTool(BaseModel): + server_id: str + tool_name: str + + +class ToolsetCreateBody(BaseModel): + toolset_name: str + description: str | None = None + tools: list[ToolsetTool] + + +class Toolset(BaseModel): + toolset_id: str + toolset_name: str + tools: list[ToolsetTool] = [] + + +class PrefixedTool(BaseModel): + name: str + + +class PrefixedToolsResponse(BaseModel): + tools: list[PrefixedTool] = [] + + +def _delete_toolset(client: McpClient, toolset_id: str) -> None: + _ = client.proxy.transport.delete( + f"/v1/mcp/toolset/{toolset_id}", + headers=client.proxy.transport.master, + json=NoBody(), + response_type=NoBody, + ) + + +def _create_toolset(client: McpClient, resources: ResourceManager, *, server_id: str) -> Toolset: + toolset = unwrap( + client.proxy.transport.post( + "/v1/mcp/toolset", + headers=client.proxy.transport.master, + json=ToolsetCreateBody( + toolset_name=f"e2e-toolset-{unique_marker()}", + tools=[ToolsetTool(server_id=server_id, tool_name=SEARCH_LOGS_TOOL)], + ), + response_type=Toolset, + ) + ) + resources.defer(lambda: _delete_toolset(client, toolset.toolset_id)) + return toolset + + +class TestMcpToolsetToolNames: + @pytest.mark.covers("mcp.toolset.api_key.persists_bare_tool_names") + def test_stored_toolset_keeps_bare_tool_name( + self, client: McpClient, resources: ResourceManager + ) -> None: + server_id = register_datadog_mcp(client, resources) + client.await_registered(server_id) + + created = _create_toolset(client, resources, server_id=server_id) + + stored = unwrap( + client.proxy.transport.get( + f"/v1/mcp/toolset/{created.toolset_id}", + headers=client.proxy.transport.master, + params=NoBody(), + response_type=Toolset, + ) + ) + assert stored.tools == [ToolsetTool(server_id=server_id, tool_name=SEARCH_LOGS_TOOL)], ( + f"toolset row must store the bare tool name paired with its server_id; a " + f"pre-prefixed or rewritten name breaks serve-time prefixing: {stored.tools}" + ) + + @pytest.mark.covers("mcp.toolset.api_key.serves_prefixed_tool_names") + def test_toolset_scoped_key_sees_prefixed_tool_name( + self, client: McpClient, resources: ResourceManager + ) -> None: + server_id = register_datadog_mcp(client, resources) + client.await_registered(server_id) + toolset = _create_toolset(client, resources, server_id=server_id) + + key = client.proxy.generate_key( + KeyGenerateBody( + user_id=f"e2e-toolset-{unique_marker()}", + object_permission=ObjectPermission(mcp_toolsets=[toolset.toolset_id]), + ) + ) + resources.defer(lambda: client.proxy.delete_key(key)) + + served = unwrap( + client.proxy.transport.get( + "/v1/mcp/tools", + headers=client.proxy.transport.bearer(key), + params=NoBody(), + response_type=PrefixedToolsResponse, + ) + ) + served_names = [tool.name for tool in served.tools] + prefixed = [ + name + for name in served_names + if name.endswith(SEARCH_LOGS_TOOL) and name != SEARCH_LOGS_TOOL + ] + assert prefixed, ( + f"a toolset-scoped key must be served {SEARCH_LOGS_TOOL} under its " + f"'{{server_prefix}}{{separator}}{SEARCH_LOGS_TOOL}' wire name; got {served_names}" + ) diff --git a/tests/e2e/models.py b/tests/e2e/models.py index bef199cdd16..a0f1e7c26d4 100644 --- a/tests/e2e/models.py +++ b/tests/e2e/models.py @@ -49,6 +49,7 @@ class KeyMetadata(BaseModel): class ObjectPermission(BaseModel): mcp_servers: list[str] | None = None mcp_access_groups: list[str] | None = None + mcp_toolsets: list[str] | None = None class KeyGenerateBody(BaseModel): @@ -110,6 +111,7 @@ class KeyInfo(BaseModel): budget_id: str | None = None litellm_budget_table: LiteLLMBudgetTable | None = None budget_limits: list[BudgetWindowState] | None = None + object_permission: ObjectPermission | None = None class KeyInfoResponse(BaseModel): @@ -711,7 +713,7 @@ class LiteLLMParamsBody(BaseModel): tpm: int | None = None -ModelMode = Literal["batch", "realtime", "image_generation"] +ModelMode = Literal["batch", "realtime", "image_generation", "responses"] class ModelInfoBody(BaseModel): @@ -777,7 +779,9 @@ class CredentialCreateResponse(BaseModel): class KeyUpdateBody(BaseModel): key: str - models: list[str] + models: list[str] | None = None + max_budget: float | None = None + object_permission: ObjectPermission | None = None class KeyBlockBody(BaseModel):