diff --git a/litellm/proxy/_experimental/mcp_server/toolset_db.py b/litellm/proxy/_experimental/mcp_server/toolset_db.py index ecaaf35e817..48bad178927 100644 --- a/litellm/proxy/_experimental/mcp_server/toolset_db.py +++ b/litellm/proxy/_experimental/mcp_server/toolset_db.py @@ -132,10 +132,18 @@ async def update_mcp_toolset( data: UpdateMCPToolsetRequest, touched_by: str, ) -> MCPToolset | None: - data_dict: Final = data.model_dump(exclude_none=True, exclude={"toolset_id"}) - if "tools" in data_dict: - data_dict["tools"] = json.dumps(data_dict["tools"]) - data_dict["updated_by"] = touched_by + """A partial update: absent keeps, null clears. A toolset always has a name and a + tool list, so a null ``toolset_name`` or ``tools`` is a no-op rather than a clear; + emptying the tool selection is an explicit ``[]``, which cannot be mistaken for a + caller that left the field out.""" + data_dict: Final = dict( # mutable-ok: Prisma requires a plain dict for JSON query serialization + ( + (field, json.dumps(value) if field == "tools" else value) + for field, value in data.model_dump(exclude_unset=True).items() + if field != "toolset_id" and (field not in ("toolset_name", "tools") or value is not None) + ), + updated_by=touched_by, + ) try: row: Final = await _toolset_table(prisma_client).update( where={"toolset_id": data.toolset_id}, diff --git a/litellm/proxy/management_endpoints/mcp_management_endpoints.py b/litellm/proxy/management_endpoints/mcp_management_endpoints.py index d5c3427f29a..2aa7fdc7393 100644 --- a/litellm/proxy/management_endpoints/mcp_management_endpoints.py +++ b/litellm/proxy/management_endpoints/mcp_management_endpoints.py @@ -2673,6 +2673,8 @@ if MCP_AVAILABLE: """ Updates the MCP Server in the db. + Partial update: a field left out of the payload keeps its stored value, and a field sent as null is cleared. + Parameters: - payload: UpdateMCPServerRequest - Required. The updated mcp server data. ``` @@ -3098,6 +3100,8 @@ if MCP_AVAILABLE: user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), litellm_changed_by: str | None = Header(None), ): + """Partial update: a field left out keeps its stored value, and a field sent as null is cleared, except + ``toolset_name`` and ``tools``, which a toolset always has; empty the tool selection with an explicit [].""" prisma_client: Final = get_prisma_client_or_throw("Database not connected. Connect a database to your proxy") if LitellmUserRoles.PROXY_ADMIN != user_api_key_dict.user_role: raise HTTPException( diff --git a/tests/e2e/coverage_registry/mcp.yaml b/tests/e2e/coverage_registry/mcp.yaml index ab644118a47..f853e9ff8d6 100644 --- a/tests/e2e/coverage_registry/mcp.yaml +++ b/tests/e2e/coverage_registry/mcp.yaml @@ -119,3 +119,11 @@ assertions: [succeeds] source: "server.py:1089" rationale: Smoke; rarely used; same auth model as tools +- id: mcp.list_tools.api_key.toolset_scoped + module: mcp + tier: P0 + operation: list_tools + auth_family: api_key + assertions: [toolset_scoped] + source: "user_api_key_auth_mcp.py:2137" + rationale: "A key granted a toolset lists exactly the toolset's tools: the rest of the server's catalog stays hidden and every stored name resolves" diff --git a/tests/e2e/coverage_registry/mgmt.yaml b/tests/e2e/coverage_registry/mgmt.yaml index 860d96a50b4..c8d7037d2fd 100644 --- a/tests/e2e/coverage_registry/mgmt.yaml +++ b/tests/e2e/coverage_registry/mgmt.yaml @@ -76,3 +76,13 @@ - {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.test_connection.happy_path, module: mgmt, tier: P0, surface: api, assertions: [happy_path], source: "_health_endpoints.py:1785", rationale: "Test Connection for a responses-mode Bedrock Mantle deployment reaches the live provider and reports success; this exact shape 500ed on an acompletion partial before v1.91.0", fail_before_fix: proven} +- {id: mgmt.mcp_server.new.persists, module: mgmt, tier: P0, surface: api, assertions: [persists], source: "mcp_management_endpoints.py:1577", rationale: "Every field of an admin-created MCP server reads back verbatim, by id and in the list, on every replica"} +- {id: mgmt.mcp_server.list.persists, module: mgmt, tier: P1, surface: api, assertions: [persists], source: "mcp_management_endpoints.py:1112", rationale: "The MCP page grid lists a created server with the same field values its detail view reports"} +- {id: mgmt.mcp_server.update.preserves_unrelated_fields, module: mgmt, tier: P0, surface: api, assertions: [preserves_unrelated_fields], source: "mcp_management_endpoints.py:2665", rationale: "A dashboard edit of one field leaves the others intact and is visible on every replica after one save; edits that took several saves to stick were a customer defect"} +- {id: mgmt.mcp_server.update.clear_persists, module: mgmt, tier: P0, surface: api, assertions: [clear_persists], source: "mcp_management_endpoints.py:2665", rationale: "An explicit null clears the stored field (absent keeps, null clears)"} +- {id: mgmt.mcp_server.delete.persists, module: mgmt, tier: P0, surface: api, assertions: [persists], source: "mcp_management_endpoints.py:2139", rationale: "A deleted server is gone by id and from the list on every replica"} +- {id: mgmt.mcp_toolset.new.persists, module: mgmt, tier: P0, surface: api, assertions: [persists], source: "mcp_management_endpoints.py:3009", rationale: "Toolset tools read back under the exact server_id and tool_name written; a toolset stored under one name and read under another granted nothing"} +- {id: mgmt.mcp_toolset.update.preserves_unrelated_fields, module: mgmt, tier: P0, surface: api, assertions: [preserves_unrelated_fields], source: "mcp_management_endpoints.py:3098", rationale: "Editing the description leaves the tools and name intact"} +- {id: mgmt.mcp_toolset.update.persists, module: mgmt, tier: P0, surface: api, assertions: [persists], source: "mcp_management_endpoints.py:3098", rationale: "Narrowing the tools to one entry reads back exactly that entry"} +- {id: mgmt.mcp_toolset.update.clear_persists, module: mgmt, tier: P0, surface: api, assertions: [clear_persists], source: "mcp_management_endpoints.py:3098", fail_before_fix: proven, rationale: "An explicit null clears the stored description; the update used to drop null and keep the old value"} +- {id: mgmt.mcp_toolset.delete.persists, module: mgmt, tier: P0, surface: api, assertions: [persists], source: "mcp_management_endpoints.py:3149", rationale: "A deleted toolset is gone by id and from the list on every replica"} diff --git a/tests/e2e/e2e_http.py b/tests/e2e/e2e_http.py index 9d5f1658e91..415c72bbb3c 100644 --- a/tests/e2e/e2e_http.py +++ b/tests/e2e/e2e_http.py @@ -49,6 +49,11 @@ class AnthropicHeaders(AuthHeaders): anthropic_version: str = Field(default="2023-06-01", alias="anthropic-version") +class PartialBody(BaseModel): + """A body for a partial-update route (absent = keep, null = clear): a field left + unset is omitted from the wire, and a field set to None is sent as JSON null.""" + + class NoBody(BaseModel): """Empty body/query for routes that take none.""" @@ -252,6 +257,13 @@ def assert_auth_denied(result: StreamingResponse, context: str) -> None: f"{context}: expected 401/403, got {result.status_code}: {result.body[:300]}" ) + +def wire_body(json: BaseModel) -> dict[str, object]: + if isinstance(json, PartialBody): + return json.model_dump(by_alias=True, exclude_unset=True) + return json.model_dump(by_alias=True, exclude_none=True) + + def _headers(headers: BaseModel) -> dict[str, str]: dumped: dict[str, object] = headers.model_dump(by_alias=True, exclude_none=True) return {key: str(value) for key, value in dumped.items()} @@ -307,9 +319,26 @@ def request_with_retry[T: RetryableResponse]( return issue() -def _classify[R: BaseModel]( - resp: requests.Response, response_type: type[R] -) -> Result[R]: +class ClassifiableResponse(Protocol): + """What classifying an outcome reads off a response. requests.Response satisfies + it, and so does a fake, so the classification rules are testable on their own.""" + + @property + def status_code(self) -> int: ... + + @property + def ok(self) -> bool: ... + + @property + def text(self) -> str: ... + + @property + def content(self) -> bytes: ... + + def json(self) -> object: ... + + +def classify[R: BaseModel](resp: ClassifiableResponse, response_type: type[R]) -> Result[R]: if resp.status_code == 401: return UnauthorizedError(body=resp.text) if resp.status_code == 429: @@ -317,7 +346,8 @@ def _classify[R: BaseModel]( if not resp.ok: return UnknownApiError(status_code=resp.status_code, body=resp.text) try: - return Success(status_code=resp.status_code, data=response_type.model_validate(resp.json())) + payload: Final[object] = resp.json() if resp.content else {} + return Success(status_code=resp.status_code, data=response_type.model_validate(payload)) except Exception as exc: # noqa: BLE001 - any parse/validation failure is a value return ValidationError(message=str(exc)) @@ -335,13 +365,13 @@ def post[R: BaseModel]( lambda: requests.post( str(url), headers=_headers(headers), - json=json.model_dump(by_alias=True, exclude_none=True), + json=wire_body(json), timeout=timeout, ) ) except requests.RequestException as exc: return NetworkError(message=str(exc)) - return _classify(resp, response_type) + return classify(resp, response_type) def get[R: BaseModel]( @@ -363,7 +393,7 @@ def get[R: BaseModel]( ) except requests.RequestException as exc: return NetworkError(message=str(exc)) - return _classify(resp, response_type) + return classify(resp, response_type) def get_external[R: BaseModel]( @@ -383,7 +413,7 @@ def get_external[R: BaseModel]( ) except requests.RequestException as exc: return NetworkError(message=str(exc)) - return _classify(resp, response_type) + return classify(resp, response_type) def delete[R: BaseModel]( @@ -400,14 +430,14 @@ def delete[R: BaseModel]( lambda: requests.delete( str(url), headers=_headers(headers), - json=json.model_dump(by_alias=True, exclude_none=True), + json=wire_body(json), params=_params(params), timeout=timeout, ) ) except requests.RequestException as exc: return NetworkError(message=str(exc)) - return _classify(resp, response_type) + return classify(resp, response_type) def patch[R: BaseModel]( @@ -423,13 +453,13 @@ def patch[R: BaseModel]( lambda: requests.patch( str(url), headers=_headers(headers), - json=json.model_dump(by_alias=True, exclude_none=True), + json=wire_body(json), timeout=timeout, ) ) except requests.RequestException as exc: return NetworkError(message=str(exc)) - return _classify(resp, response_type) + return classify(resp, response_type) def put[R: BaseModel]( @@ -445,13 +475,13 @@ def put[R: BaseModel]( lambda: requests.put( str(url), headers=_headers(headers), - json=json.model_dump(by_alias=True, exclude_none=True), + json=wire_body(json), timeout=timeout, ) ) except requests.RequestException as exc: return NetworkError(message=str(exc)) - return _classify(resp, response_type) + return classify(resp, response_type) def probe( @@ -555,7 +585,7 @@ def send( str(url), headers=_headers(headers), params=_params(params), - json=json.model_dump(by_alias=True, exclude_none=True), + json=wire_body(json), stream=stream, timeout=timeout, ) @@ -605,7 +635,7 @@ def upload[R: BaseModel]( ) except requests.RequestException as exc: return NetworkError(message=str(exc)) - return _classify(resp, response_type) + return classify(resp, response_type) def stream_binary( @@ -623,7 +653,7 @@ def stream_binary( resp = requests.post( str(url), headers=_headers(headers), - json=json.model_dump(by_alias=True, exclude_none=True), + json=wire_body(json), stream=True, timeout=timeout, ) diff --git a/tests/e2e/llm_translation/test_responses_bridge_streaming_e2e.py b/tests/e2e/llm_translation/test_responses_bridge_streaming_e2e.py index 9a45743a0cd..75817340876 100644 --- a/tests/e2e/llm_translation/test_responses_bridge_streaming_e2e.py +++ b/tests/e2e/llm_translation/test_responses_bridge_streaming_e2e.py @@ -16,8 +16,10 @@ into a chat completion chunk. Two customer-visible contracts only hold on that p from __future__ import annotations +from typing import Final, Literal + import pytest -from pydantic import BaseModel +from pydantic import BaseModel, Field from e2e_config import unique_marker from e2e_http import StreamingResponse @@ -51,7 +53,8 @@ class _BridgeChoice(BaseModel): class _BridgeChunk(BaseModel): id: str - choices: list[_BridgeChoice] = [] + object: Literal["chat.completion.chunk"] + choices: list[_BridgeChoice] = Field(default_factory=list) class _WeatherArgs(BaseModel): @@ -103,16 +106,19 @@ class TestResponsesBridgeChatCompletionsStreaming: resources.key(), ChatBody( model=bridged_model, - messages=[ChatMessage(role="user", content=f"Count from 1 to 5, one number per line. {unique_marker()}")], + messages=[ + ChatMessage(role="user", content=f"Count from 1 to 5, one number per line. {unique_marker()}") + ], max_tokens=64, stream=True, ), ) - chunks = _bridge_chunks(result) - ids = {chunk.id for chunk in chunks} + chunks: Final = _bridge_chunks(result) + assert len(chunks) > 1, "the shared-id contract needs more than one streamed chunk" + ids: Final = frozenset(chunk.id for chunk in chunks) assert len(ids) == 1, f"bridged stream used {len(ids)} different chunk ids: {sorted(ids)[:5]}" - assert ids.pop().startswith("chatcmpl-"), f"bridged chunk id is not chat-completion shaped: {chunks[0].id}" + assert chunks[0].id.strip(), "bridged stream emitted an empty chunk id" @pytest.mark.covers( "llm.chat_completions.openai.basic.stream.bridge_streams_sse", @@ -134,9 +140,9 @@ class TestResponsesBridgeChatCompletionsStreaming: chunks = _bridge_chunks(result) content = "".join(choice.delta.content or "" for chunk in chunks for choice in chunk.choices) assert content.strip(), f"bridged stream completed with no content deltas: {result.stream_events[:3]}" - assert any( - choice.finish_reason for chunk in chunks for choice in chunk.choices - ), f"bridged stream never emitted a finish_reason: {result.stream_events[-3:]}" + assert any(choice.finish_reason for chunk in chunks for choice in chunk.choices), ( + f"bridged stream never emitted a finish_reason: {result.stream_events[-3:]}" + ) assert result.stream_done, f"bridged stream did not terminate with [DONE]: {result.stream_events[-2:]}" @pytest.mark.covers( diff --git a/tests/e2e/logging/datadog_reader.py b/tests/e2e/logging/datadog_reader.py index d0f478185c2..368c20cb6aa 100644 --- a/tests/e2e/logging/datadog_reader.py +++ b/tests/e2e/logging/datadog_reader.py @@ -12,8 +12,12 @@ empty result. External reads go through ``e2e_http``. from __future__ import annotations +import math +import random import time -from dataclasses import dataclass +from collections.abc import Callable, Mapping +from dataclasses import dataclass, field +from typing import Final import pytest from pydantic import BaseModel, ConfigDict, Field @@ -27,17 +31,33 @@ from e2e_config import ( DD_SITE, POLL_TIMEOUT, ) -from e2e_http import URL, Headers, RateLimitedError, Success, post +from e2e_http import URL, Headers, StreamingResponse, send -#: How many rate-limited responses in a row one search tolerates before the -#: hard fail; each retry sleeps a full search interval, so this rides out a -#: burst from a concurrent consumer of the org-wide search budget. -_RATE_LIMIT_RETRIES = 5 +type SearchCall = Callable[[str, float], StreamingResponse] + + +def _seconds(value: str | None) -> float | None: + if value is None: + return None + try: + seconds: Final = float(value) + except ValueError: + return None + return seconds if math.isfinite(seconds) and seconds >= 0 else None + + +def _rate_limit_delay(headers: Mapping[str, str]) -> float: + delays: Final = tuple( + delay + for name in ("x-ratelimit-reset", "retry-after") + if (delay := _seconds(headers.get(name))) is not None + ) + return max(1.0, max(delays, default=DD_SEARCH_INTERVAL)) class _DdAuthHeaders(Headers): - api_key: str = Field(serialization_alias="DD-API-KEY") - app_key: str = Field(serialization_alias="DD-APPLICATION-KEY") + api_key: str = Field(serialization_alias="DD-API-KEY", repr=False) + app_key: str = Field(serialization_alias="DD-APPLICATION-KEY", repr=False) class _SearchFilter(BaseModel): @@ -88,8 +108,12 @@ class _SearchResponse(BaseModel): @dataclass(frozen=True, slots=True) class DdLogsReader: site: str - api_key: str - app_key: str + api_key: str = field(repr=False) + app_key: str = field(repr=False) + search: SearchCall | None = field(default=None, repr=False) + now: Callable[[], float] = field(default=time.monotonic, repr=False) + sleep: Callable[[float], None] = field(default=time.sleep, repr=False) + jitter: Callable[[], float] = field(default=random.random, repr=False) def events_for_marker(self, marker: str) -> list[DdLogEvent]: """Every ingested event whose attributes carry the marker. DataDog @@ -108,25 +132,28 @@ class DdLogsReader: a single event. A 429 backs off and retries - the search budget is org-wide, so another consumer can empty it under us - while any other failure stays a hard fail.""" - for _ in range(_RATE_LIMIT_RETRIES): - result = post( - URL(f"https://api.{self.site}/api/v2/logs/events/search"), - headers=_DdAuthHeaders(api_key=self.api_key, app_key=self.app_key), - json=_SearchRequest(filter=_SearchFilter(query=query)), - response_type=_SearchResponse, - timeout=30.0, - ) - match result: - case Success(data=page): - return [event.attributes for event in page.data] - case RateLimitedError(retry_after_seconds=retry_after): - time.sleep(retry_after if retry_after else DD_SEARCH_INTERVAL) - case failure: - pytest.fail(f"DataDog Logs Search API at api.{self.site} failed: {failure}") + return self._events_for_query(query, self.now() + POLL_TIMEOUT) + + def _events_for_query(self, query: str, deadline: float) -> list[DdLogEvent]: + search: Final = self.search or self._search_page + while (remaining := deadline - self.now()) > 0: + if (result := search(query, min(30.0, remaining))).ok: + return [event.attributes for event in _SearchResponse.model_validate_json(result.body).data] + if result.status_code != 429: + pytest.fail(f"DataDog Logs Search API at api.{self.site} failed with HTTP {result.status_code}") + if (delay := min(_rate_limit_delay(result.headers) + self.jitter(), deadline - self.now())) > 0: + self.sleep(delay) pytest.fail( - f"DataDog Logs Search API at api.{self.site} still rate-limited after " - f"{_RATE_LIMIT_RETRIES} retries {DD_SEARCH_INTERVAL}s apart - the org-wide " - "logs_public_search_api budget (2 requests per 10s) is exhausted by another consumer" + f"DataDog Logs Search API at api.{self.site} remained rate-limited for {POLL_TIMEOUT}s; " + "the org-wide logs_public_search_api budget is exhausted" + ) + + def _search_page(self, query: str, timeout: float) -> StreamingResponse: + return send( + URL(f"https://api.{self.site}/api/v2/logs/events/search"), + headers=_DdAuthHeaders(api_key=self.api_key, app_key=self.app_key), + json=_SearchRequest(filter=_SearchFilter(query=query)), + timeout=timeout, ) def poll_events_for_marker(self, marker: str) -> list[DdLogEvent]: @@ -140,33 +167,42 @@ class DdLogsReader: hide from the exactly-one assertion - real-DataDog jitter can surface one call's two events tens of seconds apart. Searches pace at DD_SEARCH_INTERVAL, not POLL_INTERVAL, to respect the search API's - request budget. At the deadline the last result is returned as-is.""" - deadline = time.monotonic() + POLL_TIMEOUT - while time.monotonic() < deadline: - events = self.events_for_query(query) + request budget. Discovery, quota retries, and duplicate detection share + one POLL_TIMEOUT deadline; an incomplete settle window fails closed.""" + deadline: Final = self.now() + POLL_TIMEOUT + while (remaining := deadline - self.now()) > 0: + events = self._events_for_query(query, deadline) if events: - return self._settled_events_for_query(query, events) - time.sleep(DD_SEARCH_INTERVAL) - return self.events_for_query(query) + return self._settled_events_for_query(query, events, deadline) + if (remaining := deadline - self.now()) > 0: + self.sleep(min(DD_SEARCH_INTERVAL, remaining)) + return [] - def _settled_events_for_query(self, query: str, events: list[DdLogEvent]) -> list[DdLogEvent]: + def _settled_events_for_query(self, query: str, events: list[DdLogEvent], deadline: float) -> list[DdLogEvent]: """Re-read at every search interval until the settle window closes; a duplicate ends the watch early because more waiting cannot clear it. Keep the last non-empty result: a transient empty search (index lag) must not erase events already confirmed earlier in the settle window. + A successful final search must reach the full settle window before the + shared read-back deadline; otherwise duplicate detection is incomplete. """ - settle_deadline = time.monotonic() + DD_SETTLE_SECONDS + settle_deadline: Final = self.now() + DD_SETTLE_SECONDS last_nonempty = events - while time.monotonic() < settle_deadline: - time.sleep(DD_SEARCH_INTERVAL) - latest = self.events_for_query(query) - if not latest: - continue + if len(events) > 1: + return events + while (remaining := deadline - self.now()) > 0: + self.sleep(min(DD_SEARCH_INTERVAL, remaining)) + if self.now() >= deadline: + break + latest = self._events_for_query(query, deadline) if len(latest) > 1: return latest - last_nonempty = latest - return last_nonempty + if latest: + last_nonempty = latest + if self.now() >= settle_deadline: + return last_nonempty + pytest.fail(f"DataDog log delivery could not complete its duplicate-detection window within {POLL_TIMEOUT}s") def build_dd_logs_reader() -> DdLogsReader: diff --git a/tests/e2e/logging/test_datadog_reader.py b/tests/e2e/logging/test_datadog_reader.py new file mode 100644 index 00000000000..910a1cefd42 --- /dev/null +++ b/tests/e2e/logging/test_datadog_reader.py @@ -0,0 +1,223 @@ +import json +from collections.abc import Iterator, Sequence +from dataclasses import dataclass +from typing import Final + +import pytest + +from datadog_reader import DdLogsReader +from datadog_reader import _DdAuthHeaders # pyright: ignore[reportPrivateUsage] # verifies private auth-header serialization +from e2e_config import DD_SEARCH_INTERVAL, POLL_TIMEOUT +from e2e_http import StreamingResponse + + +def test_failure_diagnostics_hide_credentials_without_changing_auth_headers() -> None: + api_key: Final = "test-datadog-api-secret" + app_key: Final = "test-datadog-app-secret" + reader: Final = DdLogsReader(site="datadoghq.com", api_key=api_key, app_key=app_key) + headers: Final = _DdAuthHeaders(api_key=api_key, app_key=app_key) + + for value in (reader, headers): + assert api_key not in repr(value) + assert app_key not in repr(value) + + assert headers.model_dump(by_alias=True) == { + "DD-API-KEY": api_key, + "DD-APPLICATION-KEY": app_key, + } + + +@dataclass +class Clock: + elapsed: float = 0.0 + + def now(self) -> float: + return self.elapsed + + def sleep(self, seconds: float) -> None: + self.elapsed += seconds + + +@dataclass +class Search: + responses: Iterator[StreamingResponse] + calls: tuple[tuple[str, float], ...] = () + + def __call__(self, query: str, timeout: float) -> StreamingResponse: + self.calls += ((query, timeout),) + return next(self.responses) + + +def _page(*event_ids: str) -> StreamingResponse: + return StreamingResponse( + status_code=200, + body=json.dumps({"data": [{"attributes": {"attributes": {"id": event_id}}} for event_id in event_ids]}), + ) + + +def _reader(responses: Sequence[StreamingResponse], clock: Clock) -> tuple[DdLogsReader, Search]: + search: Final = Search(iter(responses)) + return DdLogsReader( + site="us5.datadoghq.com", + api_key="test-api-secret", + app_key="test-app-secret", + search=search, + now=clock.now, + sleep=clock.sleep, + jitter=lambda: 0.25, + ), search + + +def test_429_honors_server_reset_and_preserves_duplicate_events() -> None: + clock: Final = Clock() + reader, search = _reader( + (StreamingResponse(status_code=429, body="", headers={"x-ratelimit-reset": "6"}), _page("first", "duplicate")), + clock, + ) + + events: Final = reader.events_for_query("test-marker") + + assert tuple(event.attributes["id"] for event in events) == ("first", "duplicate") + assert clock.elapsed == 6.25 + assert search.calls == (("test-marker", 30.0), ("test-marker", 30.0)) + + +@pytest.mark.parametrize("reset", ("", "invalid", "nan", "inf", "-1")) +def test_invalid_reset_uses_search_interval(reset: str) -> None: + clock: Final = Clock() + reader, _ = _reader( + (StreamingResponse(status_code=429, body="", headers={"x-ratelimit-reset": reset}), _page()), clock + ) + + assert reader.events_for_query("test-marker") == [] + assert clock.elapsed == DD_SEARCH_INTERVAL + 0.25 + + +def test_zero_reset_cannot_create_a_busy_retry_loop() -> None: + clock: Final = Clock() + reader, _ = _reader( + (StreamingResponse(status_code=429, body="", headers={"x-ratelimit-reset": "0"}), _page()), clock + ) + + assert reader.events_for_query("test-marker") == [] + assert clock.elapsed == 1.25 + + +def test_retry_after_is_not_shortened_by_an_earlier_reset() -> None: + clock: Final = Clock() + reader, _ = _reader( + (StreamingResponse(status_code=429, body="", headers={"x-ratelimit-reset": "2", "retry-after": "8"}), _page()), + clock, + ) + + assert reader.events_for_query("test-marker") == [] + assert clock.elapsed == 8.25 + + +def test_rate_limit_wait_stops_at_deadline_without_issuing_another_request() -> None: + clock: Final = Clock() + reader, search = _reader( + (StreamingResponse(status_code=429, body="", headers={"x-ratelimit-reset": str(POLL_TIMEOUT * 10)}),), clock + ) + + with pytest.raises(pytest.fail.Exception, match="remained rate-limited"): + reader.events_for_query("test-marker") + + assert clock.elapsed == POLL_TIMEOUT + assert search.calls == (("test-marker", 30.0),) + + +def test_late_retry_cannot_receive_a_fresh_request_timeout() -> None: + clock: Final = Clock() + reader, search = _reader( + (StreamingResponse(status_code=429, body="", headers={"x-ratelimit-reset": str(POLL_TIMEOUT - 5)}), _page()), + clock, + ) + + assert reader.events_for_query("test-marker") == [] + assert search.calls == (("test-marker", 30.0), ("test-marker", 4.75)) + + +@pytest.mark.parametrize("status", (-1, 401, 403, 500)) +def test_non_quota_failures_are_not_retried_or_treated_as_empty_results(status: int) -> None: + clock: Final = Clock() + reader, search = _reader((StreamingResponse(status_code=status, body=""), _page()), clock) + + with pytest.raises(pytest.fail.Exception, match=f"failed with HTTP {status}"): + reader.events_for_query("test-marker") + + assert search.calls == (("test-marker", 30.0),) + assert clock.elapsed == 0 + + +def test_polling_quota_retries_share_the_original_deadline() -> None: + clock: Final = Clock() + reader, search = _reader( + (_page(), StreamingResponse(status_code=429, body="", headers={"x-ratelimit-reset": str(POLL_TIMEOUT)})), + clock, + ) + + with pytest.raises(pytest.fail.Exception, match="remained rate-limited"): + reader.poll_events_for_query("test-marker") + + assert clock.elapsed == POLL_TIMEOUT + assert len(search.calls) == 2 + + +def test_empty_polling_does_not_start_a_final_search_after_its_deadline() -> None: + clock: Final = Clock() + attempts: Final = int(POLL_TIMEOUT / DD_SEARCH_INTERVAL) + reader, search = _reader((_page(),) * attempts, clock) + + assert reader.poll_events_for_query("test-marker") == [] + assert clock.elapsed == POLL_TIMEOUT + assert len(search.calls) == attempts + + +def test_settlement_quota_retries_keep_the_remaining_readback_budget() -> None: + clock: Final = Clock() + empty_reads: Final = int(POLL_TIMEOUT / DD_SEARCH_INTERVAL) - 2 + reader, search = _reader( + (_page(),) * empty_reads + + (_page("first"), StreamingResponse(status_code=429, body="", headers={"x-ratelimit-reset": str(POLL_TIMEOUT)})), + clock, + ) + + with pytest.raises(pytest.fail.Exception, match="remained rate-limited"): + reader.poll_events_for_query("test-marker") + + assert clock.elapsed == POLL_TIMEOUT + assert search.calls[-1] == ("test-marker", DD_SEARCH_INTERVAL) + assert len(search.calls) == empty_reads + 2 + + +def test_settlement_detects_a_duplicate_on_the_final_search() -> None: + clock: Final = Clock() + reader, _ = _reader((_page("first"), _page("first"), _page(), _page("first", "duplicate")), clock) + + events: Final = reader.poll_events_for_query("test-marker") + + assert tuple(event.attributes["id"] for event in events) == ("first", "duplicate") + assert clock.elapsed == 30 + + +def test_settlement_keeps_confirmed_events_through_empty_searches() -> None: + clock: Final = Clock() + reader, _ = _reader((_page("first"), _page(), _page(), _page()), clock) + + events: Final = reader.poll_events_for_query("test-marker") + + assert tuple(event.attributes["id"] for event in events) == ("first",) + assert clock.elapsed == 30 + + +def test_late_delivery_cannot_pass_without_a_complete_settle_window() -> None: + clock: Final = Clock() + empty_reads: Final = int(POLL_TIMEOUT / DD_SEARCH_INTERVAL) - 2 + reader, search = _reader((_page(),) * empty_reads + (_page("first"), _page("first")), clock) + + with pytest.raises(pytest.fail.Exception, match="duplicate-detection window"): + reader.poll_events_for_query("test-marker") + + assert clock.elapsed == POLL_TIMEOUT + assert len(search.calls) == empty_reads + 2 diff --git a/tests/e2e/management/management_client.py b/tests/e2e/management/management_client.py index 2b897f5f07f..1ef0d89a8f9 100644 --- a/tests/e2e/management/management_client.py +++ b/tests/e2e/management/management_client.py @@ -43,6 +43,9 @@ from models import ( KeyResetSpendBody, KeyResetSpendResponse, KeyUpdateBody, + McpServerCreateBody, + McpServerRow, + McpServerUpdateBody, ModelDeleteBody, OrgDeleteBody, OrgInfoParams, @@ -537,6 +540,38 @@ class ManagementClient: ).root ) + def create_mcp_server(self, body: McpServerCreateBody) -> McpServerRow: + return unwrap( + self.proxy.transport.post( + "/v1/mcp/server", + headers=self.proxy.transport.master, + json=body, + response_type=McpServerRow, + ) + ) + + def update_mcp_server(self, body: McpServerUpdateBody) -> McpServerRow: + """PUT /v1/mcp/server, the call behind the dashboard's Save Changes: a partial + update where a field left unset keeps its stored value and None clears it.""" + return unwrap( + self.proxy.transport.put( + "/v1/mcp/server", + headers=self.proxy.transport.master, + json=body, + response_type=McpServerRow, + ) + ) + + def delete_mcp_server(self, server_id: str) -> Result[NoBody]: + """DELETE /v1/mcp/server/{server_id}. Returns the outcome so the act phase can + unwrap it while a deferred teardown can ignore an already-deleted server.""" + return self.proxy.transport.delete( + f"/v1/mcp/server/{server_id}", + headers=self.proxy.transport.master, + json=NoBody(), + response_type=NoBody, + ) + def chat_status(self, key: str, model: str, content: str) -> StreamingResponse: return self.proxy.transport.send( "/chat/completions", diff --git a/tests/e2e/management/test_mcp_lifecycle_e2e.py b/tests/e2e/management/test_mcp_lifecycle_e2e.py new file mode 100644 index 00000000000..9257d697647 --- /dev/null +++ b/tests/e2e/management/test_mcp_lifecycle_e2e.py @@ -0,0 +1,294 @@ +"""Live e2e: the MCP server and toolset management routes' lifecycle contract. + +Two customer defects sit on these routes, and each step here is the read-back that +would have caught one of them: a dashboard edit that took several saves to stick +because the read landed on a replica the write had not reached, and a toolset whose +tools were stored under one name and read back under another, so it granted +nothing. Every read-back therefore polls every replica that serves the route +(ProxyClient.read_back_everywhere) and asserts the exact values written, and both +update routes are held to the same partial-update contract: a field left out of the +payload keeps its stored value, a field sent as null is cleared. The server URL is +unreachable on purpose; only persistence is under test, never a tool call. +""" + +from __future__ import annotations + +from collections.abc import Callable, Mapping +from typing import Final + +import pytest +from e2e_config import unique_marker +from e2e_http import unwrap +from lifecycle import ResourceManager +from management_client import ManagementClient +from models import ( + McpInfo, + McpServerCreateBody, + McpServerListResponse, + McpServerRow, + McpServerUpdateBody, + ToolsetCreateBody, + ToolsetListResponse, + ToolsetRow, + ToolsetTool, + ToolsetUpdateBody, +) + +pytestmark = pytest.mark.e2e + +UNREACHABLE_URL: Final = "https://e2e-fake-mcp.test.local/mcp" + + +def _create_server(client: ManagementClient, resources: ResourceManager) -> tuple[McpServerCreateBody, str]: + name: Final = f"e2e_mcp_lifecycle_{unique_marker()}" + body: Final = McpServerCreateBody( + server_name=name, + alias=name, + url=UNREACHABLE_URL, + transport="http", + description="e2e lifecycle server", + mcp_info=McpInfo( + server_name=f"{name} (display)", + description="shown on the MCP page", + logo_url="https://e2e.test.local/logo.png", + ), + ) + server_id: Final = client.create_mcp_server(body).server_id + resources.defer(lambda: client.delete_mcp_server(server_id)) + return body, server_id + + +def _assert_server_matches(row: McpServerRow, written: McpServerCreateBody, *, where: str) -> None: + stored: Final = (row.server_name, row.alias, row.url, row.transport, row.description, row.mcp_info) + expected: Final = ( + written.server_name, + written.alias, + written.url, + written.transport, + written.description, + written.mcp_info, + ) + assert stored == expected, f"{where}: stored {stored}, expected {expected}" + + +def _server_everywhere( + client: ManagementClient, server_id: str, *, settled: Callable[[McpServerRow], bool] +) -> Mapping[str, McpServerRow]: + return client.proxy.read_body_back_everywhere(f"/v1/mcp/server/{server_id}", McpServerRow, settled=settled) + + +def _listed_server_everywhere(client: ManagementClient, server_id: str) -> Mapping[str, McpServerRow]: + listings: Final = client.proxy.read_body_back_everywhere( + "/v1/mcp/server", + McpServerListResponse, + settled=lambda rows: any(row.server_id == server_id for row in rows.root), + ) + return {replica: next(row for row in rows.root if row.server_id == server_id) for replica, rows in listings.items()} + + +class TestMcpServerLifecycle: + @pytest.mark.covers("mgmt.mcp_server.new.persists") + def test_create_persists_every_field_on_every_replica( + self, client: ManagementClient, resources: ResourceManager + ) -> None: + body, server_id = _create_server(client, resources) + + by_id: Final = _server_everywhere(client, server_id, settled=lambda row: row.server_id == server_id) + for replica, row in by_id.items(): + _assert_server_matches(row, body, where=f"GET /v1/mcp/server/{server_id} on {replica}") + + @pytest.mark.skip( + reason=( + "product gap: GET /v1/mcp/server builds each row from the in-memory registry, whose " + "_build_mcp_server_table sets description from mcp_info['description'], so the list " + "reports the mcp_info description while GET /v1/mcp/server/{server_id} reports the " + "stored description column. A server created with both set to different text reads " + "back with two different descriptions depending on the route" + ) + ) + @pytest.mark.covers("mgmt.mcp_server.list.persists") + def test_created_server_is_listed_with_every_field( + self, client: ManagementClient, resources: ResourceManager + ) -> None: + body, server_id = _create_server(client, resources) + + for replica, row in _listed_server_everywhere(client, server_id).items(): + _assert_server_matches(row, body, where=f"GET /v1/mcp/server on {replica}") + + @pytest.mark.covers("mgmt.mcp_server.update.preserves_unrelated_fields") + def test_updating_only_the_alias_keeps_every_other_field_on_every_replica( + self, client: ManagementClient, resources: ResourceManager + ) -> None: + body, server_id = _create_server(client, resources) + renamed: Final = f"{body.alias}_renamed" + + _ = client.update_mcp_server(McpServerUpdateBody(server_id=server_id, alias=renamed)) + + after_one_put: Final = _server_everywhere(client, server_id, settled=lambda row: row.alias == renamed) + for replica, row in after_one_put.items(): + _assert_server_matches( + row, + body.model_copy(update={"alias": renamed}), + where=f"GET /v1/mcp/server/{server_id} on {replica} after one PUT of alias", + ) + + @pytest.mark.covers("mgmt.mcp_server.update.clear_persists") + def test_clearing_the_description_with_null_reads_back_null( + self, client: ManagementClient, resources: ResourceManager + ) -> None: + body, server_id = _create_server(client, resources) + + _ = client.update_mcp_server(McpServerUpdateBody(server_id=server_id, description=None)) + + cleared: Final = _server_everywhere(client, server_id, settled=lambda row: row.description is None) + for replica, row in cleared.items(): + _assert_server_matches( + row, + body.model_copy(update={"description": None}), + where=f"GET /v1/mcp/server/{server_id} on {replica} after PUT description=null", + ) + + @pytest.mark.covers("mgmt.mcp_server.delete.persists") + def test_delete_removes_the_server_from_every_replica( + self, client: ManagementClient, resources: ResourceManager + ) -> None: + _, server_id = _create_server(client, resources) + + _ = unwrap(client.delete_mcp_server(server_id)) + + gone: Final = client.proxy.gone_everywhere(f"/v1/mcp/server/{server_id}") + assert set(gone.values()) == {404}, f"a deleted server must 404 on every replica; got {dict(gone)}" + listings: Final = client.proxy.read_body_back_everywhere( + "/v1/mcp/server", + McpServerListResponse, + settled=lambda rows: all(row.server_id != server_id for row in rows.root), + ) + for replica, rows in listings.items(): + assert all(row.server_id != server_id for row in rows.root), ( + f"GET /v1/mcp/server on {replica} still lists the deleted server {server_id}" + ) + + +def _create_toolset( + client: ManagementClient, resources: ResourceManager, server_id: str +) -> tuple[ToolsetCreateBody, str]: + body: Final = ToolsetCreateBody( + toolset_name=f"e2e_toolset_{unique_marker()}", + description="e2e lifecycle toolset", + tools=[ + ToolsetTool(server_id=server_id, tool_name="search_datadog_logs"), + ToolsetTool(server_id=server_id, tool_name="get_datadog_metric"), + ], + ) + toolset_id: Final = client.proxy.create_toolset(body).toolset_id + resources.defer(lambda: client.proxy.delete_toolset(toolset_id)) + return body, toolset_id + + +def _assert_toolset_matches(row: ToolsetRow, written: ToolsetCreateBody, *, where: str) -> None: + stored: Final = (row.toolset_name, row.description, row.tools) + expected: Final = (written.toolset_name, written.description, written.tools) + assert stored == expected, f"{where}: stored {stored}, expected {expected}" + + +def _toolset_everywhere( + client: ManagementClient, toolset_id: str, *, settled: Callable[[ToolsetRow], bool] +) -> Mapping[str, ToolsetRow]: + return client.proxy.read_body_back_everywhere(f"/v1/mcp/toolset/{toolset_id}", ToolsetRow, settled=settled) + + +class TestMcpToolsetLifecycle: + @pytest.mark.covers("mgmt.mcp_toolset.new.persists") + def test_create_persists_both_tools_under_the_exact_names_written( + self, client: ManagementClient, resources: ResourceManager + ) -> None: + _, server_id = _create_server(client, resources) + body, toolset_id = _create_toolset(client, resources, server_id) + + by_id: Final = _toolset_everywhere(client, toolset_id, settled=lambda row: row.toolset_id == toolset_id) + for replica, row in by_id.items(): + _assert_toolset_matches(row, body, where=f"GET /v1/mcp/toolset/{toolset_id} on {replica}") + listings: Final = client.proxy.read_body_back_everywhere( + "/v1/mcp/toolset", + ToolsetListResponse, + settled=lambda rows: any(row.toolset_id == toolset_id for row in rows.root), + ) + for replica, rows in listings.items(): + _assert_toolset_matches( + next(row for row in rows.root if row.toolset_id == toolset_id), + body, + where=f"GET /v1/mcp/toolset on {replica}", + ) + + @pytest.mark.covers("mgmt.mcp_toolset.update.preserves_unrelated_fields") + def test_updating_only_the_description_keeps_the_tools_and_name( + self, client: ManagementClient, resources: ResourceManager + ) -> None: + _, server_id = _create_server(client, resources) + body, toolset_id = _create_toolset(client, resources, server_id) + + _ = client.proxy.update_toolset(ToolsetUpdateBody(toolset_id=toolset_id, description="edited")) + + edited: Final = _toolset_everywhere(client, toolset_id, settled=lambda row: row.description == "edited") + for replica, row in edited.items(): + _assert_toolset_matches( + row, + body.model_copy(update={"description": "edited"}), + where=f"GET /v1/mcp/toolset/{toolset_id} on {replica} after PUT of description", + ) + + @pytest.mark.covers("mgmt.mcp_toolset.update.persists") + def test_updating_the_tools_to_one_entry_reads_back_exactly_that_entry( + self, client: ManagementClient, resources: ResourceManager + ) -> None: + _, server_id = _create_server(client, resources) + body, toolset_id = _create_toolset(client, resources, server_id) + kept: Final = body.tools[:1] + + _ = client.proxy.update_toolset(ToolsetUpdateBody(toolset_id=toolset_id, tools=kept)) + + narrowed: Final = _toolset_everywhere(client, toolset_id, settled=lambda row: row.tools == kept) + for replica, row in narrowed.items(): + _assert_toolset_matches( + row, + body.model_copy(update={"tools": kept}), + where=f"GET /v1/mcp/toolset/{toolset_id} on {replica} after PUT of one tool", + ) + + @pytest.mark.covers("mgmt.mcp_toolset.update.clear_persists") + def test_clearing_the_description_with_null_reads_back_null( + self, client: ManagementClient, resources: ResourceManager + ) -> None: + _, server_id = _create_server(client, resources) + body, toolset_id = _create_toolset(client, resources, server_id) + + _ = client.proxy.update_toolset(ToolsetUpdateBody(toolset_id=toolset_id, description=None)) + + cleared: Final = _toolset_everywhere(client, toolset_id, settled=lambda row: row.description is None) + for replica, row in cleared.items(): + _assert_toolset_matches( + row, + body.model_copy(update={"description": None}), + where=f"GET /v1/mcp/toolset/{toolset_id} on {replica} after PUT description=null", + ) + + @pytest.mark.covers("mgmt.mcp_toolset.delete.persists") + def test_delete_removes_the_toolset_from_every_replica( + self, client: ManagementClient, resources: ResourceManager + ) -> None: + _, server_id = _create_server(client, resources) + _, toolset_id = _create_toolset(client, resources, server_id) + + _ = unwrap(client.proxy.delete_toolset(toolset_id)) + + gone: Final = client.proxy.gone_everywhere(f"/v1/mcp/toolset/{toolset_id}") + assert set(gone.values()) == {404}, f"a deleted toolset must 404 on every replica; got {dict(gone)}" + listings: Final = client.proxy.read_body_back_everywhere( + "/v1/mcp/toolset", + ToolsetListResponse, + settled=lambda rows: all(row.toolset_id != toolset_id for row in rows.root), + ) + for replica, rows in listings.items(): + assert all(row.toolset_id != toolset_id for row in rows.root), ( + f"GET /v1/mcp/toolset on {replica} still lists the deleted toolset {toolset_id}" + ) diff --git a/tests/e2e/mcp/datadog_mcp.py b/tests/e2e/mcp/datadog_mcp.py index d1ea53a0b3b..352b4446cfd 100644 --- a/tests/e2e/mcp/datadog_mcp.py +++ b/tests/e2e/mcp/datadog_mcp.py @@ -3,6 +3,7 @@ from __future__ import annotations import os +from collections.abc import Sequence from e2e_config import datadog_mcp_url, unique_marker from lifecycle import ResourceManager @@ -35,7 +36,11 @@ def register_datadog_mcp( resources: ResourceManager, *, mcp_access_groups: list[str] | None = None, + allowed_tools: Sequence[str] | None = (SEARCH_LOGS_TOOL,), ) -> str: + """Register the core Datadog toolset with its credentials from the env. By default + the server exposes only `search_datadog_logs`; pass `allowed_tools=None` to expose + every tool the core toolset serves.""" assert_dd_mcp_creds() name = f"e2e_dd_mcp_{unique_marker()}" server_id = client.register_server( @@ -47,7 +52,7 @@ def register_datadog_mcp( "DD-API-KEY": _dd_api_key(), "DD-APPLICATION-KEY": _dd_app_key(), }, - allowed_tools=[SEARCH_LOGS_TOOL], + allowed_tools=None if allowed_tools is None else list(allowed_tools), mcp_access_groups=mcp_access_groups, ) resources.defer(lambda: client.delete_server(server_id)) diff --git a/tests/e2e/mcp/mcp_client.py b/tests/e2e/mcp/mcp_client.py index 73453478e5a..210fc7a1e98 100644 --- a/tests/e2e/mcp/mcp_client.py +++ b/tests/e2e/mcp/mcp_client.py @@ -16,11 +16,11 @@ import time from collections.abc import Mapping from dataclasses import dataclass -from pydantic import BaseModel, ConfigDict, Field, RootModel +from pydantic import BaseModel, ConfigDict, Field from e2e_config import settle_propagation from e2e_http import Headers, NoBody, Result, Success, UnknownApiError, unwrap -from models import KeyGenerateBody, ObjectPermission +from models import KeyGenerateBody, McpServerListResponse, McpServerRow, ObjectPermission from proxy_client import ProxyClient McpToolArg = str | int | float | bool | list[str] | dict[str, str] @@ -46,16 +46,6 @@ class McpServerNewResponse(BaseModel): server_id: str -class McpServerRow(BaseModel): - server_id: str - alias: str | None = None - url: str | None = None - - -class McpServersListResponse(RootModel[list[McpServerRow]]): - pass - - class McpToolMcpInfo(BaseModel): server_id: str | None = None alias: str | None = None @@ -193,7 +183,7 @@ class McpClient: "/v1/mcp/server", headers=self.proxy.transport.master, params=NoBody(), - response_type=McpServersListResponse, + response_type=McpServerListResponse, ) ).root @@ -224,11 +214,16 @@ class McpClient: user_id: str, mcp_servers: list[str] | None, mcp_access_groups: list[str] | None = None, + mcp_toolsets: list[str] | None = None, models: list[str] | None = None, ) -> str: object_permission = ( - ObjectPermission(mcp_servers=mcp_servers, mcp_access_groups=mcp_access_groups) - if mcp_servers is not None or mcp_access_groups is not None + ObjectPermission( + mcp_servers=mcp_servers, + mcp_access_groups=mcp_access_groups, + mcp_toolsets=mcp_toolsets, + ) + if mcp_servers is not None or mcp_access_groups is not None or mcp_toolsets is not None else None ) return self.proxy.generate_key( @@ -272,6 +267,20 @@ class McpClient: ) time.sleep(self.proxy.poll_interval) + def await_tools(self, key: str, server_id: str, *, expected: frozenset[str]) -> frozenset[str]: + """Poll tools/list until `server_id`'s tools as `key` sees them are exactly + `expected`, and return the last listing either way, so the caller's equality + assertion names the difference. Fails at poll_timeout only when the read + itself never succeeded.""" + deadline = time.monotonic() + self.proxy.poll_timeout + while True: + result = self.list_tools(key) + if isinstance(result, Success) and result.data.tool_names_for_server(server_id) == expected: + return expected + if time.monotonic() >= deadline: + return unwrap(result).tool_names_for_server(server_id) + time.sleep(self.proxy.poll_interval) + def await_call_tool( self, key: str, diff --git a/tests/e2e/mcp/test_mcp_toolset_enforcement_e2e.py b/tests/e2e/mcp/test_mcp_toolset_enforcement_e2e.py new file mode 100644 index 00000000000..6b901145eb1 --- /dev/null +++ b/tests/e2e/mcp/test_mcp_toolset_enforcement_e2e.py @@ -0,0 +1,95 @@ +"""Live e2e: a key granted a toolset lists exactly the toolset's tools. + +An admin registers the real Datadog remote MCP server with its whole core toolset +exposed, discovers two of its tool names through a key granted the server outright, +and curates a toolset naming exactly those two. A second key is granted the server +plus that toolset, and its tools/list must come back as exactly those two names: no +more, so the rest of the server's catalog stays hidden behind the toolset, and no +fewer, so a tool stored under one name and read under another (which granted +nothing) fails here first. Requires DD_API_KEY + DD_APP_KEY (the suite's real MCP +upstream). +""" + +from __future__ import annotations + +from typing import Final + +import pytest +from datadog_mcp import SEARCH_LOGS_TOOL, register_datadog_mcp +from e2e_config import unique_marker +from e2e_http import unwrap +from lifecycle import ResourceManager +from mcp_client import McpClient +from models import ToolsetCreateBody, ToolsetTool + +pytestmark = pytest.mark.e2e + + +def _key( + client: McpClient, + resources: ResourceManager, + label: str, + *, + server_id: str, + toolset_id: str | None = None, +) -> str: + key: Final = client.generate_key( + user_id=f"e2e-mcp-{label}-{unique_marker()}", + mcp_servers=[server_id], + mcp_toolsets=None if toolset_id is None else [toolset_id], + ) + resources.defer(lambda: client.proxy.delete_key(key)) + return key + + +def _wire_prefix(wire_name: str, tool_name: str, catalog: frozenset[str]) -> str: + """The prefix tools/list puts in front of one server's tool names, measured off a + tool whose own name is known rather than guessed from the alias. A toolset grants + by the tool's own name, never the wire name, and the prefix is whatever the proxy + is configured to build (the alias, or a short server id), so measuring it is the + only way to cross between the two.""" + assert wire_name.endswith(tool_name), f"tools/list served {wire_name!r}, expected it to end with {tool_name!r}" + prefix: Final = wire_name[: len(wire_name) - len(tool_name)] + unprefixed: Final = frozenset(name for name in catalog if not name.startswith(prefix)) + assert not unprefixed, ( + f"every tool of one server shares the wire prefix {prefix!r}, so {sorted(unprefixed)} " + f"cannot be reduced to the names a toolset grants by" + ) + return prefix + + +class TestMcpToolsetEnforcement: + @pytest.mark.covers("mcp.list_tools.api_key.toolset_scoped") + def test_key_granted_a_toolset_lists_exactly_its_tools(self, client: McpClient, resources: ResourceManager) -> None: + server_id: Final = register_datadog_mcp(client, resources, allowed_tools=None) + client.await_registered(server_id) + + catalog_key: Final = _key(client, resources, "catalog", server_id=server_id) + known_wire: Final = client.await_tool(catalog_key, server_id, SEARCH_LOGS_TOOL) + catalog: Final = unwrap(client.list_tools(catalog_key)).tool_names_for_server(server_id) + assert len(catalog) > 2, ( + f"the Datadog core toolset must serve more tools than the toolset names, or the " + f"restriction has nothing to hide; got {sorted(catalog)}" + ) + prefix: Final = _wire_prefix(known_wire, SEARCH_LOGS_TOOL, catalog) + chosen_wire: Final = frozenset(sorted(catalog)[:2]) + chosen: Final = frozenset(name.removeprefix(prefix) for name in chosen_wire) + + toolset: Final = client.proxy.create_toolset( + ToolsetCreateBody( + toolset_name=f"e2e_toolset_{unique_marker()}", + description="two Datadog tools", + tools=[ToolsetTool(server_id=server_id, tool_name=name) for name in sorted(chosen)], + ) + ) + resources.defer(lambda: client.proxy.delete_toolset(toolset.toolset_id)) + assert frozenset(tool.tool_name for tool in toolset.tools) == chosen, ( + f"toolset stored {toolset.tools}, expected the two names {sorted(chosen)} verbatim" + ) + + scoped_key: Final = _key(client, resources, "toolset", server_id=server_id, toolset_id=toolset.toolset_id) + listed: Final = client.await_tools(scoped_key, server_id, expected=chosen_wire) + assert listed == chosen_wire, ( + f"a key granted the toolset must list exactly its two tools; " + f"got {sorted(listed)}, expected {sorted(chosen_wire)}" + ) diff --git a/tests/e2e/models.py b/tests/e2e/models.py index 9f2654e0eec..62810e6cfd9 100644 --- a/tests/e2e/models.py +++ b/tests/e2e/models.py @@ -10,6 +10,7 @@ from collections.abc import Sequence from datetime import datetime from typing import Final, Literal +from e2e_http import PartialBody from pydantic import AliasChoices, BaseModel, ConfigDict, Field, RootModel, model_serializer, model_validator # ---------- keys ---------- @@ -55,6 +56,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): @@ -77,7 +79,7 @@ class KeyGenerateBody(BaseModel): allowed_passthrough_routes: list[str] | None = None metadata: KeyMetadata | None = None object_permission: ObjectPermission | None = None - router_settings: "RouterSettingsOverride | None" = None + router_settings: RouterSettingsOverride | None = None class KeyGenerateResponse(BaseModel): @@ -516,6 +518,15 @@ class CountTokensResponse(BaseModel): # ---------- mcp servers ---------- +class McpInfo(BaseModel): + """The `mcp_info` display block stored on an MCP server; only the fields the + lifecycle test writes and reads back.""" + + server_name: str | None = None + description: str | None = None + logo_url: str | None = None + + class McpServerCreateBody(BaseModel): """POST /v1/mcp/server. For a gateway-managed OAuth server, `auth_type` is `oauth2` and `oauth2_flow` is `authorization_code`; the upstream endpoints @@ -530,6 +541,18 @@ class McpServerCreateBody(BaseModel): oauth2_flow: Literal["client_credentials", "authorization_code"] | None = None authorization_url: str | None = None token_url: str | None = None + server_name: str | None = None + description: str | None = None + mcp_info: McpInfo | None = None + + +class McpServerUpdateBody(PartialBody): + """PUT /v1/mcp/server: a field left unset keeps its stored value, a field set + to None is cleared.""" + + server_id: str + alias: str | None = None + description: str | None = None class McpServerInfo(BaseModel): @@ -543,6 +566,54 @@ class McpServerInfo(BaseModel): allow_all_keys: bool | None = None +class McpServerRow(McpServerInfo): + """A stored MCP server as the create, get, and list routes return it: the + fields the lifecycle test asserts survive the round trip.""" + + server_name: str | None = None + transport: str | None = None + description: str | None = None + mcp_info: McpInfo | None = None + + +class McpServerListResponse(RootModel[list[McpServerRow]]): + """GET /v1/mcp/server answers with a bare array of servers.""" + + +class ToolsetTool(BaseModel): + server_id: str + tool_name: str + + +class ToolsetCreateBody(BaseModel): + toolset_name: str + description: str | None = None + tools: list[ToolsetTool] + + +class ToolsetUpdateBody(PartialBody): + """PUT /v1/mcp/toolset: a field left unset keeps its stored value, a field set + to None is cleared.""" + + toolset_id: str + description: str | None = None + tools: list[ToolsetTool] | None = None + + +class ToolsetRow(BaseModel): + """A stored toolset as POST /v1/mcp/toolset, GET /v1/mcp/toolset/{toolset_id}, + and each row of GET /v1/mcp/toolset return it.""" + + toolset_id: str + toolset_name: str + description: str | None = None + tools: list[ToolsetTool] = Field(default_factory=list) + + +class ToolsetListResponse(RootModel[list[ToolsetRow]]): + """GET /v1/mcp/toolset answers with a bare array of toolsets.""" + + class EmbedBody(BaseModel): model: str input: str diff --git a/tests/e2e/proxy_client.py b/tests/e2e/proxy_client.py index 520cbfde5a9..1bac5116a9d 100644 --- a/tests/e2e/proxy_client.py +++ b/tests/e2e/proxy_client.py @@ -12,6 +12,7 @@ import time import warnings from collections.abc import Callable, Mapping from dataclasses import dataclass +from functools import reduce from datetime import datetime from types import MappingProxyType from typing import Final @@ -26,6 +27,7 @@ from e2e_http import ( Result, StreamingResponse, Success, + UnknownApiError, is_ok, unwrap, ) @@ -70,6 +72,9 @@ from models import ( SpendLogsPage, SpendLogsPageParams, SpendLogsParams, + ToolsetCreateBody, + ToolsetRow, + ToolsetUpdateBody, ) from e2e_config import ( CONTROL_PLANE_BASE_URL, @@ -82,7 +87,7 @@ from e2e_config import ( SLOW_PROVIDER_TIMEOUT_SECONDS, settle_propagation, ) -from transport import HttpTransport, SplitTransport, Transport +from transport import HttpTransport, SplitTransport, Transport, is_control_plane_path RowsPredicate = Callable[[list[SpendLogRow]], bool] @@ -235,6 +240,99 @@ def servable_timeout_message( ) +type ReplicaRead[T] = Callable[[float], T] + + +@dataclass(frozen=True, slots=True) +class EverywhereConverged[T]: + """Every replica answered with something `settled` accepts, keyed by replica.""" + + answers: Mapping[str, T] + + +@dataclass(frozen=True, slots=True) +class NeverConvergedOn[T]: + """`replica` ran out its budget without an answer `settled` accepts; `last` is + its final answer, so the failure can say what that replica still serves.""" + + replica: str + last: T + + +def _last_answer[T]( + read: ReplicaRead[T], + *, + settled: Callable[[T], bool], + timeout: float, + interval: float, + request_timeout: float, + now: Callable[[], float], + sleep: Callable[[float], None], +) -> T: + """Poll `read` until `settled` accepts its answer or `timeout` runs out, and + return the last answer either way. Each read's request timeout is clamped to + the budget left, and the final poll runs even when less than an interval + remains, so a deadline never skips the read that would have settled.""" + deadline: Final = now() + timeout + answer = read(min(request_timeout, timeout)) + while not settled(answer): + remaining = deadline - now() + if remaining <= 0: + return answer + sleep(min(interval, remaining)) + answer = read(min(request_timeout, remaining)) + return answer + + +def await_everywhere[T]( + reads: Mapping[str, ReplicaRead[T]], + *, + settled: Callable[[T], bool], + timeout: float, + interval: float, + request_timeout: float, + now: Callable[[], float], + sleep: Callable[[float], None], +) -> EverywhereConverged[T] | NeverConvergedOn[T]: + """`_last_answer` against every replica in turn, each with the full budget, so a + write counts as visible only once the last replica reflects it, and stop at the + first replica that never converges. Clock and sleep are injected.""" + def read_replica( + outcome: EverywhereConverged[T] | NeverConvergedOn[T], + item: tuple[str, ReplicaRead[T]], + ) -> EverywhereConverged[T] | NeverConvergedOn[T]: + if isinstance(outcome, NeverConvergedOn): + return outcome + replica, read = item + answer: Final = _last_answer( + read, + settled=settled, + timeout=timeout, + interval=interval, + request_timeout=request_timeout, + now=now, + sleep=sleep, + ) + if not settled(answer): + return NeverConvergedOn(replica=replica, last=answer) + return EverywhereConverged(answers=MappingProxyType({**outcome.answers, replica: answer})) + + initial: Final[EverywhereConverged[T] | NeverConvergedOn[T]] = EverywhereConverged(answers=MappingProxyType({})) + return reduce(read_replica, reads.items(), initial) + + +def _is_not_found[R: BaseModel](result: Result[R]) -> bool: + return isinstance(result, UnknownApiError) and result.status_code == 404 + + +def _status_of[R: BaseModel](result: Result[R]) -> int: + match result: + case Success(status_code=status_code) | UnknownApiError(status_code=status_code): + return status_code + case _: + return -1 + + type Poller[T] = Callable[[], T] @@ -321,6 +419,7 @@ def converge_timeout_message(*, what: str, replica: str, timeout: float, last_re class ProxyClient: transport: Transport replicas: Mapping[str, Transport] + control_replicas: Mapping[str, Transport] poll_timeout: float = 120.0 poll_interval: float = 5.0 model_servable_timeout: float = MODEL_SERVABLE_TIMEOUT @@ -569,6 +668,112 @@ class ProxyClient: if not is_ok(result): warnings.warn(f"delete_model({model_id!r}) failed: {result}", stacklevel=2) + # ---- replica read-back ---------------------------------------------- + + def replicas_for(self, path: str) -> Mapping[str, Transport]: + """The replicas that serve `path`: every data-plane replica for an LLM route, + and for a management route the control-plane replicas, since the data-plane + replicas trim management routes and answer them 404. A monolith serves both + from every replica, so a management read-back polls all of them; a split + deployment exposes one control-plane address (there is one backend process + behind it on the stack these suites run against), so it polls that. A + control plane fronting several backends would need its own replica list to + prove each one converged, the way PROXY_REPLICA_URLS does for the gateways. + Never empty: a read-back against no replica would assert nothing and pass.""" + replicas: Final = self.control_replicas if is_control_plane_path(path) else self.replicas + assert replicas, f"no replica is configured to serve {path}, so a read-back there would prove nothing" + return replicas + + def read_body_back_everywhere[R: BaseModel]( + self, path: str, response_type: type[R], *, settled: Callable[[R], bool] + ) -> Mapping[str, R]: + """GET `path` on every replica that serves it, polling each to poll_timeout + until `settled` accepts its body, and fail naming the first replica that + never converged. Returns each replica's settled body, keyed by replica, so + the caller can assert the rest of it.""" + outcome: Final = await_everywhere( + {url: self._reader(transport, path, response_type) for url, transport in self.replicas_for(path).items()}, + settled=lambda result: isinstance(result, Success) and settled(result.data), + timeout=self.poll_timeout, + interval=self.poll_interval, + request_timeout=REQUEST_TIMEOUT, + now=time.monotonic, + sleep=time.sleep, + ) + match outcome: + case EverywhereConverged(answers=answers): + return MappingProxyType({url: unwrap(result) for url, result in answers.items()}) + case NeverConvergedOn(replica=replica, last=last): + raise AssertionError( + f"GET {path} on {replica} never converged within {self.poll_timeout}s of the write; " + f"last read: {last}" + ) + + def gone_everywhere(self, path: str) -> Mapping[str, int]: + """Poll GET `path` on every replica that serves it until each stops serving + it, and fail naming the first replica that still does at poll_timeout. + Returns each replica's final status, so the caller asserts the 404 itself.""" + outcome: Final = await_everywhere( + {url: self._reader(transport, path, NoBody) for url, transport in self.replicas_for(path).items()}, + settled=_is_not_found, + timeout=self.poll_timeout, + interval=self.poll_interval, + request_timeout=REQUEST_TIMEOUT, + now=time.monotonic, + sleep=time.sleep, + ) + match outcome: + case EverywhereConverged(answers=answers): + return MappingProxyType({url: _status_of(result) for url, result in answers.items()}) + case NeverConvergedOn(replica=replica, last=last): + raise AssertionError( + f"GET {path} on {replica} still answers {self.poll_timeout}s after the delete; last read: {last}" + ) + + @staticmethod + def _reader[R: BaseModel](transport: Transport, path: str, response_type: type[R]) -> ReplicaRead[Result[R]]: + return lambda request_timeout: transport.get( + path, + headers=transport.master, + params=NoBody(), + response_type=response_type, + timeout=request_timeout, + ) + + # ---- mcp toolsets --------------------------------------------------- + + def create_toolset(self, body: ToolsetCreateBody) -> ToolsetRow: + return unwrap( + self.transport.post( + "/v1/mcp/toolset", + headers=self.transport.master, + json=body, + response_type=ToolsetRow, + ) + ) + + def update_toolset(self, body: ToolsetUpdateBody) -> ToolsetRow: + """PUT /v1/mcp/toolset: a partial update where a field left unset keeps its + stored value and None clears it.""" + return unwrap( + self.transport.put( + "/v1/mcp/toolset", + headers=self.transport.master, + json=body, + response_type=ToolsetRow, + ) + ) + + def delete_toolset(self, toolset_id: str) -> Result[NoBody]: + """DELETE /v1/mcp/toolset/{toolset_id}. Returns the outcome so the act phase + can unwrap it while a deferred teardown can ignore an already-deleted row.""" + return self.transport.delete( + f"/v1/mcp/toolset/{toolset_id}", + headers=self.transport.master, + json=NoBody(), + response_type=NoBody, + ) + def create_credential(self, body: CredentialCreateBody) -> None: unwrap( self.transport.post( @@ -736,7 +941,10 @@ def build_proxy_client( base URLs are the same for a monolithic proxy, so routing is then a no-op. ``replica_urls`` (PROXY_REPLICA_URLS) names every data-plane replica the model barrier polls directly; it is the data-plane URL itself unless the stack - exports each gateway's own address. + exports each gateway's own address. Management read-backs poll those same + replicas when the two planes share a base URL (a monolith, where every replica + serves every route) and the control plane alone when they differ (a split + deployment, where the data-plane replicas do not serve management routes). The endpoints are injectable for callers that resolve the proxy some other way than ``e2e_config``'s env names (see ``claude_code/_env.py``); they must @@ -764,9 +972,13 @@ def build_proxy_client( for url in replica_urls } ) + control_replicas: Final = ( + replicas if control_plane_base_url == base_url else MappingProxyType({control_plane_base_url: split.control}) + ) return ProxyClient( transport=split, replicas=replicas, + control_replicas=control_replicas, poll_timeout=POLL_TIMEOUT, poll_interval=POLL_INTERVAL, ) diff --git a/tests/e2e/router/test_auto_router_regressions_e2e.py b/tests/e2e/router/test_auto_router_regressions_e2e.py index 188db2a8eb5..374badcf5fc 100644 --- a/tests/e2e/router/test_auto_router_regressions_e2e.py +++ b/tests/e2e/router/test_auto_router_regressions_e2e.py @@ -41,6 +41,7 @@ which stores either the registered alias or the provider-prefixed form. import json import os from collections.abc import Iterator +from contextlib import ExitStack from dataclasses import dataclass from typing import Final @@ -120,19 +121,10 @@ class ResponsesApiResponse(BaseModel): @dataclass(frozen=True, slots=True) -class TagSplitDeployments: - """Scenario A mirrors the customer-shaped config from GitHub issue #36619: - plain deployment registered first, tier deployment and marker both tagged. - Scenario B flips both axes for GitHub issue #36621: marker registered first - and its tier deployment left untagged, so routing depends neither on - registration order nor on tier deployments carrying tags.""" - - tag_a: str - shared_a: str - tier_a: str - tag_b: str - shared_b: str - tier_b: str +class TagSplitDeployment: + tag: str + shared: str + tier: str @dataclass(frozen=True, slots=True) @@ -173,9 +165,7 @@ def _uniform_tier_config(tier_model: str) -> dict[str, object]: } -def _key_for( - proxy: ProxyClient, resources: ResourceManager, models: list[str], tag_filtering: bool = False -) -> str: +def _key_for(proxy: ProxyClient, resources: ResourceManager, models: list[str], tag_filtering: bool = False) -> str: key: Final = proxy.generate_key( KeyGenerateBody( models=models, @@ -211,46 +201,61 @@ def _assert_served_only_by(rows: list[SpendLogRow], allowed: frozenset[str], con ) -@pytest.fixture(scope="module") -def split(proxy: ProxyClient) -> Iterator[TagSplitDeployments]: +@pytest.fixture(scope="class") +def router_stack() -> Iterator[ExitStack]: + with ExitStack() as stack: + yield stack + + +def _register_models( + proxy: ProxyClient, stack: ExitStack, registrations: tuple[tuple[str, LiteLLMParamsBody], ...] +) -> None: + for name, params in registrations: + stack.callback(proxy.delete_model, proxy.create_model(name, params)) + + +def _tag_split(proxy: ProxyClient, stack: ExitStack, *, marker_first: bool) -> TagSplitDeployment: marker: Final = unique_marker() - deployments: Final = TagSplitDeployments( - tag_a=f"e2e-split-a-{marker}", - shared_a=f"e2e-autoroute-a-{marker}", - tier_a=f"e2e-tier-a-{marker}", - tag_b=f"e2e-split-b-{marker}", - shared_b=f"e2e-autoroute-b-{marker}", - tier_b=f"e2e-tier-b-{marker}", + named: Final = TagSplitDeployment( + tag=f"e2e-split-{marker}", + shared=f"e2e-autoroute-{marker}", + tier=f"e2e-tier-{marker}", ) anthropic_key: Final = _provider_key("ANTHROPIC_API_KEY") - marker_params_a: Final = LiteLLMParamsBody( - model="auto_router/complexity_router", - complexity_router_config=_uniform_tier_config(deployments.tier_a), - tags=[deployments.tag_a], + marker_registration: Final = ( + named.shared, + LiteLLMParamsBody( + model="auto_router/complexity_router", + complexity_router_config=_uniform_tier_config(named.tier), + tags=[named.tag], + ), ) - marker_params_b: Final = LiteLLMParamsBody( - model="auto_router/complexity_router", - complexity_router_config=_uniform_tier_config(deployments.tier_b), - tags=[deployments.tag_b], + tier_registration: Final = ( + named.tier, + LiteLLMParamsBody(model=CHEAP_MODEL, api_key=anthropic_key, tags=None if marker_first else [named.tag]), ) - registrations: Final[tuple[tuple[str, LiteLLMParamsBody], ...]] = ( - (deployments.shared_a, LiteLLMParamsBody(model=PLAIN_MODEL, api_key=anthropic_key)), - (deployments.tier_a, LiteLLMParamsBody(model=CHEAP_MODEL, api_key=anthropic_key, tags=[deployments.tag_a])), - (deployments.shared_a, marker_params_a), - (deployments.shared_b, marker_params_b), - (deployments.tier_b, LiteLLMParamsBody(model=CHEAP_MODEL, api_key=anthropic_key)), - (deployments.shared_b, LiteLLMParamsBody(model=PLAIN_MODEL, api_key=anthropic_key)), + plain_registration: Final = (named.shared, LiteLLMParamsBody(model=PLAIN_MODEL, api_key=anthropic_key)) + registrations: Final = ( + (marker_registration, tier_registration, plain_registration) + if marker_first + else (plain_registration, tier_registration, marker_registration) ) - created: Final = tuple(proxy.create_model(name, params) for name, params in registrations) - try: - yield deployments - finally: - for model_id in created: - proxy.delete_model(model_id) + _register_models(proxy, stack, registrations) + return named -@pytest.fixture(scope="module") -def zero_priced_alias(proxy: ProxyClient) -> Iterator[ZeroPricedAlias]: +@pytest.fixture(scope="class") +def plain_first_split(proxy: ProxyClient, router_stack: ExitStack) -> TagSplitDeployment: + return _tag_split(proxy, router_stack, marker_first=False) + + +@pytest.fixture(scope="class") +def marker_first_split(proxy: ProxyClient, router_stack: ExitStack) -> TagSplitDeployment: + return _tag_split(proxy, router_stack, marker_first=True) + + +@pytest.fixture(scope="class") +def zero_priced_alias(proxy: ProxyClient, router_stack: ExitStack) -> ZeroPricedAlias: marker: Final = unique_marker() named: Final = ZeroPricedAlias(alias=f"e2e-priced-alias-{marker}", tier=f"e2e-priced-tier-{marker}") alias_params: Final = LiteLLMParamsBody( @@ -263,16 +268,12 @@ def zero_priced_alias(proxy: ProxyClient) -> Iterator[ZeroPricedAlias]: (named.tier, LiteLLMParamsBody(model=CHEAP_MODEL, api_key=_provider_key("ANTHROPIC_API_KEY"))), (named.alias, alias_params), ) - created: Final = tuple(proxy.create_model(name, params) for name, params in registrations) - try: - yield named - finally: - for model_id in created: - proxy.delete_model(model_id) + _register_models(proxy, router_stack, registrations) + return named -@pytest.fixture(scope="module") -def heuristic_split(proxy: ProxyClient) -> Iterator[HeuristicSplit]: +@pytest.fixture(scope="class") +def heuristic_split(proxy: ProxyClient, router_stack: ExitStack) -> HeuristicSplit: marker: Final = unique_marker() named: Final = HeuristicSplit( alias=f"e2e-heuristic-router-{marker}", @@ -289,16 +290,12 @@ def heuristic_split(proxy: ProxyClient) -> Iterator[HeuristicSplit]: (named.strong, LiteLLMParamsBody(model=STRONG_MODEL, api_key=_provider_key("OPENAI_API_KEY"))), (named.alias, LiteLLMParamsBody(model="auto_router/complexity_router", complexity_router_config=config)), ) - created: Final = tuple(proxy.create_model(name, params) for name, params in registrations) - try: - yield named - finally: - for model_id in created: - proxy.delete_model(model_id) + _register_models(proxy, router_stack, registrations) + return named -@pytest.fixture(scope="module") -def semantic_auto_router(proxy: ProxyClient) -> Iterator[SemanticAutoRouter]: +@pytest.fixture(scope="class") +def semantic_auto_router(proxy: ProxyClient, router_stack: ExitStack) -> SemanticAutoRouter: marker: Final = unique_marker() named: Final = SemanticAutoRouter( marker=f"e2e-semantic-router-{marker}", @@ -321,16 +318,12 @@ def semantic_auto_router(proxy: ProxyClient) -> Iterator[SemanticAutoRouter]: (named.fallback, LiteLLMParamsBody(model=PLAIN_MODEL, api_key=_provider_key("ANTHROPIC_API_KEY"))), (named.marker, marker_params), ) - created: Final = tuple(proxy.create_model(name, params) for name, params in registrations) - try: - yield named - finally: - for model_id in created: - proxy.delete_model(model_id) + _register_models(proxy, router_stack, registrations) + return named -@pytest.fixture(scope="module") -def credentialed_alias(proxy: ProxyClient) -> Iterator[CredentialedAlias]: +@pytest.fixture(scope="class") +def credentialed_alias(proxy: ProxyClient, router_stack: ExitStack) -> CredentialedAlias: marker: Final = unique_marker() named: Final = CredentialedAlias(alias=f"e2e-cred-alias-{marker}", tier=f"e2e-cred-tier-{marker}") alias_params: Final = LiteLLMParamsBody( @@ -342,104 +335,110 @@ def credentialed_alias(proxy: ProxyClient) -> Iterator[CredentialedAlias]: (named.tier, LiteLLMParamsBody(model=CHEAP_MODEL, api_key=_provider_key("ANTHROPIC_API_KEY"))), (named.alias, alias_params), ) - created: Final = tuple(proxy.create_model(name, params) for name, params in registrations) - try: - yield named - finally: - for model_id in created: - proxy.delete_model(model_id) + _register_models(proxy, router_stack, registrations) + return named class TestTagSplitRouting: @pytest.mark.covers("reliability.routing.tagged_marker.request_tag_selects_marker") def test_body_tagged_chat_routes_through_the_marker_to_its_tier( - self, proxy: ProxyClient, resources: ResourceManager, split: TagSplitDeployments + self, proxy: ProxyClient, resources: ResourceManager, plain_first_split: TagSplitDeployment ) -> None: """Pins GitHub issue #36619: with tag filtering on, a chat request whose body metadata tags match the tagged marker under a shared model name is answered by the marker's tier deployment, not by the plain deployment that was registered under the name first.""" - key: Final = _key_for(proxy, resources, [split.shared_a, split.tier_a], tag_filtering=True) - chat: Final = unwrap(proxy.chat(key, _hello_chat_body(split.shared_a, tags=[split.tag_a]))) + key: Final = _key_for(proxy, resources, [plain_first_split.shared, plain_first_split.tier], tag_filtering=True) + chat: Final = unwrap(proxy.chat(key, _hello_chat_body(plain_first_split.shared, tags=[plain_first_split.tag]))) assert chat.choices, "tagged chat through the shared name returned no choices" rows: Final = proxy.poll_logs_for_key(key, min_rows=1) - _assert_served_only_by(rows, CHEAP_SERVED | {split.tier_a}, "body-tagged chat on the shared name") + _assert_served_only_by(rows, CHEAP_SERVED | {plain_first_split.tier}, "body-tagged chat on the shared name") @pytest.mark.covers("reliability.routing.tagged_marker.untagged_request_served_by_plain_deployment") def test_untagged_chat_is_always_served_by_the_plain_deployment( - self, proxy: ProxyClient, resources: ResourceManager, split: TagSplitDeployments + self, proxy: ProxyClient, resources: ResourceManager, plain_first_split: TagSplitDeployment ) -> None: """Pins GitHub issue #36620: untagged chat requests to the shared name succeed on every call and are all served by the plain deployment; the tagged marker never captures them, so no intermittent auto-router errors and no tier hijacking.""" - key: Final = _key_for(proxy, resources, [split.shared_a, split.tier_a], tag_filtering=True) + key: Final = _key_for(proxy, resources, [plain_first_split.shared, plain_first_split.tier], tag_filtering=True) for _ in range(5): - chat = unwrap(proxy.chat(key, _hello_chat_body(split.shared_a))) + chat = unwrap(proxy.chat(key, _hello_chat_body(plain_first_split.shared))) assert chat.choices, "untagged chat through the shared name returned no choices" rows: Final = proxy.poll_logs_for_key(key, min_rows=5) - _assert_served_only_by(rows, PLAIN_SERVED | {split.shared_a}, "untagged chat on the shared name") + _assert_served_only_by(rows, PLAIN_SERVED | {plain_first_split.shared}, "untagged chat on the shared name") @pytest.mark.covers("reliability.routing.tagged_marker.untagged_request_served_by_plain_deployment") def test_untagged_messages_is_served_by_the_plain_deployment( - self, proxy: ProxyClient, resources: ResourceManager, split: TagSplitDeployments + self, proxy: ProxyClient, resources: ResourceManager, plain_first_split: TagSplitDeployment ) -> None: """Pins GitHub issue #36620 on the /v1/messages surface: an untagged Anthropic-native request to the shared name is served by the plain deployment, not captured by the tagged marker.""" - key: Final = _key_for(proxy, resources, [split.shared_a, split.tier_a], tag_filtering=True) - answer: Final = unwrap(proxy.messages(key, _hello_messages_body(split.shared_a))) + key: Final = _key_for(proxy, resources, [plain_first_split.shared, plain_first_split.tier], tag_filtering=True) + answer: Final = unwrap(proxy.messages(key, _hello_messages_body(plain_first_split.shared))) assert answer.content or answer.choices, "untagged /v1/messages returned neither content nor choices" rows: Final = proxy.poll_logs_for_key(key, min_rows=1) - _assert_served_only_by(rows, PLAIN_SERVED | {split.shared_a}, "untagged /v1/messages on the shared name") + _assert_served_only_by( + rows, PLAIN_SERVED | {plain_first_split.shared}, "untagged /v1/messages on the shared name" + ) class TestUntaggedTierDeployments: @pytest.mark.covers("reliability.routing.tagged_marker.header_tag_selects_marker") def test_header_tagged_messages_routes_through_the_marker_to_an_untagged_tier( - self, proxy: ProxyClient, resources: ResourceManager, split: TagSplitDeployments + self, proxy: ProxyClient, resources: ResourceManager, marker_first_split: TagSplitDeployment ) -> None: """Pins GitHub issue #36621: a /v1/messages request tagged only via the x-litellm-tags header selects the tagged marker, and the rewrite still lands on the tier deployment even though that deployment carries no tags, because the marker consumed the routing tags.""" - key: Final = _key_for(proxy, resources, [split.shared_b, split.tier_b], tag_filtering=True) - headers: Final = TaggedAnthropicHeaders(authorization=f"Bearer {key}", x_litellm_tags=split.tag_b) + key: Final = _key_for( + proxy, resources, [marker_first_split.shared, marker_first_split.tier], tag_filtering=True + ) + headers: Final = TaggedAnthropicHeaders(authorization=f"Bearer {key}", x_litellm_tags=marker_first_split.tag) answer: Final = unwrap( proxy.transport.post( "/v1/messages", headers=headers, - json=_hello_messages_body(split.shared_b), + json=_hello_messages_body(marker_first_split.shared), response_type=AnthropicMessagesResponse, ) ) assert answer.content or answer.choices, "header-tagged /v1/messages returned neither content nor choices" rows: Final = proxy.poll_logs_for_key(key, min_rows=1) - _assert_served_only_by(rows, CHEAP_SERVED | {split.tier_b}, "header-tagged /v1/messages on the shared name") + _assert_served_only_by( + rows, CHEAP_SERVED | {marker_first_split.tier}, "header-tagged /v1/messages on the shared name" + ) @pytest.mark.covers("reliability.routing.tagged_marker.untagged_tier_deployments_still_served") def test_body_tagged_chat_reaches_the_untagged_tier_after_marker_rewrite( - self, proxy: ProxyClient, resources: ResourceManager, split: TagSplitDeployments + self, proxy: ProxyClient, resources: ResourceManager, marker_first_split: TagSplitDeployment ) -> None: """Pins the tag-consumption half of GitHub issue #36621: after the tagged marker rewrites the request to its tier model, the consumed routing tags no longer constrain deployment selection, so the untagged tier deployment serves the request instead of a strict-tag denial.""" - key: Final = _key_for(proxy, resources, [split.shared_b, split.tier_b], tag_filtering=True) - chat: Final = unwrap(proxy.chat(key, _hello_chat_body(split.shared_b, tags=[split.tag_b]))) + key: Final = _key_for( + proxy, resources, [marker_first_split.shared, marker_first_split.tier], tag_filtering=True + ) + chat: Final = unwrap( + proxy.chat(key, _hello_chat_body(marker_first_split.shared, tags=[marker_first_split.tag])) + ) assert chat.choices, "body-tagged chat through the marker-first shared name returned no choices" rows: Final = proxy.poll_logs_for_key(key, min_rows=1) - _assert_served_only_by(rows, CHEAP_SERVED | {split.tier_b}, "body-tagged chat with untagged tier") + _assert_served_only_by(rows, CHEAP_SERVED | {marker_first_split.tier}, "body-tagged chat with untagged tier") @pytest.mark.covers("reliability.routing.tagged_marker.tag_semantics_stay_strict") def test_tagged_call_straight_at_an_untagged_deployment_stays_denied( - self, proxy: ProxyClient, resources: ResourceManager, split: TagSplitDeployments + self, proxy: ProxyClient, resources: ResourceManager, marker_first_split: TagSplitDeployment ) -> None: """The tag-consumption fix must not loosen strict tag semantics: a tagged request aimed directly at an untagged deployment (no marker involved) is still rejected with the 401 tags-configuration error.""" - key: Final = _key_for(proxy, resources, [split.tier_b], tag_filtering=True) - result: Final = proxy.chat(key, _hello_chat_body(split.tier_b, tags=[split.tag_b])) + key: Final = _key_for(proxy, resources, [marker_first_split.tier], tag_filtering=True) + result: Final = proxy.chat(key, _hello_chat_body(marker_first_split.tier, tags=[marker_first_split.tag])) assert isinstance(result, UnauthorizedError), ( f"expected the tagged direct call to an untagged deployment to be denied with 401, got {result}" ) @@ -451,37 +450,39 @@ class TestUntaggedTierDeployments: class TestResponsesApiTagRouting: @pytest.mark.covers("reliability.routing.tagged_marker.responses_input_routes_through_marker") def test_header_tagged_responses_with_string_input_routes_to_the_tier( - self, proxy: ProxyClient, resources: ResourceManager, split: TagSplitDeployments + self, proxy: ProxyClient, resources: ResourceManager, plain_first_split: TagSplitDeployment ) -> None: """Pins the /v1/responses surface of the tag split (GitHub issues #36620/#36621): a /v1/responses request with string input, tagged via the x-litellm-tags header, succeeds and routes through the tagged marker to its tier.""" - key: Final = _key_for(proxy, resources, [split.shared_a, split.tier_a], tag_filtering=True) - headers: Final = TaggedAuthHeaders(authorization=f"Bearer {key}", x_litellm_tags=split.tag_a) + key: Final = _key_for(proxy, resources, [plain_first_split.shared, plain_first_split.tier], tag_filtering=True) + headers: Final = TaggedAuthHeaders(authorization=f"Bearer {key}", x_litellm_tags=plain_first_split.tag) body: Final = ResponsesBody( - model=split.shared_a, input=f"say hello {unique_marker()}", max_output_tokens=64 + model=plain_first_split.shared, input=f"say hello {unique_marker()}", max_output_tokens=64 ) answer: Final = unwrap( proxy.transport.post("/v1/responses", headers=headers, json=body, response_type=ResponsesApiResponse) ) assert answer.id, "header-tagged /v1/responses returned no response id" rows: Final = proxy.poll_logs_for_key(key, min_rows=1) - _assert_served_only_by(rows, CHEAP_SERVED | {split.tier_a}, "header-tagged /v1/responses string input") + _assert_served_only_by( + rows, CHEAP_SERVED | {plain_first_split.tier}, "header-tagged /v1/responses string input" + ) @pytest.mark.covers("reliability.routing.tagged_marker.responses_input_routes_through_marker") def test_body_tagged_responses_with_list_input_routes_to_the_tier( - self, proxy: ProxyClient, resources: ResourceManager, split: TagSplitDeployments + self, proxy: ProxyClient, resources: ResourceManager, plain_first_split: TagSplitDeployment ) -> None: """Pins the body-tag and list-input combination of the same split: /v1/responses with litellm_metadata.tags and structured input items routes through the tagged marker to its tier.""" - key: Final = _key_for(proxy, resources, [split.shared_a, split.tier_a], tag_filtering=True) + key: Final = _key_for(proxy, resources, [plain_first_split.shared, plain_first_split.tier], tag_filtering=True) body: Final = ResponsesBody( - model=split.shared_a, + model=plain_first_split.shared, input=[ResponsesInputItem(role="user", content=f"say hello {unique_marker()}")], max_output_tokens=64, - litellm_metadata=ResponsesTagMetadata(tags=[split.tag_a]), + litellm_metadata=ResponsesTagMetadata(tags=[plain_first_split.tag]), ) answer: Final = unwrap( proxy.transport.post( @@ -493,18 +494,18 @@ class TestResponsesApiTagRouting: ) assert answer.id, "body-tagged /v1/responses returned no response id" rows: Final = proxy.poll_logs_for_key(key, min_rows=1) - _assert_served_only_by(rows, CHEAP_SERVED | {split.tier_a}, "body-tagged /v1/responses list input") + _assert_served_only_by(rows, CHEAP_SERVED | {plain_first_split.tier}, "body-tagged /v1/responses list input") @pytest.mark.covers("reliability.routing.tagged_marker.untagged_request_served_by_plain_deployment") def test_untagged_responses_is_served_by_the_plain_deployment( - self, proxy: ProxyClient, resources: ResourceManager, split: TagSplitDeployments + self, proxy: ProxyClient, resources: ResourceManager, plain_first_split: TagSplitDeployment ) -> None: """Pins the untagged half of the /v1/responses tag split: an untagged request to the shared name is served by the plain deployment, matching the chat and messages surfaces.""" - key: Final = _key_for(proxy, resources, [split.shared_a, split.tier_a], tag_filtering=True) + key: Final = _key_for(proxy, resources, [plain_first_split.shared, plain_first_split.tier], tag_filtering=True) body: Final = ResponsesBody( - model=split.shared_a, input=f"say hello {unique_marker()}", max_output_tokens=64 + model=plain_first_split.shared, input=f"say hello {unique_marker()}", max_output_tokens=64 ) answer: Final = unwrap( proxy.transport.post( @@ -516,7 +517,9 @@ class TestResponsesApiTagRouting: ) assert answer.id, "untagged /v1/responses returned no response id" rows: Final = proxy.poll_logs_for_key(key, min_rows=1) - _assert_served_only_by(rows, PLAIN_SERVED | {split.shared_a}, "untagged /v1/responses on the shared name") + _assert_served_only_by( + rows, PLAIN_SERVED | {plain_first_split.shared}, "untagged /v1/responses on the shared name" + ) class TestStrategyAliasPricing: @@ -551,9 +554,7 @@ class TestComplexityHeuristicScope: while the accompanying ~2KB agent system prompt is packed with enough reasoning and complexity keywords that scoring the combined text lands in REASONING; only ask-only scoring keeps this on the cheap tier.""" - key: Final = _key_for( - proxy, resources, [heuristic_split.alias, heuristic_split.cheap, heuristic_split.strong] - ) + key: Final = _key_for(proxy, resources, [heuristic_split.alias, heuristic_split.cheap, heuristic_split.strong]) body: Final = ChatBody( model=heuristic_split.alias, messages=[ diff --git a/tests/e2e/test_e2e_http.py b/tests/e2e/test_e2e_http.py index 66841725d1d..81cd6c8d3d1 100644 --- a/tests/e2e/test_e2e_http.py +++ b/tests/e2e/test_e2e_http.py @@ -13,13 +13,24 @@ monkeypatches anything. from __future__ import annotations from collections.abc import Callable, Iterator, Mapping, Sequence -from dataclasses import dataclass, field +from dataclasses import dataclass from types import MappingProxyType from typing import Final import pytest - -from e2e_http import RETRY_ATTEMPTS, TRANSIENT_STATUSES, request_with_retry, streaming_outcome +from e2e_http import ( + RETRY_ATTEMPTS, + TRANSIENT_STATUSES, + NoBody, + PartialBody, + Success, + ValidationError, + classify, + request_with_retry, + streaming_outcome, + wire_body, +) +from pydantic import BaseModel, TypeAdapter @dataclass @@ -33,10 +44,10 @@ class FakeResponse: @dataclass class SleepRecorder: - delays: list[float] = field(default_factory=list) + delays: tuple[float, ...] = () def __call__(self, seconds: float) -> None: - self.delays.append(seconds) + self.delays += (seconds,) def _issue_from(responses: Sequence[FakeResponse]) -> Callable[[], FakeResponse]: @@ -55,7 +66,7 @@ class TestTransientRetryPolicy: sleep = SleepRecorder() result = request_with_retry(_issue_from(responses), sleep=sleep) assert result is responses[0] - assert sleep.delays == [] + assert sleep.delays == () assert responses[0].close_calls == 0 def test_429_is_never_retried(self) -> None: @@ -63,7 +74,7 @@ class TestTransientRetryPolicy: sleep = SleepRecorder() result = request_with_retry(_issue_from(responses), sleep=sleep) assert result is responses[0] - assert sleep.delays == [] + assert sleep.delays == () assert responses[0].close_calls == 0 def test_overloaded_529_retries_with_backoff_then_returns_the_success(self) -> None: @@ -71,7 +82,7 @@ class TestTransientRetryPolicy: sleep = SleepRecorder() result = request_with_retry(_issue_from(responses), sleep=sleep) assert result is responses[1] - assert sleep.delays == [0.5] + assert sleep.delays == (0.5,) assert responses[0].close_calls == 1 assert responses[1].close_calls == 0 @@ -80,7 +91,7 @@ class TestTransientRetryPolicy: sleep = SleepRecorder() result = request_with_retry(_issue_from(responses), sleep=sleep) assert result is responses[RETRY_ATTEMPTS - 1] - assert sleep.delays == [0.5, 1.0] + assert sleep.delays == (0.5, 1.0) assert [r.close_calls for r in responses] == [1, 1, 0, 0] @@ -134,3 +145,65 @@ class TestStreamEventArrivals: assert result.stream_events == [] assert result.stream_event_arrivals == [] assert result.body == "bad request" + + +class _ServerUpdate(PartialBody): + server_id: str + alias: str | None = None + description: str | None = None + + +class _ServerCreate(BaseModel): + alias: str + description: str | None = None + + +class TestWireBody: + """A partial-update body must put exactly the caller's choice on the wire: an + omitted field stays off it so the route keeps the stored value, and an explicit + None goes out as JSON null so the route clears it. Plain bodies keep dropping + None, which is what every create route expects.""" + + def test_partial_body_omits_unset_fields_and_sends_explicit_none_as_null(self) -> None: + assert wire_body(_ServerUpdate(server_id="s1", description=None)) == {"server_id": "s1", "description": None} + assert wire_body(_ServerUpdate(server_id="s1", alias="renamed")) == {"server_id": "s1", "alias": "renamed"} + + def test_plain_body_drops_none_fields(self) -> None: + assert wire_body(_ServerCreate(alias="a", description=None)) == {"alias": "a"} + + +_JSON: Final[TypeAdapter[object]] = TypeAdapter(object) + + +@dataclass +class FakeJsonResponse: + """The `classify` view of a response: a status, the raw body bytes, and the + parse that would raise on an empty one.""" + + status_code: int + content: bytes + + @property + def ok(self) -> bool: + return self.status_code < 400 + + @property + def text(self) -> str: + return self.content.decode() + + def json(self) -> object: + return _JSON.validate_json(self.content) + + +class TestClassifyEmptyBody: + """A delete that answers 202 with no body is a success, not a parse failure: + the MCP server and toolset delete routes both answer that way, and reading it + as a failure would hide a delete that did not happen behind one that did.""" + + def test_empty_2xx_body_is_a_success(self) -> None: + result: Final = classify(FakeJsonResponse(status_code=202, content=b""), NoBody) + assert isinstance(result, Success) and result.status_code == 202 + + def test_body_that_is_not_json_is_still_a_validation_failure(self) -> None: + result: Final = classify(FakeJsonResponse(status_code=200, content=b""), NoBody) + assert isinstance(result, ValidationError) diff --git a/tests/e2e/test_proxy_client.py b/tests/e2e/test_proxy_client.py index 2caac58333f..3b84a47e3cc 100644 --- a/tests/e2e/test_proxy_client.py +++ b/tests/e2e/test_proxy_client.py @@ -15,28 +15,35 @@ from collections.abc import Iterable, Mapping from dataclasses import dataclass from itertools import chain, repeat from types import MappingProxyType -from typing import Final +from typing import Final, cast import pytest - from e2e_config import parse_replica_urls from e2e_http import Result, Success from models import KeyInfo, KeyInfoResponse, ModelListEntry, ModelsListResponse from proxy_client import ( - Poller, ConvergeOutcome, Converged, + EverywhereConverged, ModelsPoller, + NeverConvergedOn, NotConverged, NotServableOn, + Poller, + ProxyClient, + ReplicaRead, Servable, await_converged_everywhere, + await_everywhere, await_servable_everywhere, - first_lagging_replica, + build_proxy_client, converge_timeout_message, + first_lagging_replica, ) +from transport import Transport MODEL: Final = "gpt-under-test" +_NO_TRANSPORTS: Final = cast(Transport, None) TIMEOUT: Final = 10.0 INTERVAL: Final = 2.0 RPM_BEFORE_UPDATE: Final = 100 @@ -187,3 +194,83 @@ class TestParseReplicaUrls: def test_falls_back_to_the_data_plane_address_when_unset(self) -> None: assert parse_replica_urls("", "http://lb") == ("http://lb",) + + +def _answers(answers: Iterable[str]) -> ReplicaRead[str]: + it: Final = iter(answers) + return lambda _timeout: next(it) + + +def _await_everywhere(reads: Mapping[str, ReplicaRead[str]]) -> EverywhereConverged[str] | NeverConvergedOn[str]: + clock: Final = FakeClock() + return await_everywhere( + reads, + settled=lambda answer: answer == "renamed", + timeout=TIMEOUT, + interval=INTERVAL, + request_timeout=5.0, + now=clock.now, + sleep=clock.sleep, + ) + + +class TestAwaitEverywhere: + def test_waits_for_the_lagging_replica_and_returns_every_settled_answer(self) -> None: + reads: Final = { + "gateway-1": _answers(repeat("renamed")), + "gateway-2": _answers(chain(repeat("stale", 2), repeat("renamed"))), + } + outcome: Final = _await_everywhere(reads) + assert isinstance(outcome, EverywhereConverged) + assert dict(outcome.answers) == {"gateway-1": "renamed", "gateway-2": "renamed"} + + def test_names_the_replica_that_never_converges_with_what_it_last_served(self) -> None: + reads: Final = { + "gateway-1": _answers(repeat("renamed")), + "gateway-2": _answers(repeat("stale")), + } + assert _await_everywhere(reads) == NeverConvergedOn(replica="gateway-2", last="stale") + + def test_polls_until_the_deadline_before_giving_up(self) -> None: + lagging: Final = chain(repeat("stale", int(TIMEOUT / INTERVAL)), repeat("renamed")) + outcome: Final = _await_everywhere({"gateway-1": _answers(lagging)}) + assert isinstance(outcome, EverywhereConverged), outcome + + +class TestReplicasFor: + def test_split_deployment_reads_management_routes_back_from_the_control_plane(self) -> None: + client: Final = build_proxy_client( + base_url="http://lb", + control_plane_base_url="http://backend", + replica_urls=("http://gateway-1", "http://gateway-2"), + ) + assert set(client.replicas_for("/key/info")) == {"http://backend"} + assert set(client.replicas_for("/v1/models")) == {"http://gateway-1", "http://gateway-2"} + + def test_monolith_reads_management_routes_back_from_every_replica(self) -> None: + client: Final = build_proxy_client( + base_url="http://lb", + control_plane_base_url="http://lb", + replica_urls=("http://pod-1", "http://pod-2"), + ) + assert set(client.replicas_for("/key/info")) == {"http://pod-1", "http://pod-2"} + + def test_mcp_admin_routes_read_back_from_every_data_plane_replica(self) -> None: + """/v1/mcp/* is a lazily mounted feature, so a data-plane replica serves it + too and answers from its own in-memory registry. Routing it to the control + plane would leave every replica but that one unproven, and would move the + tools/list barrier in mcp_client off the plane that serves tools/list.""" + client: Final = build_proxy_client( + base_url="http://lb", + control_plane_base_url="http://backend", + replica_urls=("http://gateway-1", "http://gateway-2"), + ) + assert set(client.replicas_for("/v1/mcp/server/abc")) == {"http://gateway-1", "http://gateway-2"} + assert set(client.replicas_for("/v1/mcp/toolset/abc")) == {"http://gateway-1", "http://gateway-2"} + + def test_a_route_no_replica_serves_is_refused_rather_than_read_back_vacuously(self) -> None: + """A read-back over zero replicas would satisfy every predicate and assert + nothing, so asking for one fails instead of passing silently.""" + client: Final = ProxyClient(transport=_NO_TRANSPORTS, replicas={}, control_replicas={}) + with pytest.raises(AssertionError, match="no replica is configured"): + _ = client.replicas_for("/v1/models") diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_partial_update.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_partial_update.py index c0d055edb7f..669e094fee4 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_partial_update.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_partial_update.py @@ -1,10 +1,11 @@ """ -Tests for partial-update semantics of PUT /v1/mcp/server. +Tests for partial-update semantics of PUT /v1/mcp/server and PUT /v1/mcp/toolset. A partial update must only write the fields the caller explicitly provided. Omitting a field must NOT reset it to its Pydantic schema default (e.g. ``transport=sse``, ``mcp_access_groups=[]``, ``allow_all_keys=False``), which -would silently overwrite the existing DB row. +would silently overwrite the existing DB row, and a field the caller sent as null +must be cleared rather than left at its stored value. """ import json @@ -850,3 +851,69 @@ async def test_cf_pair_switch_does_not_clear_dcr_bridge(): data = UpdateMCPServerRequest(server_id="s", auth_type="oauth_delegate") data_dict = await _run_update_with_existing(data, existing_auth_type="true_passthrough") assert "dcr_bridge" not in data_dict + + +def _mock_toolset_prisma(): + """A prisma double whose update answers with a row the reader can expand, so the + call under test returns instead of failing inside the row mapper.""" + updated_row = MagicMock() + updated_row.model_dump.return_value = { + "toolset_id": "ts-1", + "toolset_name": "ops", + "description": None, + "tools": "[]", + } + mock_prisma = MagicMock() + mock_prisma.db.litellm_mcptoolsettable = AsyncMock() + mock_prisma.db.litellm_mcptoolsettable.update = AsyncMock(return_value=updated_row) + return mock_prisma + + +async def _run_toolset_update(payload: dict) -> dict: + """The columns PUT /v1/mcp/toolset writes for this payload, minus the audit stamp + every write carries. The prisma double is injected, so nothing is patched.""" + from litellm.proxy._experimental.mcp_server.toolset_db import update_mcp_toolset + from litellm.types.mcp_server.mcp_toolset import UpdateMCPToolsetRequest + + mock_prisma = _mock_toolset_prisma() + await update_mcp_toolset(mock_prisma, UpdateMCPToolsetRequest.model_validate(payload), "test-user") + written = dict(mock_prisma.db.litellm_mcptoolsettable.update.call_args[1]["data"]) + assert written["updated_by"] == "test-user" + return {name: value for name, value in written.items() if name != "updated_by"} + + +@pytest.mark.asyncio +async def test_toolset_partial_update_clears_description_on_explicit_null(): + """The dump used to drop None, so a null description could never clear the stored + one: the toolset kept a description its owner had deleted.""" + assert await _run_toolset_update({"toolset_id": "ts-1", "description": None}) == {"description": None} + + +@pytest.mark.asyncio +async def test_toolset_partial_update_omits_the_fields_the_caller_left_out(): + tools = [{"server_id": "s1", "tool_name": "alpha"}] + assert await _run_toolset_update({"toolset_id": "ts-1", "tools": tools}) == {"tools": json.dumps(tools)} + + +@pytest.mark.asyncio +async def test_toolset_partial_update_ignores_null_tools_rather_than_revoking_them(): + """A client that sends tools=null means "leave the selection alone", so the grants + survive. Clearing them is an explicit [], which cannot be confused with an omitted + field; treating null as a clear would silently revoke every tool the toolset grants.""" + assert await _run_toolset_update({"toolset_id": "ts-1", "tools": None, "description": "kept"}) == { + "description": "kept" + } + + +@pytest.mark.asyncio +async def test_toolset_partial_update_empties_the_selection_on_an_explicit_empty_list(): + assert await _run_toolset_update({"toolset_id": "ts-1", "tools": []}) == {"tools": "[]"} + + +@pytest.mark.asyncio +async def test_toolset_partial_update_ignores_a_null_name(): + """A toolset always has a name, so a null toolset_name is a no-op, not a clear + that would write a NOT NULL column to null.""" + assert await _run_toolset_update({"toolset_id": "ts-1", "toolset_name": None, "description": "kept"}) == { + "description": "kept" + }