chore: merge latest staging into OCR cutover

Co-Authored-By: Ishaan Jaffer <155045088+ishaan-berri@users.noreply.github.com>

litellm/proxy/vector_store_files_endpoints/endpoints.py

@@ -227,6 +227,8 @@ async def _update_request_data_with_model_routing_hint(
     model_hint = data.get("model") or user_controlled_model_hint
     should_authorize_model_hint = isinstance(model_hint, str) and model_hint == user_controlled_model_hint
 
+    caller_team_id = getattr(user_api_key_dict, "team_id", None) if user_api_key_dict else None
+
     should_route = False
     credentials = None
     if isinstance(model_hint, str) and "*" in model_hint:
@@ -237,7 +239,9 @@ async def _update_request_data_with_model_routing_hint(
                     llm_router=llm_router,
                     user_api_key_dict=user_api_key_dict,
                 )
-            credentials = llm_router.get_deployment_credentials_with_provider(model_id=model_hint)
+            credentials = llm_router.get_deployment_credentials_with_provider(
+                model_id=model_hint, team_id=caller_team_id
+            )
             should_route = credentials is not None
     else:
         if isinstance(model_hint, str) and should_authorize_model_hint:
@@ -285,7 +289,7 @@ async def _update_request_data_with_model_routing_hint(
 
     openai_credentials = None
     for model_name in model_names_to_check:
-        credentials = llm_router.get_deployment_credentials_with_provider(model_id=model_name)
+        credentials = llm_router.get_deployment_credentials_with_provider(model_id=model_name, team_id=caller_team_id)
         if credentials is None:
             continue
 

litellm/router.py

@@ -7940,6 +7940,7 @@ class Router:
         self.model_id_to_deployment_index_map = {}  # Reset the index
         self.model_name_to_deployment_indices = {}  # Reset the model_name index
         self.team_model_to_deployment_indices = {}  # Reset the team_model index
+        self.team_pattern_routers = {}
         self.team_public_model_names = frozenset()
         # Reset per-strategy router registries so hot-reload doesn't leave
         # stale routers pointing at the old model_list.
@@ -8296,6 +8297,12 @@ class Router:
             public_model_name for _, public_model_name in self.team_model_to_deployment_indices
         )
 
+        for team_id in list(self.team_pattern_routers.keys()):
+            team_pattern_router = self.team_pattern_routers[team_id]
+            team_pattern_router.remove_deployment(model_id)
+            if not team_pattern_router.patterns:
+                del self.team_pattern_routers[team_id]
+
     def _update_team_model_index(self, model: dict, idx: int) -> None:
         """
         Helper to update team_model_to_deployment_indices for a single deployment.
@@ -8544,7 +8551,9 @@ class Router:
             int(max_output) if max_output is not None else None,
         )
 
-    def get_deployment_credentials_with_provider(self, model_id: str) -> Optional[Dict[str, Any]]:
+    def get_deployment_credentials_with_provider(
+        self, model_id: str, team_id: str | None = None
+    ) -> dict[str, Any] | None:
         """
         Get API credentials and provider info from a model name in model_list.
         Useful for passthrough endpoints (files, batches, etc.) that need credentials.
@@ -8554,6 +8563,9 @@ class Router:
 
         Args:
             model_id: Model ID or model name from model_list (e.g., "gpt-4o-litellm")
+            team_id: Optional team id of the caller. When set, team-scoped
+                deployments (indexed by team public model name, including team
+                wildcard models like "openai/*") are also considered.
 
         Returns:
             Dictionary containing api_key, api_base, custom_llm_provider, etc.
@@ -8572,9 +8584,22 @@ class Router:
         if deployment is None:
             deployment = self.get_deployment_by_model_group_name(model_group_name=model_id)
 
-        # If still not found, check for wildcard pattern matches
+        # If not found, check team-scoped deployments whose team public model
+        # name exactly matches model_id (wildcard team names are matched via
+        # team_pattern_routers below).
+        if deployment is None and team_id is not None:
+            team_indices = self.team_model_to_deployment_indices.get((team_id, model_id), [])
+            if team_indices:
+                team_model = self.model_list[team_indices[0]]
+                deployment = Deployment(**team_model) if isinstance(team_model, dict) else team_model
+
+        # If still not found, check for wildcard pattern matches. Team wildcard
+        # matches take priority so a global pattern (e.g. "openai/*") doesn't
+        # shadow the team's own entry.
         if deployment is None:
-            potential_wildcard_models = self.pattern_router.route(model_id) or []
+            team_pattern_router = self.team_pattern_routers.get(team_id) if team_id is not None else None
+            team_wildcard_models = (team_pattern_router.route(model_id) or []) if team_pattern_router else []
+            potential_wildcard_models = team_wildcard_models or self.pattern_router.route(model_id) or []
             if potential_wildcard_models:
                 # Use the first matching wildcard deployment
                 deployment_dict = potential_wildcard_models[0]

litellm/router_utils/pattern_match_deployments.py

@@ -73,6 +73,17 @@ class PatternMatchRouter:
             self.patterns[regex] = []
         self.patterns[regex].append(llm_deployment)
 
+    def remove_deployment(self, model_id: str) -> None:
+        """
+        Remove every deployment with the given model id from the pattern registry,
+        dropping any pattern whose deployment list becomes empty.
+        """
+        self.patterns = {
+            regex: remaining
+            for regex, deployments in self.patterns.items()
+            if (remaining := [d for d in deployments if (d.get("model_info") or {}).get("id") != model_id])
+        }
+
     def _pattern_to_regex(self, pattern: str) -> str:
         """
         Convert a wildcard pattern to a regex pattern

tests/e2e/CLAUDE.md

@@ -12,14 +12,24 @@ Each subdirectory under `tests/e2e/` is one suite, scoped to an endpoint family
 - `batches/` - the `/batches` endpoint (placeholder until the first test lands)
 - `realtime/` - realtime websocket sessions, including the pipecat audio path
 - `quota_management/` - quota enforcement and accounting, one subfolder per behavior: `ratelimit/` (rpm/tpm blocks, window reset, pacing headers on live traffic), `budgets/` (budget definition, enforcement, and reset windows: key, team, tag, soft, multi-window), and `spend_tracking/` (spend logging and cost attribution on `/spend/*`)
-- `management/` - key/team/user/organization management routes: create/update/delete persistence via the info routes, team membership, and llm-only-key route denials; also the dashboard UI behavior on top of them, driven through the proxy-served UI at /ui with playwright (optional dep behind importorskip)
-- `mcp/` - the MCP server surface over api_key auth: an admin registers an upstream MCP server through the management API and grants keys access via `object_permission.mcp_servers`, then the suite asserts tool listing and calling honor that permission (a key without the grant sees none of the server's tools and is refused a `tools/call` with a 403)
+- `management/` - key/team/user/organization management routes: create/update/delete persistence via the info routes, team membership, and llm-only-key route denials (API surface; not Playwright)
+- `mcp/` - the MCP server surface over api_key auth against the real Datadog remote MCP server only (see "MCP suite: real Datadog only" below)
 - `logging/` - logging-integration delivery (datadog and friends)
 - `security/` - secret handling and log-leak protection
 - `router/` - routing and reliability behavior (fallbacks, cooldowns)
 - `load/` - throughput/performance under concurrency: drives real concurrent traffic through the whole stack with Locust and asserts a throughput SLO; marked `load` so the parent conftest collects it last and it never perturbs latency-sensitive suites
 - `gateway/` - proxy configuration only (`litellm-config.yml`); no tests
-- `claude_code/` - the Claude Code compatibility matrix: drives the real `claude` CLI (and HTTP probes) against a proxy for each feature x provider cell, reporting tagged-union outcomes via the `compat_result` fixture; ships its own driver/builder/publisher plus `_*_unit_tests/` trees, and does not use the shared transport harness
+- `claude_code/` - the Claude Code compatibility matrix: drives the real `claude` CLI (and HTTP probes) against a proxy for each feature x provider cell, reporting tagged-union outcomes via the `compat_result` fixture; ships its own driver/builder/publisher plus `_*_unit_tests/` trees. The HTTP probes ride the shared transport (`ProxyClient.count_tokens` / `ProxyClient.messages`); the CLI-driving path stays bespoke
+
+## MCP suite: real Datadog only
+
+Every test under `tests/e2e/mcp/` must exercise the proxy against the real Datadog remote MCP server. Do not add a compose service, FastMCP fixture, mock upstream, or any other fake MCP host for this suite
+
+- Register via `register_datadog_mcp` in `tests/e2e/mcp/datadog_mcp.py` (or extend that helper if you need a different `toolsets=` / `allowed_tools` slice of the same Datadog endpoint). That posts `/v1/mcp/server` with `url=datadog_mcp_url(...)` and static headers `DD-API-KEY` / `DD-APPLICATION-KEY` from the process env
+- Auth is Datadog's documented CI/header path, not a browser OAuth authorize/token dance. Hard-fail when `DD_API_KEY` or `DD_APP_KEY` is missing (`assert_dd_mcp_creds`); never skip for a missing fake upstream
+- Prefer calling real Datadog tools that prove the product path (e.g. `search_datadog_logs` for list/call and permission denials). Seed a unique marker (`e2e-datadog-mcp-*`) in a chat completion when you need a log the tool can find; dual-read with `dd_logs` from conftest when delivery matters
+- Delete the MCP server (and any keys) through `resources.defer` the same way every other suite tears down
+- If a new MCP behavior cannot be covered with Datadog's tool surface, say so in the PR and get agreement before inventing another upstream; the default is always Datadog
 
 ## Lay the pattern down in a class
 

tests/e2e/batches/capabilities.py

@@ -7,8 +7,15 @@ import os
 from dataclasses import dataclass
 from typing import Literal
 
+from e2e_config import unique_marker
 from models import LiteLLMParamsBody
 
+_BATCH_RUN = unique_marker()
+
+
+def batch_model_name(base: str) -> str:
+    return f"{base}-{_BATCH_RUN}"
+
 
 def _env_ref(*names: str) -> str:
     for name in names:
@@ -91,24 +98,53 @@ class Capability:
 
     @property
     def jsonl_model(self) -> str:
-        return self.model if self.scenario == "unified" else self.raw_model
+        # Always the provider deployment name. Unified routes via
+        # target_model_names; the JSONL body.model must still be a name Azure /
+        # Vertex accept. Putting the proxy alias here used to depend on a perfect
+        # rewrite, and a stale or mis-selected deployment produced model_not_found.
+        return self.raw_model
 
 
 PROVIDERS: tuple[Provider, ...] = (
-    Provider("openai", "openai-batch", "gpt-4o-mini", can_cancel=True, can_list=True),
-    Provider("azure", "azure-batch", "gpt-5.4-mini-batch", can_cancel=True, can_list=True),
     Provider(
-        "vertex_ai", "vertex-batch", "gemini-2.5-flash", can_cancel=True, can_list=True
+        "openai", batch_model_name("openai-batch"), "gpt-4o-mini", can_cancel=True, can_list=True
+    ),
+    Provider(
+        "azure",
+        batch_model_name("azure-batch"),
+        "gpt-5.4-mini-batch",
+        can_cancel=True,
+        can_list=True,
+    ),
+    Provider(
+        "vertex_ai",
+        batch_model_name("vertex-batch"),
+        "gemini-2.5-flash",
+        can_cancel=True,
+        can_list=True,
     ),
     Provider(
         "bedrock",
-        "bedrock-batch",
+        batch_model_name("bedrock-batch"),
         "bedrock/us.anthropic.claude-haiku-4-5-20251001-v1:0",
         can_cancel=False,
         can_list=False,
     ),
 )
 
+def _model_for(provider_name: str) -> str:
+    for provider in PROVIDERS:
+        if provider.name == provider_name:
+            return provider.model
+    raise ValueError(
+        f"no batch provider named {provider_name!r} in PROVIDERS; "
+        f"known={[p.name for p in PROVIDERS]}"
+    )
+
+
+OPENAI_BATCH_MODEL = _model_for("openai")
+AZURE_BATCH_MODEL = _model_for("azure")
+
 BEDROCK_SCENARIOS: tuple[Scenario, ...] = ("unified",)
 
 

tests/e2e/batches/test_batches_e2e.py

@@ -33,9 +33,11 @@ from batch_client import (
     is_result_access_denied,
 )
 from capabilities import (
+    AZURE_BATCH_MODEL,
     BATCH_ID_SHAPE,
     CAPABILITIES,
     FILE_ID_SHAPE,
+    OPENAI_BATCH_MODEL,
     Capability,
     coverage_cells_for_lifecycle,
     matches_id_shape,
@@ -58,25 +60,69 @@ pytestmark = pytest.mark.e2e
 CREATED_BATCH_STATUSES = {"validating", "in_progress", "finalizing"}
 BATCH_CANCEL_DELAY_SECONDS = 2
 BATCH_TERMINAL_BEFORE_CANCEL = {"failed", "cancelled", "expired"}
-BATCH_CANCEL_RETRIES = 3
+BATCH_OP_RETRIES = 5
+# Azure / Vertex cancel and the pre-cancel re-retrieve are provider-side flakes
+# (connection refused, brief 500s) and the registry only has one basic cell per
+# provider (shared across scenarios). Create + retrieve already prove routing;
+# cancel is still deferred for cleanup, just not asserted for these two.
+_CANCEL_ASSERTED_PROVIDERS = frozenset({"openai"})
+
+
+def _transient_status(status_code: int) -> bool:
+    return status_code in {408, 429, 500, 502, 503, 504}
+
+
+def _backoff_seconds(attempt: int) -> float:
+    delays: tuple[float, ...] = (1.0, 2.0, 4.0, 8.0, 8.0)
+    return delays[min(attempt, len(delays) - 1)]
 
 
 def cancel_batch(
     client: BatchClient, batch_id: str, *, key: str, provider: str | None
 ) -> BatchObject:
     last = client.cancel_batch(batch_id, key=key, provider=provider)
-    for _ in range(BATCH_CANCEL_RETRIES - 1):
+    for attempt in range(BATCH_OP_RETRIES - 1):
         match last:
             case Success(data=data):
                 return data
-            case UnknownApiError(status_code=500):
-                time.sleep(1)
+            case UnknownApiError(status_code=code) if _transient_status(code):
+                time.sleep(_backoff_seconds(attempt))
                 last = client.cancel_batch(batch_id, key=key, provider=provider)
             case _:
                 break
     return unwrap(last)
 
 
+def retrieve_batch(
+    client: BatchClient, batch_id: str, *, key: str, provider: str | None
+) -> BatchObject:
+    last = client.retrieve_batch(batch_id, key=key, provider=provider)
+    for attempt in range(BATCH_OP_RETRIES - 1):
+        match last:
+            case Success(data=data):
+                return data
+            case UnknownApiError(status_code=code) if _transient_status(code):
+                time.sleep(_backoff_seconds(attempt))
+                last = client.retrieve_batch(batch_id, key=key, provider=provider)
+            case _:
+                break
+    return unwrap(last)
+
+
+def create_batch_resilient(
+    client: BatchClient, cap: Capability, file_id: str, key: str
+) -> StreamingResponse:
+    last = create_for_scenario(client, cap, file_id, key)
+    for attempt in range(BATCH_OP_RETRIES - 1):
+        if last.ok:
+            return last
+        if not _transient_status(last.status_code):
+            return last
+        time.sleep(_backoff_seconds(attempt))
+        last = create_for_scenario(client, cap, file_id, key)
+    return last
+
+
 def render_jsonl(model: str) -> bytes:
     line = {
         "custom_id": "req-1",
@@ -198,7 +244,7 @@ def test_batch_lifecycle(
         FILE_ID_SHAPE[cap.scenario], file.id
     ), f"{cap.id}: file id {file.id!r} is not a {FILE_ID_SHAPE[cap.scenario]} id"
 
-    created = create_for_scenario(client, cap, file.id, key)
+    created = create_batch_resilient(client, cap, file.id, key)
     require_successful_call(created)
     batch = BatchObject.model_validate_json(created.body)
     resources.defer(
@@ -218,7 +264,7 @@ def test_batch_lifecycle(
             cap.provider, batch.id
         ), f"{cap.provider} batch id {batch.id!r} not in that provider's native shape; misrouted?"
 
-    fetched = unwrap(client.retrieve_batch(batch.id, key=key, provider=provider))
+    fetched = retrieve_batch(client, batch.id, key=key, provider=provider)
     assert_batch_object(fetched)
     assert fetched.id == batch.id
     assert (
@@ -226,9 +272,9 @@ def test_batch_lifecycle(
     ), "retrieve changed input_file_id"
     assert fetched.status, "retrieved batch has no status"
 
-    if cap.can_cancel:
+    if cap.can_cancel and cap.provider in _CANCEL_ASSERTED_PROVIDERS:
         time.sleep(BATCH_CANCEL_DELAY_SECONDS)
-        pre_cancel = unwrap(client.retrieve_batch(batch.id, key=key, provider=provider))
+        pre_cancel = retrieve_batch(client, batch.id, key=key, provider=provider)
         assert (
             pre_cancel.status not in BATCH_TERMINAL_BEFORE_CANCEL
         ), (
@@ -240,10 +286,7 @@ def test_batch_lifecycle(
         cancelled = cancel_batch(client, batch.id, key=key, provider=provider)
         assert cancelled.id == batch.id
         assert cancelled.object == "batch"
-        valid_post_cancel = {"cancelling", "cancelled"}
-        if cap.provider == "vertex_ai":
-            valid_post_cancel |= CREATED_BATCH_STATUSES
-        assert cancelled.status in valid_post_cancel, (
+        assert cancelled.status in {"cancelling", "cancelled"}, (
             f"unexpected post-cancel status {cancelled.status!r}"
         )
 
@@ -281,12 +324,12 @@ def test_batch_lifecycle(
 def test_batch_key_model_access_denied(
     client: BatchClient, resources: ResourceManager, batch_deployments: None
 ) -> None:
-    key = resources.key(models=["openai-batch"])
+    key = resources.key(models=[OPENAI_BATCH_MODEL])
 
     denied_upload = client.upload_file(
-        content=render_jsonl("azure-batch"),
+        content=render_jsonl(AZURE_BATCH_MODEL),
         form=FileUploadForm(purpose="batch"),
-        model="azure-batch",
+        model=AZURE_BATCH_MODEL,
         key=key,
     )
     assert is_result_access_denied(
@@ -295,7 +338,7 @@ def test_batch_key_model_access_denied(
 
     raw_file = unwrap(
         client.upload_file(
-            content=render_jsonl("openai-batch"),
+            content=render_jsonl(OPENAI_BATCH_MODEL),
             form=FileUploadForm(purpose="batch"),
             key=key,
             provider="openai",
@@ -306,7 +349,7 @@ def test_batch_key_model_access_denied(
     )
 
     denied_create = client.create_batch(
-        body=BatchCreateBody(input_file_id=raw_file, model="azure-batch"), key=key
+        body=BatchCreateBody(input_file_id=raw_file, model=AZURE_BATCH_MODEL), key=key
     )
     assert is_model_access_denied(
         denied_create
@@ -323,9 +366,9 @@ def test_file_upload_and_delete_outputs(
     key = resources.key()
     file = unwrap(
         client.upload_file(
-            content=render_jsonl("openai-batch"),
+            content=render_jsonl(OPENAI_BATCH_MODEL),
             form=FileUploadForm(purpose="batch"),
-            model="openai-batch",
+            model=OPENAI_BATCH_MODEL,
             key=key,
         )
     )
@@ -390,7 +433,7 @@ def test_rate_limited_batch_create_leaves_no_unattributed_spend_row(
         client.upload_file(
             content=render_jsonl("gpt-4o-mini"),
             form=FileUploadForm(purpose="batch"),
-            model="openai-batch",
+            model=OPENAI_BATCH_MODEL,
             key=key,
         )
     )

tests/e2e/claude_code/_env.py

@@ -13,6 +13,8 @@ from typing import Mapping, NamedTuple
 
 import pytest
 
+from proxy_client import ProxyClient, build_proxy_client
+
 
 class ProxyConfig(NamedTuple):
     base_url: str
@@ -72,3 +74,30 @@ def require_proxy(
     if cfg is None:
         _fail_missing_proxy_env(compat_result)
     return cfg
+
+
+class ProxyClientConfig(NamedTuple):
+    client: ProxyClient
+    api_key: str
+
+
+def require_proxy_client(
+    compat_result,
+    *,
+    env: Mapping[str, str] | None = None,
+) -> ProxyClientConfig:
+    """Return the shared ``ProxyClient`` plus the master key the HTTP probes
+    authenticate with, or hard-fail the test.
+
+    Both planes of the built ``ProxyClient`` point at the one resolved base URL,
+    so the probes reuse the shared transport (split control/data-plane routing,
+    timeout, typed ``Result``) rather than hand-rolling ``httpx``. The api_key is
+    returned alongside because the probes call ``/v1/messages`` with the master
+    key (the compat matrix's credential), the same way the CLI rows do."""
+    cfg = require_proxy(compat_result, env=env)
+    client = build_proxy_client(
+        base_url=cfg.base_url,
+        master_key=cfg.api_key,
+        control_plane_base_url=cfg.base_url,
+    )
+    return ProxyClientConfig(client=client, api_key=cfg.api_key)

tests/e2e/claude_code/count_tokens/test_anthropic.py

@@ -39,7 +39,7 @@ from __future__ import annotations
 
 import pytest
 
-from claude_code._env import require_proxy
+from claude_code._env import require_proxy_client
 from claude_code.http_probe import (
     assert_count_tokens_shape,
     probe_count_tokens,
@@ -57,12 +57,12 @@ ANTHROPIC_MODELS = [
 def test_count_tokens_anthropic(compat_result):
     """Probe `/v1/messages/count_tokens` for each Anthropic tier and
     assert the response shape."""
-    base_url, api_key = require_proxy(compat_result)
+    client, api_key = require_proxy_client(compat_result)
 
     failures = []
     for model in ANTHROPIC_MODELS:
         result = probe_count_tokens(
-            base_url=base_url, api_key=api_key, model=model
+            client=client, api_key=api_key, model=model
         )
         shape_error = assert_count_tokens_shape(result)
         if shape_error is not None:

tests/e2e/claude_code/count_tokens/test_azure.py

@@ -39,7 +39,7 @@ from __future__ import annotations
 
 import pytest
 
-from claude_code._env import require_proxy
+from claude_code._env import require_proxy_client
 from claude_code.http_probe import (
     assert_count_tokens_shape,
     probe_count_tokens,
@@ -57,12 +57,12 @@ AZURE_MODELS = [
 def test_count_tokens_azure(compat_result):
     """Probe `/v1/messages/count_tokens` for each Azure (Microsoft Foundry) tier and
     assert the response shape."""
-    base_url, api_key = require_proxy(compat_result)
+    client, api_key = require_proxy_client(compat_result)
 
     failures = []
     for model in AZURE_MODELS:
         result = probe_count_tokens(
-            base_url=base_url, api_key=api_key, model=model
+            client=client, api_key=api_key, model=model
         )
         shape_error = assert_count_tokens_shape(result)
         if shape_error is not None:

tests/e2e/claude_code/count_tokens/test_bedrock_converse.py

@@ -39,7 +39,7 @@ from __future__ import annotations
 
 import pytest
 
-from claude_code._env import require_proxy
+from claude_code._env import require_proxy_client
 from claude_code.http_probe import (
     assert_count_tokens_shape,
     probe_count_tokens,
@@ -57,12 +57,12 @@ BEDROCK_CONVERSE_MODELS = [
 def test_count_tokens_bedrock_converse(compat_result):
     """Probe `/v1/messages/count_tokens` for each Bedrock (Converse) tier and
     assert the response shape."""
-    base_url, api_key = require_proxy(compat_result)
+    client, api_key = require_proxy_client(compat_result)
 
     failures = []
     for model in BEDROCK_CONVERSE_MODELS:
         result = probe_count_tokens(
-            base_url=base_url, api_key=api_key, model=model
+            client=client, api_key=api_key, model=model
         )
         shape_error = assert_count_tokens_shape(result)
         if shape_error is not None:

tests/e2e/claude_code/count_tokens/test_bedrock_invoke.py

@@ -39,7 +39,7 @@ from __future__ import annotations
 
 import pytest
 
-from claude_code._env import require_proxy
+from claude_code._env import require_proxy_client
 from claude_code.http_probe import (
     assert_count_tokens_shape,
     probe_count_tokens,
@@ -57,12 +57,12 @@ BEDROCK_INVOKE_MODELS = [
 def test_count_tokens_bedrock_invoke(compat_result):
     """Probe `/v1/messages/count_tokens` for each Bedrock (Invoke) tier and
     assert the response shape."""
-    base_url, api_key = require_proxy(compat_result)
+    client, api_key = require_proxy_client(compat_result)
 
     failures = []
     for model in BEDROCK_INVOKE_MODELS:
         result = probe_count_tokens(
-            base_url=base_url, api_key=api_key, model=model
+            client=client, api_key=api_key, model=model
         )
         shape_error = assert_count_tokens_shape(result)
         if shape_error is not None:

tests/e2e/claude_code/count_tokens/test_vertex_ai.py

@@ -39,7 +39,7 @@ from __future__ import annotations
 
 import pytest
 
-from claude_code._env import require_proxy
+from claude_code._env import require_proxy_client
 from claude_code.http_probe import (
     assert_count_tokens_shape,
     probe_count_tokens,
@@ -58,12 +58,12 @@ VERTEX_AI_MODELS = [
 def test_count_tokens_vertex_ai(compat_result):
     """Probe `/v1/messages/count_tokens` for each Vertex AI tier and
     assert the response shape."""
-    base_url, api_key = require_proxy(compat_result)
+    client, api_key = require_proxy_client(compat_result)
 
     failures = []
     for model in VERTEX_AI_MODELS:
         result = probe_count_tokens(
-            base_url=base_url, api_key=api_key, model=model
+            client=client, api_key=api_key, model=model
         )
         shape_error = assert_count_tokens_shape(result)
         if shape_error is not None:

tests/e2e/claude_code/http_probe.py

@@ -16,21 +16,43 @@ feature can be tested via the CLI, it should be, because the CLI path
 is closer to what real Claude Code users hit. HTTP probes are only for
 features the CLI can't reach.
 
-The probe deliberately uses a short timeout (30s) and small payloads:
-this is a "did the request shape survive the proxy's
-provider-specific transformations" test, not a load test, and a real
-endpoint regression typically surfaces in well under a second of wall
-time (400 / 500 from the upstream, or LiteLLM 500 on a transformation
-bug).
+The probes ride the shared transport: each takes an injected `ProxyClient`
+and issues its request through the shared `count_tokens` / `messages`
+methods, so they reuse the split control/data-plane routing, timeout,
+and typed `Result` handling the rest of `tests/e2e/` uses. This is a
+"did the request shape survive the proxy's provider-specific
+transformations" test, not a load test, and a real endpoint regression
+typically surfaces in well under a second of wall time (400 / 500 from
+the upstream, or LiteLLM 500 on a transformation bug).
 """
 
 from __future__ import annotations
 
-import json
-from dataclasses import dataclass
-from typing import Any, Mapping, Optional
+from typing import TYPE_CHECKING
 
-import httpx
+from pydantic import BaseModel
+
+from e2e_http import (
+    NetworkError,
+    RateLimitedError,
+    Result,
+    Success,
+    UnauthorizedError,
+    UnknownApiError,
+    ValidationError,
+)
+from models import (
+    AnthropicCustomTool,
+    AnthropicMessagesBody,
+    AnthropicMessagesResponse,
+    AnthropicTool,
+    AnthropicToolSearchTool,
+    ChatMessage,
+    CountTokensBody,
+    CountTokensResponse,
+    JsonSchemaProperty,
+    ToolInputSchema,
+)
 
 from claude_code.rate_limiter import (
     RateLimiter,
@@ -38,255 +60,169 @@ from claude_code.rate_limiter import (
     infer_provider,
 )
 
-
-DEFAULT_TIMEOUT_SECONDS = 30.0
+if TYPE_CHECKING:
+    from proxy_client import ProxyClient
 
 
-@dataclass
-class ProbeResult:
-    """Structured outcome of a single HTTP probe.
+# The tool_search discovery tool plus one trivial user tool, matching the
+# `tools` array real Claude Code emits when its MCP-tool-search beta is active.
+# The discovery tool's `_20251119`-suffixed type is what LiteLLM keys its
+# per-provider beta-header translation on; the user tool is included so the wire
+# shape mirrors what Claude Code sends rather than a semantically empty request.
+_TOOL_SEARCH_TOOLS: tuple[AnthropicTool, ...] = (
+    AnthropicToolSearchTool(
+        type="tool_search_tool_regex_20251119",
+        name="tool_search_tool_regex",
+    ),
+    AnthropicCustomTool(
+        name="add_numbers",
+        description="Add two integers",
+        input_schema=ToolInputSchema(
+            properties={
+                "a": JsonSchemaProperty(type="integer"),
+                "b": JsonSchemaProperty(type="integer"),
+            },
+            required=["a", "b"],
+        ),
+    ),
+)
 
-    `status_code` and `body` are the wire response; `payload` is the
-    parsed JSON body if the response was JSON, else None. Tests assert
-    on `status_code` + `payload` shape; `body` is preserved so failure
-    diagnostics can echo the raw error string (which is the only thing
-    a maintainer needs to triage a red cell).
-    """
+_TOOL_SEARCH_PROMPT = (
+    "If you have a tool to discover other tools, use it to "
+    "find one. Otherwise reply with the word 'done'."
+)
 
-    status_code: int
-    body: str
-    payload: Optional[Mapping[str, Any]] = None
-    error: Optional[str] = None
+
+def _acquire(model: str, rate_limiter: RateLimiter | None) -> None:
+    """Take one token from the cross-process per-provider limiter so probe
+    traffic counts against the same aggregate budget as the CLI rows. Without
+    this, an HTTP-probe row would fire unthrottled requests in parallel with
+    throttled CLI rows and silently violate the limiter's aggregate-rate
+    guarantee. `rate_limiter` is an injection seam for unit tests; production
+    callers leave it unset to use the process-wide default."""
+    limiter = rate_limiter if rate_limiter is not None else get_default_limiter()
+    limiter.acquire(infer_provider(model))
 
 
 def probe_count_tokens(
     *,
-    base_url: str,
+    client: ProxyClient,
     api_key: str,
     model: str,
     message: str = "hello world",
-    timeout: float = DEFAULT_TIMEOUT_SECONDS,
-    rate_limiter: Optional[RateLimiter] = None,
-) -> ProbeResult:
-    """POST to `{base_url}/v1/messages/count_tokens` for `model` and return the parsed result.
+    rate_limiter: RateLimiter | None = None,
+) -> Result[CountTokensResponse]:
+    """POST to `/v1/messages/count_tokens` for `model` and return the typed result.
 
-    The Anthropic / LiteLLM `count_tokens` endpoint accepts a request
-    body whose shape mirrors `/v1/messages` (model + messages), and
-    returns `{"input_tokens": N}` for a successful response. Anything
-    else -- non-200 status, non-JSON body, missing/non-int
-    `input_tokens` -- is a regression we want the cell to flip red on.
-
-    The same cross-process token-bucket limiter `cli_driver.run_claude`
-    uses is acquired here too, so probe rows count against the
-    aggregate per-provider budget. Without this, an HTTP-probe row
-    would fire unthrottled requests in parallel with throttled CLI
-    rows and silently violate the limiter's aggregate-rate guarantee.
-    `rate_limiter` is an injection seam for unit tests; production
-    callers should leave it unset to use the process-wide default.
+    The Anthropic / LiteLLM `count_tokens` endpoint accepts a request body whose
+    shape mirrors `/v1/messages` (model + messages) and returns
+    `{"input_tokens": N}` for a successful response. Anything else -- non-200
+    status, non-JSON body, missing/non-int `input_tokens` -- is a regression the
+    cell flips red on (see `assert_count_tokens_shape`).
     """
-    limiter = rate_limiter if rate_limiter is not None else get_default_limiter()
-    limiter.acquire(infer_provider(model))
-
-    url = base_url.rstrip("/") + "/v1/messages/count_tokens"
-    try:
-        response = httpx.post(
-            url,
-            headers={
-                "Authorization": f"Bearer {api_key}",
-                "Content-Type": "application/json",
-                # `anthropic-version` is required by Anthropic's native
-                # API and harmless on every other provider the proxy
-                # routes to. Matches what the Claude Code CLI sends
-                # for its own internal `count_tokens` calls.
-                "anthropic-version": "2023-06-01",
-            },
-            json={"model": model, "messages": [{"role": "user", "content": message}]},
-            timeout=timeout,
-        )
-    except httpx.HTTPError as exc:
-        return ProbeResult(status_code=0, body="", error=f"transport: {exc}")
-
-    body = response.text or ""
-    try:
-        payload = response.json() if body else None
-    except (json.JSONDecodeError, ValueError):
-        payload = None
-
-    return ProbeResult(
-        status_code=response.status_code,
-        body=body,
-        payload=payload,
+    _acquire(model, rate_limiter)
+    return client.count_tokens(
+        api_key,
+        CountTokensBody(model=model, messages=[ChatMessage(role="user", content=message)]),
     )
 
 
 def probe_tool_search(
     *,
-    base_url: str,
+    client: ProxyClient,
     api_key: str,
     model: str,
-    timeout: float = DEFAULT_TIMEOUT_SECONDS,
-    rate_limiter: Optional[RateLimiter] = None,
-) -> ProbeResult:
-    """POST to `{base_url}/v1/messages` with a `tool_search_tool_regex_20251119`
-    tool definition and return the result.
+    rate_limiter: RateLimiter | None = None,
+) -> Result[AnthropicMessagesResponse]:
+    """POST to `/v1/messages` with a `tool_search_tool_regex_20251119` tool
+    definition and return the typed result.
 
-    The shape of the tools array is the one Claude Code emits when its
-    MCP-tool-search beta is active: a `tool_search_tool_regex_20251119`
-    discovery tool (name `tool_search_tool_regex`) plus at least one
-    regular user tool to be searched. LiteLLM's
-    `is_tool_search_used` helper keys on the `_20251119`-suffixed type
-    string to decide whether to attach the provider-specific tool-search
-    beta header (`advanced-tool-use-2025-11-20` for Anthropic/Azure,
-    `tool-search-tool-2025-10-19` for Vertex/Bedrock). A proxy
-    regression in that translation will surface here as a 400 from
-    the upstream complaining about the tool type or beta header.
+    LiteLLM's `is_tool_search_used` helper keys on the `_20251119`-suffixed type
+    string to decide whether to attach the provider-specific tool-search beta
+    header (`advanced-tool-use-2025-11-20` for Anthropic/Azure,
+    `tool-search-tool-2025-10-19` for Vertex/Bedrock). A proxy regression in that
+    translation surfaces here as a 400 from the upstream complaining about the
+    tool type or beta header.
 
-    The prompt deliberately does not force a tool call -- the goal is
-    to verify the *request* round-trips without 400 and produces some
-    response, not to test whether the model decided to invoke
-    tool_search. That kind of behavior test would couple this row to
-    Claude Code's model behavior heuristics, which change weekly.
-
-    Like `probe_count_tokens`, this acquires one token from the
-    process-wide rate limiter so probe traffic counts against the
-    same aggregate per-provider budget as the CLI rows. `rate_limiter`
-    is a test seam; production callers should leave it unset.
+    The prompt deliberately does not force a tool call -- the goal is to verify
+    the *request* round-trips without 400 and produces some response, not to test
+    whether the model decided to invoke tool_search. That kind of behavior test
+    would couple this row to Claude Code's model behavior heuristics, which change
+    weekly.
     """
-    limiter = rate_limiter if rate_limiter is not None else get_default_limiter()
-    limiter.acquire(infer_provider(model))
-
-    url = base_url.rstrip("/") + "/v1/messages"
-    payload = {
-        "model": model,
-        "max_tokens": 64,
-        "messages": [
-            {
-                "role": "user",
-                "content": (
-                    "If you have a tool to discover other tools, use it to "
-                    "find one. Otherwise reply with the word 'done'."
-                ),
-            }
-        ],
-        "tools": [
-            # The tool_search discovery tool itself. Type is the SDK-
-            # version-pinned `_20251119` suffix; name is the canonical
-            # `tool_search_tool_regex` (no suffix) Anthropic accepts.
-            # LiteLLM keys its beta-header translation on the type.
-            {
-                "type": "tool_search_tool_regex_20251119",
-                "name": "tool_search_tool_regex",
-            },
-            # A trivial user tool for the discovery tool to potentially
-            # surface. Without at least one non-search tool the request
-            # is shape-valid but semantically empty; we include one so
-            # the wire shape mirrors what real Claude Code sends.
-            {
-                "name": "add_numbers",
-                "description": "Add two integers",
-                "input_schema": {
-                    "type": "object",
-                    "properties": {
-                        "a": {"type": "integer"},
-                        "b": {"type": "integer"},
-                    },
-                    "required": ["a", "b"],
-                },
-            },
-        ],
-    }
-    try:
-        response = httpx.post(
-            url,
-            headers={
-                "Authorization": f"Bearer {api_key}",
-                "Content-Type": "application/json",
-                "anthropic-version": "2023-06-01",
-            },
-            json=payload,
-            timeout=timeout,
-        )
-    except httpx.HTTPError as exc:
-        return ProbeResult(status_code=0, body="", error=f"transport: {exc}")
-
-    body = response.text or ""
-    try:
-        payload_out = response.json() if body else None
-    except (json.JSONDecodeError, ValueError):
-        payload_out = None
-
-    return ProbeResult(
-        status_code=response.status_code,
-        body=body,
-        payload=payload_out,
+    _acquire(model, rate_limiter)
+    return client.messages(
+        api_key,
+        AnthropicMessagesBody(
+            model=model,
+            max_tokens=64,
+            messages=[ChatMessage(role="user", content=_TOOL_SEARCH_PROMPT)],
+            tools=list(_TOOL_SEARCH_TOOLS),
+        ),
     )
 
 
-def assert_tool_search_shape(result: ProbeResult) -> Optional[str]:
+def _failure_diagnostic[R: BaseModel](result: Result[R], route: str) -> str:
+    """Map a non-success `Result` to a one-line diagnostic. The `status 429`
+    wording is load-bearing: the compat conftest classifies a rate-limited cell
+    by matching the failure text against `RATE_LIMIT_SHAPED_RE`, so the literal
+    `429` must survive into the reported error."""
+    match result:
+        case Success():
+            return ""
+        case UnauthorizedError():
+            return "status 401 (unauthorized)"
+        case RateLimitedError(body=body):
+            return f"status 429: {body[:400]}"
+        case UnknownApiError(status_code=status_code, body=body):
+            return f"status {status_code}: {body[:400]}"
+        case ValidationError(message=message):
+            return f"unexpected {route} response body: {message}"
+        case NetworkError(message=message):
+            return f"transport error: {message}"
+        case _:
+            return f"unexpected result: {result!r}"
+
+
+def assert_tool_search_shape(result: Result[AnthropicMessagesResponse]) -> str | None:
     """Return None on success, else describe the first violation.
 
     Acceptance criteria:
 
-      1. HTTP status is 200 (no 400 from the upstream rejecting the
-         tool_search tool type or a missing beta header).
-      2. Body is valid JSON.
-      3. Body has either `content` (Anthropic-shape passthrough) or
-         `choices` (LiteLLM normalized openai-shape, used by Bedrock
-         Converse). Either is acceptable -- the matrix cares that the
-         proxy *accepts and forwards* tool_search, not that the model
-         actually chose to invoke it. Tool-invocation behavior is a
-         model decision the matrix has no business asserting on.
-
-    The cell goes red when the upstream rejects the tool type, the
-    proxy drops the beta header, or the response shape is unusable.
-    Anything else (model decided to call or not call tool_search) is
-    irrelevant for this row.
+      1. The call succeeded (HTTP 200, no 400 from the upstream rejecting the
+         tool_search tool type or a missing beta header, no 429/401/transport
+         error).
+      2. The body has either `content` (Anthropic-shape passthrough) or `choices`
+         (LiteLLM normalized OpenAI-shape, used by Bedrock Converse). Either is
+         acceptable -- the matrix cares that the proxy *accepts and forwards*
+         tool_search, not that the model actually chose to invoke it.
     """
-    if result.error is not None:
-        return f"transport error: {result.error}"
-    if result.status_code != 200:
-        return f"status {result.status_code}: {result.body[:400]}"
-    if result.payload is None:
-        return f"non-JSON body: {result.body[:400]}"
-    if not isinstance(result.payload, Mapping):
-        return f"body is not a JSON object: {type(result.payload).__name__}"
-    # LiteLLM normalizes some provider responses to OpenAI shape
-    # (`choices`) and passes others through Anthropic-shape (`content`).
-    # Accept either; both prove the proxy round-tripped the request.
-    if "content" not in result.payload and "choices" not in result.payload:
-        return (
-            f"response has neither `content` nor `choices`: "
-            f"keys={sorted(result.payload.keys())}"
-        )
-    return None
+    match result:
+        case Success(data=data):
+            if data.content is None and data.choices is None:
+                keys = sorted(data.model_dump(exclude_none=True).keys())
+                return f"response has neither `content` nor `choices`: keys={keys}"
+            return None
+        case _:
+            return _failure_diagnostic(result, "/v1/messages")
 
 
-def assert_count_tokens_shape(result: ProbeResult) -> Optional[str]:
+def assert_count_tokens_shape(result: Result[CountTokensResponse]) -> str | None:
     """Return None on success, or an error string describing the first violation.
 
     Acceptance criteria are intentionally minimal:
 
-      1. HTTP status is 200.
-      2. Body is valid JSON.
-      3. Body has an `input_tokens` key whose value is a positive int.
+      1. The call succeeded (HTTP 200, valid JSON parsing into `input_tokens`).
+      2. `input_tokens` is a positive int.
 
-    Anything beyond that (cache token fields, server metadata) is
-    optional and varies by provider/transport. Asserting on extras
-    would create a brittle test that flips red on neutral protocol
-    drift; matrix cells should only go red on functional regressions
-    a Claude Code user would feel.
+    Anything beyond that (cache token fields, server metadata) is optional and
+    varies by provider/transport; asserting on extras would create a brittle test
+    that flips red on neutral protocol drift.
     """
-    if result.error is not None:
-        return f"transport error: {result.error}"
-    if result.status_code != 200:
-        return f"status {result.status_code}: {result.body[:400]}"
-    if result.payload is None:
-        return f"non-JSON body: {result.body[:400]}"
-    if not isinstance(result.payload, Mapping):
-        return f"body is not a JSON object: {type(result.payload).__name__}"
-    tokens = result.payload.get("input_tokens")
-    if not isinstance(tokens, int) or isinstance(tokens, bool):
-        return f"input_tokens missing or not an int: got {tokens!r}"
-    if tokens <= 0:
-        return f"input_tokens must be positive; got {tokens}"
-    return None
+    match result:
+        case Success(data=data):
+            if data.input_tokens <= 0:
+                return f"input_tokens must be positive; got {data.input_tokens}"
+            return None
+        case _:
+            return _failure_diagnostic(result, "/v1/messages/count_tokens")

tests/e2e/claude_code/tool_search/test_anthropic.py

@@ -45,7 +45,7 @@ from __future__ import annotations
 
 import pytest
 
-from claude_code._env import require_proxy
+from claude_code._env import require_proxy_client
 from claude_code.http_probe import (
     assert_tool_search_shape,
     probe_tool_search,
@@ -64,11 +64,11 @@ def test_tool_search_anthropic(compat_result):
     """Probe `/v1/messages` with a `tool_search_tool_regex_20251119`
     tool and assert the proxy + upstream accept it for every Anthropic
     tier."""
-    base_url, api_key = require_proxy(compat_result)
+    client, api_key = require_proxy_client(compat_result)
 
     failures = []
     for model in ANTHROPIC_MODELS:
-        result = probe_tool_search(base_url=base_url, api_key=api_key, model=model)
+        result = probe_tool_search(client=client, api_key=api_key, model=model)
         shape_error = assert_tool_search_shape(result)
         if shape_error is not None:
             error = f"[{model}] tool_search probe failed: {shape_error}"

tests/e2e/claude_code/tool_search/test_azure.py

@@ -45,7 +45,7 @@ from __future__ import annotations
 
 import pytest
 
-from claude_code._env import require_proxy
+from claude_code._env import require_proxy_client
 from claude_code.http_probe import (
     assert_tool_search_shape,
     probe_tool_search,
@@ -65,11 +65,11 @@ def test_tool_search_azure(compat_result):
     """Probe `/v1/messages` with a `tool_search_tool_regex_20251119`
     tool and assert the proxy + upstream accept it for every Azure (Microsoft Foundry)
     tier."""
-    base_url, api_key = require_proxy(compat_result)
+    client, api_key = require_proxy_client(compat_result)
 
     failures = []
     for model in AZURE_MODELS:
-        result = probe_tool_search(base_url=base_url, api_key=api_key, model=model)
+        result = probe_tool_search(client=client, api_key=api_key, model=model)
         shape_error = assert_tool_search_shape(result)
         if shape_error is not None:
             error = f"[{model}] tool_search probe failed: {shape_error}"

tests/e2e/claude_code/tool_search/test_bedrock_converse.py

@@ -45,7 +45,7 @@ from __future__ import annotations
 
 import pytest
 
-from claude_code._env import require_proxy
+from claude_code._env import require_proxy_client
 from claude_code.http_probe import (
     assert_tool_search_shape,
     probe_tool_search,
@@ -64,11 +64,11 @@ def test_tool_search_bedrock_converse(compat_result):
     """Probe `/v1/messages` with a `tool_search_tool_regex_20251119`
     tool and assert the proxy + upstream accept it for every Bedrock (Converse)
     tier."""
-    base_url, api_key = require_proxy(compat_result)
+    client, api_key = require_proxy_client(compat_result)
 
     failures = []
     for model in BEDROCK_CONVERSE_MODELS:
-        result = probe_tool_search(base_url=base_url, api_key=api_key, model=model)
+        result = probe_tool_search(client=client, api_key=api_key, model=model)
         shape_error = assert_tool_search_shape(result)
         if shape_error is not None:
             error = f"[{model}] tool_search probe failed: {shape_error}"

tests/e2e/claude_code/tool_search/test_bedrock_invoke.py

@@ -45,7 +45,7 @@ from __future__ import annotations
 
 import pytest
 
-from claude_code._env import require_proxy
+from claude_code._env import require_proxy_client
 from claude_code.http_probe import (
     assert_tool_search_shape,
     probe_tool_search,
@@ -68,11 +68,11 @@ def test_tool_search_bedrock_invoke(compat_result):
     """Probe `/v1/messages` with a `tool_search_tool_regex_20251119`
     tool and assert the proxy + upstream accept it for every Bedrock (Invoke)
     tier."""
-    base_url, api_key = require_proxy(compat_result)
+    client, api_key = require_proxy_client(compat_result)
 
     failures = []
     for model in BEDROCK_INVOKE_MODELS:
-        result = probe_tool_search(base_url=base_url, api_key=api_key, model=model)
+        result = probe_tool_search(client=client, api_key=api_key, model=model)
         shape_error = assert_tool_search_shape(result)
         if shape_error is not None:
             error = f"[{model}] tool_search probe failed: {shape_error}"

tests/e2e/claude_code/tool_search/test_vertex_ai.py

@@ -45,7 +45,7 @@ from __future__ import annotations
 
 import pytest
 
-from claude_code._env import require_proxy
+from claude_code._env import require_proxy_client
 from claude_code.http_probe import (
     assert_tool_search_shape,
     probe_tool_search,
@@ -65,11 +65,11 @@ def test_tool_search_vertex_ai(compat_result):
     """Probe `/v1/messages` with a `tool_search_tool_regex_20251119`
     tool and assert the proxy + upstream accept it for every Vertex AI
     tier."""
-    base_url, api_key = require_proxy(compat_result)
+    client, api_key = require_proxy_client(compat_result)
 
     failures = []
     for model in VERTEX_AI_MODELS:
-        result = probe_tool_search(base_url=base_url, api_key=api_key, model=model)
+        result = probe_tool_search(client=client, api_key=api_key, model=model)
         shape_error = assert_tool_search_shape(result)
         if shape_error is not None:
             error = f"[{model}] tool_search probe failed: {shape_error}"

tests/e2e/docker-compose.yml

@@ -1,7 +1,5 @@
 # local setup to run e2e tests
 configs:
-  mcp_upstream_server:
-    file: ../mcp_tests/mcp_e2e_upstream_server.py
   litellm_config:
     content: |
       general_settings:
@@ -133,26 +131,6 @@ services:
         target: /app/config.yaml
     command: ["--config", "/app/config.yaml", "--port", "4000"]
 
-# deterministic self-hosted upstream MCP server (FastMCP add/multiply over
-# streamable-http), reachable by the litellm container at mcp-upstream:8090/mcp.
-# Not a depends_on of litellm on purpose: only the mcp suite needs it, and it
-# boots long before the proxy is live, so it must not gate the other suites'
-# stack. The suite registers it through /v1/mcp/server at test time.
-  mcp-upstream:
-    image: ghcr.io/berriai/litellm:main-latest
-    entrypoint: ["python3", "/app/mcp_upstream_server.py"]
-    environment:
-      MCP_HOST: 0.0.0.0
-      MCP_PORT: "8090"
-    configs:
-      - source: mcp_upstream_server
-        target: /app/mcp_upstream_server.py
-    healthcheck:
-      test: ["CMD", "python3", "-c", "import socket; socket.create_connection(('127.0.0.1', 8090), 2).close()"]
-      interval: 3s
-      timeout: 3s
-      retries: 40
-
 # throwaway db
   db:
     image: postgres:16

tests/e2e/e2e_config.py

@@ -4,8 +4,18 @@ Shared by every e2e suite under tests/e2e/. Values come from the
 environment so the same tests run against localhost or a deployed proxy.
 """
 
+from __future__ import annotations
+
 import os
 import uuid
+from pathlib import Path
+
+from dotenv import load_dotenv
+
+# Local runs keep provider / DataDog keys in tests/e2e/.env (see CONTRIBUTING.md).
+# Compose injects them into the proxy container, but pytest on the host does not
+# inherit that file unless we load it. override=False so a real shell export wins.
+load_dotenv(Path(__file__).resolve().parent / ".env", override=False)
 
 PROXY_BASE_URL = os.environ.get("LITELLM_PROXY_URL", "http://localhost:4000").rstrip("/")
 MASTER_KEY = os.environ.get("LITELLM_MASTER_KEY", "sk-1234")
@@ -41,6 +51,7 @@ OTEL_QUERY_URL = os.environ.get("E2E_OTEL_QUERY_URL", "http://localhost:16686").
 DD_SITE = os.environ.get("DD_SITE", "datadoghq.com").strip()
 DD_API_KEY = os.environ.get("DD_API_KEY", "").strip()
 DD_APP_KEY = os.environ.get("DD_APP_KEY", "").strip()
+
 # After the first event is searchable, keep watching this long for a late
 # duplicate before the exactly-one assertion: real-DataDog ingestion jitter can
 # make one call's two events searchable tens of seconds apart, and a duplicate
@@ -68,6 +79,23 @@ LOAD_MIN_RPS = float(os.environ.get("E2E_LOAD_MIN_RPS", "355"))
 LOAD_MAX_FAILURE_RATIO = float(os.environ.get("E2E_LOAD_MAX_FAILURE_RATIO", "0.01"))
 
 
+def datadog_mcp_url(*, toolsets: str = "core") -> str:
+    """Regional Datadog remote MCP endpoint for this process's DD_SITE.
+
+    US1 is mcp.datadoghq.com; every other site is mcp.<site> (e.g. us5 ->
+    mcp.us5.datadoghq.com). A fixed mcp.datadoghq.com URL 403s when the keys
+    belong to a non-US1 org.
+    """
+    site = (
+        os.environ.get("DD_SITE", DD_SITE) or "datadoghq.com"
+    ).strip().removeprefix("https://").removeprefix("http://").rstrip("/")
+    if site.startswith("app."):
+        site = site[len("app.") :]
+    host = "mcp.datadoghq.com" if site in ("", "datadoghq.com") else f"mcp.{site}"
+    base = f"https://{host}/v1/mcp"
+    return f"{base}?toolsets={toolsets}" if toolsets else base
+
+
 def unique_marker() -> str:
     """A short unique token per call/run, so concurrent runs and the shared
     response cache never collide on prompts, tags, or customer ids."""

tests/e2e/e2e_http.py

@@ -32,6 +32,15 @@ class AuthHeaders(Headers):
     x_litellm_api_key: str | None = Field(default=None, alias="x-litellm-api-key")
 
 
+class AnthropicHeaders(AuthHeaders):
+    """Auth plus the ``anthropic-version`` header the Anthropic-native
+    /v1/messages and /v1/messages/count_tokens routes expect. It is harmless on
+    the other providers the proxy routes to, and matches what Claude Code sends
+    on its own internal calls."""
+
+    anthropic_version: str = Field(default="2023-06-01", alias="anthropic-version")
+
+
 class NoBody(BaseModel):
     """Empty body/query for routes that take none."""
 

tests/e2e/load/conftest.py

@@ -5,12 +5,12 @@ from collections.abc import Iterator
 import pytest
 from requests import RequestException
 
-from proxy_client import ProxyClient
 from e2e_http import NoBody, Success
 from load_client import LoadClient, build_client
 from load_constants import LOAD_MODEL
 from models import KeyGenerateBody, LiteLLMParamsBody, ModelsListResponse
 from lifecycle import ResourceManager
+from proxy_client import ProxyClient
 
 LOAD_MODEL_PARAMS = LiteLLMParamsBody(
     model="openai/load-mock",
@@ -19,14 +19,14 @@ LOAD_MODEL_PARAMS = LiteLLMParamsBody(
 
 
 @pytest.fixture(scope="session")
-def client() -> LoadClient:
-    return build_client()
+def client(proxy: ProxyClient) -> LoadClient:
+    return build_client(proxy)
 
 
-def _model_is_servable(gateway: ProxyClient, model_name: str) -> bool:
-    result = gateway.transport.get(
+def _model_is_servable(proxy: ProxyClient, model_name: str) -> bool:
+    result = proxy.transport.get(
         "/v1/models",
-        headers=gateway.transport.master,
+        headers=proxy.transport.master,
         params=NoBody(),
         response_type=ModelsListResponse,
     )
@@ -37,15 +37,15 @@ def _model_is_servable(gateway: ProxyClient, model_name: str) -> bool:
 def _ensure_load_model(  # pyright: ignore[reportUnusedFunction]  # pytest autouse session fixture, wired by name
     client: LoadClient,
 ) -> Iterator[None]:
-    gateway = client.gateway
-    if _model_is_servable(gateway, LOAD_MODEL):
+    proxy = client.proxy
+    if _model_is_servable(proxy, LOAD_MODEL):
         yield
         return
 
     try:
-        model_id = gateway.create_model(LOAD_MODEL, LOAD_MODEL_PARAMS)
+        model_id = proxy.create_model(LOAD_MODEL, LOAD_MODEL_PARAMS)
     except (AssertionError, RequestException) as exc:
-        if _model_is_servable(gateway, LOAD_MODEL):
+        if _model_is_servable(proxy, LOAD_MODEL):
             yield
             return
         raise AssertionError(
@@ -56,11 +56,11 @@ def _ensure_load_model(  # pyright: ignore[reportUnusedFunction]  # pytest autou
     try:
         yield
     finally:
-        gateway.delete_model(model_id)
+        proxy.delete_model(model_id)
 
 
 @pytest.fixture
 def load_key(resources: ResourceManager, client: LoadClient) -> str:
-    key = client.gateway.generate_key(KeyGenerateBody(models=[LOAD_MODEL], user_id="e2e-load"))
-    resources.defer(lambda: client.gateway.delete_key(key))
+    key = client.proxy.generate_key(KeyGenerateBody(models=[LOAD_MODEL], user_id="e2e-load"))
+    resources.defer(lambda: client.proxy.delete_key(key))
     return key

tests/e2e/load/load_client.py

@@ -2,13 +2,13 @@ from __future__ import annotations
 
 from dataclasses import dataclass
 
-from proxy_client import ProxyClient, build_proxy_client
+from proxy_client import ProxyClient
 
 
 @dataclass(frozen=True, slots=True)
 class LoadClient:
-    gateway: ProxyClient
+    proxy: ProxyClient
 
 
-def build_client() -> LoadClient:
-    return LoadClient(gateway=build_proxy_client())
+def build_client(proxy: ProxyClient) -> LoadClient:
+    return LoadClient(proxy=proxy)

tests/e2e/management/conftest.py

@@ -1,24 +1,15 @@
-"""Management suite fixtures: the client plus a logged-in dashboard page.
+"""Management suite's `client` fixture.
 
-Lifecycle/liveness gate/marker live in the parent conftest. The browser fixtures drive
-the dashboard the proxy serves at /ui, so browser tests exercise exactly what an
-end user sees. playwright is an optional dependency loaded behind importorskip
-inside the fixture, so the API tests in this suite collect and run without it:
-
-    uv pip install playwright && uv run playwright install chromium
+Lifecycle/liveness gate/marker live in the parent conftest. ManagementClient
+holds the shared ProxyClient so `resources` / `scoped_key` clean up keys, teams,
+users, and orgs this suite creates.
 """
 
-from typing import TYPE_CHECKING, Iterator
-
 import pytest
 
-from e2e_config import UI_BASE_URL, UI_PASSWORD, UI_USERNAME
 from management_client import ManagementClient, build_client
 from proxy_client import ProxyClient
 
-if TYPE_CHECKING:
-    from playwright.sync_api import Browser, Page
-
 
 def pytest_configure(config: pytest.Config) -> None:
     config.addinivalue_line(
@@ -30,35 +21,3 @@ def pytest_configure(config: pytest.Config) -> None:
 @pytest.fixture(scope="session")
 def client(proxy: ProxyClient) -> ManagementClient:
     return build_client(proxy)
-
-
-@pytest.fixture(scope="session")
-def browser() -> "Iterator[Browser]":
-    pytest.importorskip("playwright.sync_api", reason="playwright not installed")
-    from playwright.sync_api import sync_playwright
-
-    with sync_playwright() as playwright:
-        launched = playwright.chromium.launch()
-        yield launched
-        launched.close()
-
-
-@pytest.fixture
-def ui_page(browser: "Browser") -> "Iterator[Page]":
-    context = browser.new_context()
-    try:
-        page = context.new_page()
-        # Split deploys serve the Next.js dashboard on the UI service, not the
-        # data-plane gateway (which 404s /ui). Login is a client-rendered form
-        # that appears after LoadingScreen; wait on the placeholder, not #id
-        # (Ant Design Input does not always set id="username").
-        page.goto(f"{UI_BASE_URL}/ui/login")
-        username = page.get_by_placeholder("Enter your username")
-        username.wait_for(state="visible", timeout=30_000)
-        username.fill(UI_USERNAME)
-        page.get_by_placeholder("Enter your password").fill(UI_PASSWORD)
-        page.get_by_role("button", name="Login", exact=True).click()
-        page.wait_for_function("() => document.cookie.includes('token=')")
-        yield page
-    finally:
-        context.close()

tests/e2e/management/test_key_models_dropdown_e2e.py

@@ -1,183 +0,0 @@
-"""The dashboard's key create/edit Models dropdown scopes its options to the key's team.
-
-A teamless key offers All Proxy Models but not the all-team-models sentinel (the
-backend expands the latter to the full proxy model list when no team is attached),
-and a team key offers all-team-models plus the team's own models but never the
-all-proxy-models sentinel, even when the team's model list carries it. The create
-cases also walk the full product path: submit the modal with the offered sentinel
-and read the persisted key back through /key/info.
-
-The tests drive gpt-5.5, one of the example models prewired in the proxy config in
-tests/e2e/docker-compose.yml; the dropdown wait fails with a pointer there when the
-proxy under test does not serve it.
-"""
-
-import pytest
-
-from e2e_config import UI_BASE_URL, unique_marker
-from lifecycle import ResourceManager
-from management_client import ManagementClient
-from models import KeyGenerateBody, TeamNewBody
-
-pytest.importorskip("playwright.sync_api", reason="playwright not installed")
-
-from playwright.sync_api import Locator, Page, expect  # noqa: E402  # import must follow the importorskip guard above
-
-
-def _form_item(page: Page, label: str) -> Locator:
-    return page.locator(".ant-form-item").filter(has=page.get_by_text(label, exact=True)).first
-
-
-def _open_dropdown(page: Page, label: str) -> Locator:
-    _form_item(page, label).locator(".ant-select-selector").first.click()
-    dropdown = page.locator(".ant-select-dropdown:not(.ant-select-dropdown-hidden)").last
-    expect(dropdown).to_be_visible()
-    return dropdown
-
-
-def _models_dropdown_texts(page: Page, must_contain: str) -> list[str]:
-    dropdown = _open_dropdown(page, "Models")
-    expect(
-        dropdown.locator(".ant-select-item-option-content", has_text=must_contain).first,
-        f"{must_contain!r} never appeared in the Models dropdown; the proxy must serve it "
-        f"(see the model_list in tests/e2e/docker-compose.yml)",
-    ).to_be_visible()
-    return dropdown.locator(".ant-select-item-option-content").all_inner_texts()
-
-
-def _open_create_key_modal(page: Page) -> None:
-    # Avoid /ui/api-keys/?create=true: on stage the SPA auth redirect often
-    # aborts that navigation mid-flight ("interrupted by another navigation").
-    # Land on the list, wait for the shell, then open create via the button.
-    page.goto(f"{UI_BASE_URL}/ui/api-keys/", wait_until="domcontentloaded")
-    create_btn = page.get_by_role("button", name="+ Create New Key")
-    expect(create_btn).to_be_visible(timeout=60_000)
-    create_btn.click()
-    expect(page.locator(".ant-modal").first).to_be_visible(timeout=15_000)
-
-
-def _select_team(page: Page, alias: str) -> None:
-    dropdown = _open_dropdown(page, "Team")
-    dropdown.get_by_text(alias).first.click()
-
-
-def _submit_create_modal(page: Page, sentinel_label: str) -> str:
-    dropdown = page.locator(".ant-select-dropdown:not(.ant-select-dropdown-hidden)").last
-    dropdown.locator(".ant-select-item-option-content", has_text=sentinel_label).first.click()
-    page.keyboard.press("Escape")
-    _form_item(page, "Key Name").locator("input").first.fill(f"e2e-ui-key-{unique_marker()}")
-    page.get_by_role("button", name="Create Key", exact=True).click()
-
-    expect(page.get_by_text("Save your Key")).to_be_visible()
-    key = page.locator(".ant-modal pre").last.inner_text().strip()
-    assert key.startswith("sk-"), f"expected the created key in the success modal, got {key!r}"
-    return key
-
-
-def _open_key_edit_form(page: Page, key_alias: str) -> None:
-    page.goto(f"{UI_BASE_URL}/ui/api-keys/")
-    # The list is async; wait for the provisioned row before opening detail.
-    row = page.locator("tr").filter(has_text=key_alias).first
-    expect(row).to_be_visible(timeout=60_000)
-    # Key Alias is plain text. KeyInfoView opens from the Key ID control in the
-    # same row (mono hash button on the tremor table / IdCell on the newer
-    # DataTable). Prefer that button; fall back to the alias text for layouts
-    # where the Key column itself is the click target.
-    key_id_button = row.locator("button.font-mono").first
-    if key_id_button.count() == 0:
-        key_id_button = row.locator("button").first
-    if key_id_button.count() > 0:
-        key_id_button.click()
-    else:
-        row.get_by_text(key_alias, exact=True).click()
-    page.get_by_role("tab", name="Settings").click()
-    page.get_by_role("button", name="Edit Settings").click()
-    expect(_form_item(page, "Models")).to_be_visible()
-
-
-def _provision_team(client: ManagementClient, resources: ResourceManager, alias: str) -> str:
-    team_id = client.create_team(TeamNewBody(team_alias=alias, models=["all-proxy-models", "gpt-5.5"]))
-    resources.defer(lambda: client.delete_team(team_id))
-    return team_id
-
-
-def _provision_key(
-    client: ManagementClient, resources: ResourceManager, alias: str, team_id: str | None = None
-) -> str:
-    key = client.proxy.generate_key(KeyGenerateBody(key_alias=alias, models=["gpt-5.5"], team_id=team_id))
-    resources.defer(lambda: client.proxy.delete_key(key))
-    return key
-
-
-@pytest.mark.e2e
-class TestKeyModelsDropdownUI:
-    @pytest.mark.covers("mgmt.key.generate.happy_path", exercised_on=[])
-    def test_create_teamless_key_offers_proxy_scope_and_persists(
-        self, ui_page: Page, client: ManagementClient, resources: ResourceManager
-    ) -> None:
-        _open_create_key_modal(ui_page)
-
-        options = _models_dropdown_texts(ui_page, must_contain="gpt-5.5")
-        assert "All Proxy Models" in options, f"teamless create lost 'All Proxy Models': {options}"
-        assert "All Team Models" not in options, f"teamless create offered 'All Team Models': {options}"
-
-        key = _submit_create_modal(ui_page, sentinel_label="All Proxy Models")
-        resources.defer(lambda: client.proxy.delete_key(key))
-
-        info = client.proxy.key_info(key)
-        assert info.models == ["all-proxy-models"], f"persisted models {info.models}"
-        assert info.team_id is None, f"teamless key persisted with team {info.team_id}"
-
-    @pytest.mark.covers("mgmt.key.generate.happy_path", exercised_on=[])
-    def test_create_team_key_offers_team_scope_and_persists(
-        self, ui_page: Page, client: ManagementClient, resources: ResourceManager
-    ) -> None:
-        team_alias = f"e2e-ui-team-{unique_marker()}"
-        team_id = _provision_team(client, resources, team_alias)
-
-        _open_create_key_modal(ui_page)
-        _select_team(ui_page, team_alias)
-
-        options = _models_dropdown_texts(ui_page, must_contain="All Team Models")
-        assert "gpt-5.5" in options, f"team key create lost the team's own model: {options}"
-        assert "All Proxy Models" not in options, f"team key create offered 'All Proxy Models': {options}"
-        assert "all-proxy-models" not in options, f"team key create offered the raw sentinel: {options}"
-
-        key = _submit_create_modal(ui_page, sentinel_label="All Team Models")
-        resources.defer(lambda: client.proxy.delete_key(key))
-
-        info = client.proxy.key_info(key)
-        assert info.models == ["all-team-models"], f"persisted models {info.models}"
-        assert info.team_id == team_id, f"persisted team {info.team_id}, expected {team_id}"
-
-    @pytest.mark.covers("mgmt.key.update.happy_path", exercised_on=[])
-    def test_edit_teamless_key_offers_proxy_scope(
-        self, ui_page: Page, client: ManagementClient, resources: ResourceManager
-    ) -> None:
-        key_alias = f"e2e-ui-teamless-{unique_marker()}"
-        _provision_key(client, resources, key_alias)
-
-        _open_key_edit_form(ui_page, key_alias)
-
-        options = _models_dropdown_texts(ui_page, must_contain="gpt-5.5")
-        assert "All Proxy Models" in options, f"teamless edit lost 'All Proxy Models': {options}"
-        assert "All Team Models" not in options, f"teamless edit offered 'All Team Models': {options}"
-
-    @pytest.mark.covers("mgmt.key.update.happy_path", exercised_on=[])
-    def test_edit_team_key_offers_team_scope_only(
-        self, ui_page: Page, client: ManagementClient, resources: ResourceManager
-    ) -> None:
-        team_alias = f"e2e-ui-team-{unique_marker()}"
-        team_id = _provision_team(client, resources, team_alias)
-        key_alias = f"e2e-ui-teamkey-{unique_marker()}"
-        _provision_key(client, resources, key_alias, team_id=team_id)
-
-        _open_key_edit_form(ui_page, key_alias)
-
-        # Wait on a real team model: All Team Models is rendered immediately while
-        # availableModels is still fetching, so requiring only the sentinel races
-        # the async team-model load and can read an incomplete dropdown.
-        options = _models_dropdown_texts(ui_page, must_contain="gpt-5.5")
-        assert "All Team Models" in options, f"team key edit lost 'All Team Models': {options}"
-        assert "All Proxy Models" not in options, f"team key edit offered 'All Proxy Models': {options}"
-        assert "all-proxy-models" not in options, f"team key edit offered the raw sentinel: {options}"

tests/e2e/mcp/conftest.py

@@ -6,12 +6,48 @@ the shared ProxyClient, so the `resources` fixture tears down whatever this suit
 creates (keys via the ProxyClient, MCP servers via the deferred cleanups).
 """
 
+from __future__ import annotations
+
+import importlib.util
+import sys
+from pathlib import Path
+from typing import Protocol, cast
+
 import pytest
 
 from mcp_client import McpClient, build_client
 from proxy_client import ProxyClient
 
 
+class DdLogsReader(Protocol):
+    def poll_events_for_marker(self, marker: str) -> list[object]: ...
+
+
+class _DdLogsReaderBuilder(Protocol):
+    def __call__(self) -> DdLogsReader: ...
+
+
+def _build_dd_logs_reader() -> DdLogsReader:
+    # Load logging/datadog_reader.py by path so basedpyright does not require a
+    # package layout. Register the module in sys.modules before exec so
+    # dataclasses inside it can resolve cls.__module__ (otherwise Python 3.12
+    # raises AttributeError: 'NoneType' object has no attribute '__dict__').
+    path = Path(__file__).resolve().parent.parent / "logging" / "datadog_reader.py"
+    name = "e2e_logging_datadog_reader"
+    spec = importlib.util.spec_from_file_location(name, path)
+    assert spec is not None and spec.loader is not None
+    module = importlib.util.module_from_spec(spec)
+    sys.modules[name] = module
+    spec.loader.exec_module(module)
+    builder = cast(_DdLogsReaderBuilder, getattr(module, "build_dd_logs_reader"))
+    return builder()
+
+
 @pytest.fixture(scope="session")
 def client(proxy: ProxyClient) -> McpClient:
     return build_client(proxy)
+
+
+@pytest.fixture(scope="session")
+def dd_logs() -> DdLogsReader:
+    return _build_dd_logs_reader()

tests/e2e/mcp/datadog_mcp.py

@@ -0,0 +1,48 @@
+"""Shared helpers for e2e tests that register the real Datadog remote MCP server."""
+
+from __future__ import annotations
+
+import os
+
+from e2e_config import datadog_mcp_url, unique_marker
+from lifecycle import ResourceManager
+from mcp_client import McpClient
+
+SEARCH_LOGS_TOOL = "search_datadog_logs"
+
+
+def _dd_api_key() -> str:
+    return os.environ.get("DD_API_KEY", "").strip()
+
+
+def _dd_app_key() -> str:
+    return os.environ.get("DD_APP_KEY", "").strip()
+
+
+def assert_dd_mcp_creds() -> None:
+    if not _dd_api_key() or not _dd_app_key():
+        import pytest
+
+        pytest.fail(
+            "Datadog MCP e2e requires DD_API_KEY and DD_APP_KEY "
+            "(header auth to mcp.<site>/v1/mcp; on the cluster the secret manager "
+            "injects them, locally tests/e2e/.env)"
+        )
+
+
+def register_datadog_mcp(client: McpClient, resources: ResourceManager) -> str:
+    assert_dd_mcp_creds()
+    name = f"e2e_dd_mcp_{unique_marker()}"
+    server_id = client.register_server(
+        server_name=name,
+        alias=name,
+        url=datadog_mcp_url(toolsets="core"),
+        transport="http",
+        static_headers={
+            "DD-API-KEY": _dd_api_key(),
+            "DD-APPLICATION-KEY": _dd_app_key(),
+        },
+        allowed_tools=[SEARCH_LOGS_TOOL],
+    )
+    resources.defer(lambda: client.delete_server(server_id))
+    return server_id

tests/e2e/mcp/mcp_client.py

@@ -11,6 +11,7 @@ request/response bodies are co-located here because only this suite speaks MCP.
 
 from __future__ import annotations
 
+from collections.abc import Mapping
 from dataclasses import dataclass
 
 from pydantic import BaseModel, ConfigDict, Field, RootModel
@@ -19,6 +20,9 @@ from e2e_http import Headers, NoBody, Result, unwrap
 from models import KeyGenerateBody, ObjectPermission
 from proxy_client import ProxyClient
 
+McpToolArg = str | int | float | bool | list[str] | dict[str, str]
+McpToolArguments = Mapping[str, McpToolArg]
+
 
 class ApiKeyHeaders(Headers):
     x_litellm_api_key: str = Field(serialization_alias="x-litellm-api-key")
@@ -29,6 +33,9 @@ class McpServerNewBody(BaseModel):
     alias: str
     url: str
     transport: str = "http"
+    auth_type: str | None = None
+    static_headers: dict[str, str] | None = None
+    allowed_tools: list[str] | None = None
 
 
 class McpServerNewResponse(BaseModel):
@@ -68,10 +75,19 @@ class McpToolsListResponse(BaseModel):
             if tool.mcp_info is not None and tool.mcp_info.server_id == server_id
         )
 
+    def tool_name_containing(self, server_id: str, needle: str) -> str | None:
+        needle_l = needle.lower()
+        for tool in self.tools:
+            if tool.mcp_info is None or tool.mcp_info.server_id != server_id:
+                continue
+            if needle_l in tool.name.lower() or tool.name.lower().endswith(needle_l):
+                return tool.name
+        return None
+
 
 class McpCallToolBody(BaseModel):
     name: str
-    arguments: dict[str, int]
+    arguments: dict[str, McpToolArg]
     server_id: str
 
 
@@ -89,17 +105,39 @@ class McpCallToolResponse(BaseModel):
     def first_text(self) -> str | None:
         return self.content[0].text if self.content else None
 
+    @property
+    def all_text(self) -> str:
+        return "\n".join(part.text for part in self.content if part.text)
+
 
 @dataclass(frozen=True, slots=True)
 class McpClient:
     proxy: ProxyClient
 
-    def register_server(self, *, server_name: str, alias: str, url: str) -> str:
+    def register_server(
+        self,
+        *,
+        server_name: str,
+        alias: str,
+        url: str,
+        transport: str = "http",
+        auth_type: str | None = None,
+        static_headers: dict[str, str] | None = None,
+        allowed_tools: list[str] | None = None,
+    ) -> str:
         return unwrap(
             self.proxy.transport.post(
                 "/v1/mcp/server",
                 headers=self.proxy.transport.master,
-                json=McpServerNewBody(server_name=server_name, alias=alias, url=url),
+                json=McpServerNewBody(
+                    server_name=server_name,
+                    alias=alias,
+                    url=url,
+                    transport=transport,
+                    auth_type=auth_type,
+                    static_headers=static_headers,
+                    allowed_tools=allowed_tools,
+                ),
                 response_type=McpServerNewResponse,
             )
         ).server_id
@@ -122,12 +160,22 @@ class McpClient:
             )
         ).root
 
-    def generate_key(self, *, user_id: str, mcp_servers: list[str] | None) -> str:
+    def generate_key(
+        self,
+        *,
+        user_id: str,
+        mcp_servers: list[str] | None,
+        models: list[str] | None = None,
+    ) -> str:
         object_permission = (
             ObjectPermission(mcp_servers=mcp_servers) if mcp_servers is not None else None
         )
         return self.proxy.generate_key(
-            KeyGenerateBody(models=[], user_id=user_id, object_permission=object_permission)
+            KeyGenerateBody(
+                models=models if models is not None else [],
+                user_id=user_id,
+                object_permission=object_permission,
+            )
         )
 
     def list_tools(self, key: str) -> Result[McpToolsListResponse]:
@@ -139,12 +187,19 @@ class McpClient:
         )
 
     def call_tool(
-        self, key: str, *, server_id: str, name: str, arguments: dict[str, int]
+        self,
+        key: str,
+        *,
+        server_id: str,
+        name: str,
+        arguments: McpToolArguments,
     ) -> Result[McpCallToolResponse]:
         return self.proxy.transport.post(
             "/mcp-rest/tools/call",
             headers=ApiKeyHeaders(x_litellm_api_key=key),
-            json=McpCallToolBody(name=name, arguments=arguments, server_id=server_id),
+            json=McpCallToolBody(
+                name=name, arguments=dict(arguments), server_id=server_id
+            ),
             response_type=McpCallToolResponse,
         )
 

tests/e2e/mcp/test_mcp_datadog_e2e.py

@@ -0,0 +1,108 @@
+"""Live e2e: the proxy brokers the real Datadog remote MCP server.
+
+Seeds a chat completion whose prompt carries a unique `e2e-datadog-mcp-*`
+marker so the proxy's DataDogLogger ships a StandardLoggingPayload the org can
+search. Registers the regional Datadog MCP endpoint with DD_API_KEY /
+DD_APP_KEY as static headers (Datadog's documented CI/header auth). A key
+granted that server lists tools, calls search_datadog_logs for the marker, and
+the response must contain it. The dual read via datadog_reader proves the log
+is also in the Logs Search API. The MCP server row is deleted on teardown.
+"""
+
+from __future__ import annotations
+
+import pytest
+
+from conftest import DdLogsReader
+from datadog_mcp import SEARCH_LOGS_TOOL, assert_dd_mcp_creds, register_datadog_mcp
+from e2e_config import CHEAP_ANTHROPIC_MODEL, DD_SEARCH_FROM, unique_marker
+from e2e_http import NoBody, unwrap
+from lifecycle import ResourceManager
+from mcp_client import McpClient
+from models import ChatBody, ChatMessage
+from proxy_client import ProxyClient
+
+pytestmark = pytest.mark.e2e
+
+DD_LOGGER_NAME = "DataDogLogger"
+MARKER_PREFIX = "e2e-datadog-mcp-"
+
+
+def _assert_datadog_logger_active(proxy: ProxyClient) -> None:
+    result = proxy.probe("/health/readiness/details", params=NoBody())
+    assert result.status_code == 200, (
+        f"/health/readiness/details must answer 200, got {result.status_code}: {result.body[:300]}"
+    )
+    assert DD_LOGGER_NAME in result.body, (
+        f"the proxy must report the {DD_LOGGER_NAME} callback active "
+        f"(callbacks + DD_* env); got: {result.body[:400]}"
+    )
+
+
+def _seed_completion(proxy: ProxyClient, *, key: str, marker: str) -> None:
+    body = ChatBody(
+        model=CHEAP_ANTHROPIC_MODEL,
+        messages=[ChatMessage(role="user", content=f"reply with one word {marker}")],
+        max_tokens=16,
+    )
+    unwrap(proxy.chat(key, body))
+
+
+class TestDatadogMcpRoundTrip:
+    @pytest.mark.covers("mcp.list_tools.api_key.succeeds", "mcp.call_tool.api_key.succeeds")
+    def test_search_logs_finds_seeded_completion(
+        self,
+        client: McpClient,
+        dd_logs: DdLogsReader,
+        resources: ResourceManager,
+    ) -> None:
+        assert_dd_mcp_creds()
+        _assert_datadog_logger_active(client.proxy)
+
+        server_id = register_datadog_mcp(client, resources)
+        marker = f"{MARKER_PREFIX}{unique_marker()}"
+
+        key = client.generate_key(
+            user_id=f"e2e-dd-mcp-{unique_marker()}",
+            mcp_servers=[server_id],
+            models=[CHEAP_ANTHROPIC_MODEL],
+        )
+        resources.defer(lambda: client.proxy.delete_key(key))
+
+        _seed_completion(client.proxy, key=key, marker=marker)
+
+        shipped = dd_logs.poll_events_for_marker(marker)
+        assert shipped, (
+            f"proxy DataDogLogger never shipped a log containing {marker!r} "
+            "within the poll deadline; MCP search would have nothing to find"
+        )
+
+        tools = unwrap(client.list_tools(key))
+        tool_name = tools.tool_name_containing(server_id, SEARCH_LOGS_TOOL)
+        assert tool_name is not None, (
+            f"granted key never saw {SEARCH_LOGS_TOOL} on server {server_id}; "
+            f"tools={tools.tool_names_for_server(server_id)}"
+        )
+
+        call = unwrap(
+            client.call_tool(
+                key,
+                server_id=server_id,
+                name=tool_name,
+                arguments={
+                    "query": marker,
+                    "from": DD_SEARCH_FROM,
+                    "to": "now",
+                    "max_tokens": 5000,
+                    "telemetry": {
+                        "intent": "e2e assert seeded litellm completion log is searchable via MCP"
+                    },
+                },
+            )
+        )
+        assert call.is_error is not True, f"search_datadog_logs errored: {call}"
+        body = call.all_text
+        assert marker in body, (
+            f"search_datadog_logs response must include the seeded marker {marker!r}; "
+            f"got: {body[:800]!r}"
+        )

tests/e2e/mcp/test_mcp_key_access_e2e.py

@@ -1,41 +1,27 @@
-"""Live e2e: a virtual key without MCP access is denied an MCP server's tools.
+"""Live e2e: a virtual key without MCP access is denied a real MCP server's tools.
 
-An admin registers an upstream MCP server through the management API (persisted in
-the DB, picked up without a restart) and queues its deletion. Two keys are created
-against that one server: one granted access through `object_permission.mcp_servers`
-and one with no MCP grant at all. The permitted key is the control that proves the
-upstream is alive and the tool is callable, so a failure on the denied key is an
-authorization denial rather than a dead server. The denied key must then see none
-of the server's tools on `tools/list` and must be refused with a 403 on
-`tools/call`.
-
-Both the recorded state (the server is registered; the permitted key resolves its
-tools) and the enforced behavior (the unpermitted key sees nothing and is blocked)
-are asserted, so a regression that leaks tools to an ungranted key or drops the
-call-time permission check fails here.
+An admin registers the Datadog remote MCP server through the management API
+(persisted in the DB, picked up without a restart) and queues its deletion. Two
+keys are created against that one server: one granted access through
+`object_permission.mcp_servers` and one with no MCP grant at all. The permitted
+key is the control that proves the upstream is alive and the tool is callable,
+so a failure on the denied key is an authorization denial rather than a dead
+server. The denied key must then see none of the server's tools on `tools/list`
+and must be refused with a 403 on `tools/call`.
 """
 
-import os
+from __future__ import annotations
 
 import pytest
 
-from e2e_config import unique_marker
+from datadog_mcp import SEARCH_LOGS_TOOL, register_datadog_mcp
+from e2e_config import DD_SEARCH_FROM, unique_marker
 from e2e_http import UnknownApiError, unwrap
 from lifecycle import ResourceManager
 from mcp_client import McpClient
 
 pytestmark = pytest.mark.e2e
 
-MCP_UPSTREAM_URL = os.environ.get("E2E_MCP_UPSTREAM_URL", "http://mcp-upstream:8090/mcp")
-MATH_TOOLS = frozenset({"add", "multiply"})
-
-
-def _register_math_server(client: McpClient, resources: ResourceManager) -> str:
-    name = f"e2e_math_{unique_marker()}"
-    server_id = client.register_server(server_name=name, alias=name, url=MCP_UPSTREAM_URL)
-    resources.defer(lambda: client.delete_server(server_id))
-    return server_id
-
 
 def _key(client: McpClient, resources: ResourceManager, *, mcp_servers: list[str] | None) -> str:
     label = "allowed" if mcp_servers else "denied"
@@ -52,18 +38,21 @@ def _assert_registered(client: McpClient, server_id: str) -> None:
 class TestMcpKeyWithoutAccessIsDenied:
     @pytest.mark.covers("mcp.list_tools.api_key.denied_without_permission")
     def test_list_tools_denied_without_permission(
-        self, client: McpClient, resources: ResourceManager
+        self,
+        client: McpClient,
+        resources: ResourceManager,
     ) -> None:
-        server_id = _register_math_server(client, resources)
+        server_id = register_datadog_mcp(client, resources)
         _assert_registered(client, server_id)
 
         permitted_key = _key(client, resources, mcp_servers=[server_id])
         denied_key = _key(client, resources, mcp_servers=None)
 
-        permitted_tools = unwrap(client.list_tools(permitted_key)).tool_names_for_server(server_id)
-        assert MATH_TOOLS <= permitted_tools, (
-            f"granted key did not see the server's tools (upstream dead or grant not applied): "
-            f"{permitted_tools}"
+        permitted = unwrap(client.list_tools(permitted_key))
+        tool_name = permitted.tool_name_containing(server_id, SEARCH_LOGS_TOOL)
+        assert tool_name is not None, (
+            f"granted key did not see {SEARCH_LOGS_TOOL} (upstream dead or grant not applied): "
+            f"{permitted.tool_names_for_server(server_id)}"
         )
 
         denied_tools = unwrap(client.list_tools(denied_key)).tool_names_for_server(server_id)
@@ -74,29 +63,36 @@ class TestMcpKeyWithoutAccessIsDenied:
 
     @pytest.mark.covers("mcp.call_tool.api_key.denied_without_permission")
     def test_call_tool_denied_without_permission(
-        self, client: McpClient, resources: ResourceManager
+        self,
+        client: McpClient,
+        resources: ResourceManager,
     ) -> None:
-        server_id = _register_math_server(client, resources)
+        server_id = register_datadog_mcp(client, resources)
         _assert_registered(client, server_id)
 
         permitted_key = _key(client, resources, mcp_servers=[server_id])
         denied_key = _key(client, resources, mcp_servers=None)
 
-        permitted_tools = unwrap(client.list_tools(permitted_key)).tool_names_for_server(server_id)
-        assert "add" in permitted_tools, (
-            f"granted key did not discover the add tool (upstream dead or grant not applied): "
-            f"{permitted_tools}"
+        permitted = unwrap(client.list_tools(permitted_key))
+        tool_name = permitted.tool_name_containing(server_id, SEARCH_LOGS_TOOL)
+        assert tool_name is not None, (
+            f"granted key did not discover {SEARCH_LOGS_TOOL} (upstream dead or grant not applied): "
+            f"{permitted.tool_names_for_server(server_id)}"
         )
 
+        search_args = {
+            "query": "service:litellm",
+            "from": DD_SEARCH_FROM,
+            "to": "now",
+            "max_tokens": 1000,
+            "telemetry": {"intent": "e2e control call proving granted key can invoke Datadog MCP"},
+        }
         permitted_call = unwrap(
-            client.call_tool(permitted_key, server_id=server_id, name="add", arguments={"a": 3, "b": 4})
+            client.call_tool(permitted_key, server_id=server_id, name=tool_name, arguments=search_args)
         )
         assert permitted_call.is_error is not True, f"granted key's tool call errored: {permitted_call}"
-        assert permitted_call.first_text == "7", (
-            f"granted key's add(3, 4) did not return 7 (upstream not reachable): {permitted_call}"
-        )
 
-        match client.call_tool(denied_key, server_id=server_id, name="add", arguments={"a": 3, "b": 4}):
+        match client.call_tool(denied_key, server_id=server_id, name=tool_name, arguments=search_args):
             case UnknownApiError(status_code=403, body=body):
                 assert "access_denied" in body, f"403 was not an MCP access denial: {body}"
             case other:

tests/e2e/models.py

@@ -155,17 +155,6 @@ class ChatBody(BaseModel):
     guardrails: list[str] | None = None
 
 
-class AnthropicMessagesBody(BaseModel):
-    model: str
-    messages: list[ChatMessage]
-    max_tokens: int
-    stream: bool | None = None
-
-
-class AnthropicMessagesResponse(BaseModel):
-    model: str | None = None
-
-
 class OutMessage(BaseModel):
     content: str | None = None
     reasoning_content: str | None = None
@@ -196,6 +185,82 @@ class ChatResponse(BaseModel):
     service_tier: str | None = None
 
 
+# ---------- anthropic /v1/messages + count_tokens ----------
+
+
+class JsonSchemaProperty(BaseModel):
+    """One property in a tool's JSON-Schema `input_schema`. Only `type` is
+    modelled; the endpoints under test read no further into the schema."""
+
+    type: str
+
+
+class ToolInputSchema(BaseModel):
+    type: str = "object"
+    properties: dict[str, JsonSchemaProperty] = {}
+    required: list[str] = []
+
+
+class AnthropicToolSearchTool(BaseModel):
+    """The tool_search discovery tool. `type` carries the SDK-version-pinned
+    suffix (e.g. ``tool_search_tool_regex_20251119``) that LiteLLM keys its
+    per-provider beta-header translation on; `name` is the unsuffixed
+    canonical name the upstream accepts."""
+
+    type: str
+    name: str
+
+
+class AnthropicCustomTool(BaseModel):
+    name: str
+    description: str
+    input_schema: ToolInputSchema
+
+
+type AnthropicTool = AnthropicToolSearchTool | AnthropicCustomTool
+
+
+class AnthropicMessagesBody(BaseModel):
+    model: str
+    messages: list[ChatMessage]
+    max_tokens: int
+    stream: bool | None = None
+    tools: list[AnthropicTool] | None = None
+
+
+class CountTokensBody(BaseModel):
+    """POST /v1/messages/count_tokens body: the /v1/messages shape minus
+    max_tokens (the endpoint only counts the prompt)."""
+
+    model: str
+    messages: list[ChatMessage]
+
+
+class AnthropicContentBlock(BaseModel):
+    type: str | None = None
+
+
+class AnthropicMessagesResponse(BaseModel):
+    """A /v1/messages answer. `content` is the Anthropic-native passthrough
+    shape; `choices` is the OpenAI-normalized shape LiteLLM emits for some
+    providers (e.g. Bedrock Converse). Presence of either proves the proxy
+    accepted and round-tripped the request. `extra="allow"` keeps the other
+    top-level keys so a shape-check failure can report the actual response keys
+    for triage."""
+
+    model_config = ConfigDict(extra="allow")
+    model: str | None = None
+    content: list[AnthropicContentBlock] | None = None
+    choices: list[ChatChoice] | None = None
+
+
+class CountTokensResponse(BaseModel):
+    """`/v1/messages/count_tokens` answer. `input_tokens` is required so a 200
+    whose body lacks it fails validation instead of passing vacuously."""
+
+    input_tokens: int
+
+
 class EmbedBody(BaseModel):
     model: str
     input: str

tests/e2e/proxy_client.py

@@ -15,6 +15,7 @@ from dataclasses import dataclass
 from datetime import datetime
 
 from e2e_http import (
+    AnthropicHeaders,
     NoBody,
     ProbeResult,
     Result,
@@ -24,8 +25,12 @@ from e2e_http import (
     unwrap,
 )
 from models import (
+    AnthropicMessagesBody,
+    AnthropicMessagesResponse,
     ChatBody,
     ChatResponse,
+    CountTokensBody,
+    CountTokensResponse,
     CustomerDeleteBody,
     EmbedBody,
     EmbedResponse,
@@ -243,6 +248,31 @@ class ProxyClient:
             response_type=OcrResponse,
         )
 
+    def count_tokens(self, key: str, body: CountTokensBody) -> Result[CountTokensResponse]:
+        """POST /v1/messages/count_tokens (Anthropic-native). Sends the
+        anthropic-version header so the native path accepts it; harmless on the
+        other providers the proxy fronts."""
+        return self.transport.post(
+            "/v1/messages/count_tokens",
+            headers=self._anthropic_headers(key),
+            json=body,
+            response_type=CountTokensResponse,
+        )
+
+    def messages(self, key: str, body: AnthropicMessagesBody) -> Result[AnthropicMessagesResponse]:
+        """POST /v1/messages (Anthropic-native). The response is either the
+        Anthropic-shape passthrough (`content`) or the OpenAI-normalized shape
+        (`choices`); AnthropicMessagesResponse models both."""
+        return self.transport.post(
+            "/v1/messages",
+            headers=self._anthropic_headers(key),
+            json=body,
+            response_type=AnthropicMessagesResponse,
+        )
+
+    def _anthropic_headers(self, key: str) -> AnthropicHeaders:
+        return AnthropicHeaders(authorization=self.transport.bearer(key).authorization)
+
     # ---- spend read-back ------------------------------------------------
 
     def spend_logs(self, params: SpendLogsParams) -> list[SpendLogRow]:

tests/test_litellm/proxy/vector_store_endpoints/test_vector_store_endpoints.py

@@ -210,7 +210,7 @@ async def test_vector_store_file_list_resolves_single_openai_team_deployment():
     assert result["model"] == "openai/gpt-4o-mini"
     assert "custom_llm_provider" not in result
     llm_router.get_deployment_credentials_with_provider.assert_called_once_with(
-        model_id="team-openai"
+        model_id="team-openai", team_id=None
     )
 
 

tests/test_litellm/test_router.py

@@ -3714,6 +3714,125 @@ def test_get_deployment_credentials_with_provider_resolves_credential_name():
     litellm.credential_list = []
 
 
+def _team_wildcard_model(api_key: str, model_id: str = "team-wildcard-id") -> dict:
+    return {
+        "model_name": f"model_name_team-1_{model_id}",
+        "litellm_params": {"model": "openai/*", "api_key": api_key},
+        "model_info": {
+            "id": model_id,
+            "team_id": "team-1",
+            "team_public_model_name": "openai/*",
+        },
+    }
+
+
+def test_get_deployment_credentials_with_provider_team_wildcard_priority():
+    """
+    Regression: a global wildcard pattern (e.g. "openai/*") must not shadow a
+    team's own wildcard entry. When team_id is provided, the team wildcard
+    deployment's credentials win; without team_id the global one is used.
+    """
+    router = litellm.Router(
+        model_list=[
+            {
+                "model_name": "openai/*",
+                "litellm_params": {"model": "openai/*", "api_key": "global-key"},
+            },
+            _team_wildcard_model(api_key="team-key"),
+        ],
+    )
+
+    team_credentials = router.get_deployment_credentials_with_provider(
+        model_id="openai/gpt-5.2", team_id="team-1"
+    )
+    assert team_credentials is not None
+    assert team_credentials["api_key"] == "team-key"
+
+    global_credentials = router.get_deployment_credentials_with_provider(
+        model_id="openai/gpt-5.2"
+    )
+    assert global_credentials is not None
+    assert global_credentials["api_key"] == "global-key"
+
+
+def test_team_wildcard_credentials_not_usable_after_delete_deployment():
+    """
+    Regression: team_pattern_routers retained deleted deployments, so a team
+    user could keep resolving credentials of a deleted wildcard deployment.
+    """
+    router = litellm.Router(model_list=[_team_wildcard_model(api_key="old-key")])
+
+    assert (
+        router.get_deployment_credentials_with_provider(
+            model_id="openai/gpt-5.2", team_id="team-1"
+        )
+        is not None
+    )
+
+    router.delete_deployment(id="team-wildcard-id")
+
+    assert (
+        router.get_deployment_credentials_with_provider(
+            model_id="openai/gpt-5.2", team_id="team-1"
+        )
+        is None
+    )
+
+
+def test_pattern_match_router_remove_deployment():
+    """
+    remove_deployment must drop only the deployment with the given model id and
+    delete patterns whose deployment list becomes empty.
+    """
+    from litellm.router_utils.pattern_match_deployments import PatternMatchRouter
+
+    pattern_router = PatternMatchRouter()
+    pattern_router.add_pattern(
+        "openai/*",
+        {"litellm_params": {"model": "openai/*", "api_key": "key-a"}, "model_info": {"id": "dep-a"}},
+    )
+    pattern_router.add_pattern(
+        "openai/*",
+        {"litellm_params": {"model": "openai/*", "api_key": "key-b"}, "model_info": {"id": "dep-b"}},
+    )
+
+    pattern_router.remove_deployment(model_id="dep-a")
+    matches = pattern_router.route("openai/gpt-5.2")
+    assert matches is not None
+    assert [m["model_info"]["id"] for m in matches] == ["dep-b"]
+
+    pattern_router.remove_deployment(model_id="dep-b")
+    assert pattern_router.patterns == {}
+    assert pattern_router.route("openai/gpt-5.2") is None
+
+
+def test_team_wildcard_credentials_refreshed_on_upsert_and_set_model_list():
+    """
+    Regression: replacing a team wildcard deployment (upsert or model list
+    reload) must serve the new credentials, not the stale cached ones.
+    """
+    from litellm.types.router import Deployment
+
+    router = litellm.Router(model_list=[_team_wildcard_model(api_key="old-key")])
+
+    router.upsert_deployment(
+        deployment=Deployment(**_team_wildcard_model(api_key="new-key"))
+    )
+    credentials = router.get_deployment_credentials_with_provider(
+        model_id="openai/gpt-5.2", team_id="team-1"
+    )
+    assert credentials is not None
+    assert credentials["api_key"] == "new-key"
+
+    router.set_model_list(model_list=[])
+    assert (
+        router.get_deployment_credentials_with_provider(
+            model_id="openai/gpt-5.2", team_id="team-1"
+        )
+        is None
+    )
+
+
 def test_get_available_guardrail_single_deployment():
     """
     Test get_available_guardrail returns the single guardrail when only one exists.