diff --git a/tests/e2e/CLAUDE.md b/tests/e2e/CLAUDE.md index 58a330775c0..481893ed714 100644 --- a/tests/e2e/CLAUDE.md +++ b/tests/e2e/CLAUDE.md @@ -19,7 +19,7 @@ Each subdirectory under `tests/e2e/` is one suite, scoped to an endpoint family - `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 ## Lay the pattern down in a class diff --git a/tests/e2e/claude_code/_env.py b/tests/e2e/claude_code/_env.py index 4f93cb57fdc..ed2da1fe1e2 100644 --- a/tests/e2e/claude_code/_env.py +++ b/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) diff --git a/tests/e2e/claude_code/count_tokens/test_anthropic.py b/tests/e2e/claude_code/count_tokens/test_anthropic.py index 2fbdd4212c4..05110d24e86 100644 --- a/tests/e2e/claude_code/count_tokens/test_anthropic.py +++ b/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: diff --git a/tests/e2e/claude_code/count_tokens/test_azure.py b/tests/e2e/claude_code/count_tokens/test_azure.py index a9aa168ccea..c60c623ae89 100644 --- a/tests/e2e/claude_code/count_tokens/test_azure.py +++ b/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: diff --git a/tests/e2e/claude_code/count_tokens/test_bedrock_converse.py b/tests/e2e/claude_code/count_tokens/test_bedrock_converse.py index 6dcff3ecae3..0cb4766bb31 100644 --- a/tests/e2e/claude_code/count_tokens/test_bedrock_converse.py +++ b/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: diff --git a/tests/e2e/claude_code/count_tokens/test_bedrock_invoke.py b/tests/e2e/claude_code/count_tokens/test_bedrock_invoke.py index ae89067dc00..f1389574527 100644 --- a/tests/e2e/claude_code/count_tokens/test_bedrock_invoke.py +++ b/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: diff --git a/tests/e2e/claude_code/count_tokens/test_vertex_ai.py b/tests/e2e/claude_code/count_tokens/test_vertex_ai.py index 2bf75063590..0894214d4f0 100644 --- a/tests/e2e/claude_code/count_tokens/test_vertex_ai.py +++ b/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: diff --git a/tests/e2e/claude_code/http_probe.py b/tests/e2e/claude_code/http_probe.py index d95307db3b9..c77020acd6e 100644 --- a/tests/e2e/claude_code/http_probe.py +++ b/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") diff --git a/tests/e2e/claude_code/tool_search/test_anthropic.py b/tests/e2e/claude_code/tool_search/test_anthropic.py index 7b8ea07aa07..a23d1b3d2bb 100644 --- a/tests/e2e/claude_code/tool_search/test_anthropic.py +++ b/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}" diff --git a/tests/e2e/claude_code/tool_search/test_azure.py b/tests/e2e/claude_code/tool_search/test_azure.py index 4353a73be90..b094a35ea63 100644 --- a/tests/e2e/claude_code/tool_search/test_azure.py +++ b/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}" diff --git a/tests/e2e/claude_code/tool_search/test_bedrock_converse.py b/tests/e2e/claude_code/tool_search/test_bedrock_converse.py index 7951f8ecdb4..f395122a5ab 100644 --- a/tests/e2e/claude_code/tool_search/test_bedrock_converse.py +++ b/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}" diff --git a/tests/e2e/claude_code/tool_search/test_bedrock_invoke.py b/tests/e2e/claude_code/tool_search/test_bedrock_invoke.py index f01dc3e84f1..12f8909e3e8 100644 --- a/tests/e2e/claude_code/tool_search/test_bedrock_invoke.py +++ b/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}" diff --git a/tests/e2e/claude_code/tool_search/test_vertex_ai.py b/tests/e2e/claude_code/tool_search/test_vertex_ai.py index 00487797221..7d0d35b1c1d 100644 --- a/tests/e2e/claude_code/tool_search/test_vertex_ai.py +++ b/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}" diff --git a/tests/e2e/e2e_http.py b/tests/e2e/e2e_http.py index ff296969079..009529e09db 100644 --- a/tests/e2e/e2e_http.py +++ b/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.""" diff --git a/tests/e2e/models.py b/tests/e2e/models.py index daf85b7fc74..8d19e2f8965 100644 --- a/tests/e2e/models.py +++ b/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 diff --git a/tests/e2e/proxy_client.py b/tests/e2e/proxy_client.py index 0dbb80990c0..c466d415d0e 100644 --- a/tests/e2e/proxy_client.py +++ b/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]: