scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-08-28 05:25:59 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 36

refactor(e2e): fold claude_code HTTP probes onto shared Gateway methods (#33760)

Browse source 
* refactor(e2e): fold claude_code HTTP probes onto shared Gateway methods

Migrate tests/e2e/claude_code/http_probe.py off its own httpx client onto the
shared transport, and promote count_tokens and native anthropic messages to
first-class Gateway methods (Gateway.count_tokens / Gateway.messages) with typed
request/response models in the shared models.py so other suites reuse them.

The probes now take an injected Gateway and issue their request through the
shared count_tokens/messages methods, reusing the split control/data-plane
routing, timeout, and typed Result handling the rest of tests/e2e uses. The wire
shape is preserved: the pydantic bodies serialize byte-for-byte to what the old
httpx probes sent, and the anthropic-version header is carried by a small
AnthropicHeaders model. httpx is gone from the module.

* test(e2e): drop unit-level probe harness test

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

---------

Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
Yassin Kortam 2026-07-18 12:03:01 -07:00 • committed by GitHub
parent e8aef29d0c
commit 0439bcbfed
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
16 changed files with 334 additions and 265 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

2
tests/e2e/CLAUDE.md
View file

@ -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

29
tests/e2e/claude_code/_env.py
View file

@ -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)

6
tests/e2e/claude_code/count_tokens/test_anthropic.py
View file

@ -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:

6
tests/e2e/claude_code/count_tokens/test_azure.py
View file

@ -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:

6
tests/e2e/claude_code/count_tokens/test_bedrock_converse.py
View file

@ -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:

6
tests/e2e/claude_code/count_tokens/test_bedrock_invoke.py
View file

@ -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:

6
tests/e2e/claude_code/count_tokens/test_vertex_ai.py
View file

@ -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:

382
tests/e2e/claude_code/http_probe.py
View file

@ -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")

6
tests/e2e/claude_code/tool_search/test_anthropic.py
View file

@ -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}"

6
tests/e2e/claude_code/tool_search/test_azure.py
View file

@ -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}"

6
tests/e2e/claude_code/tool_search/test_bedrock_converse.py
View file

@ -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}"

6
tests/e2e/claude_code/tool_search/test_bedrock_invoke.py
View file

@ -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}"

6
tests/e2e/claude_code/tool_search/test_vertex_ai.py
View file

@ -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}"

9
tests/e2e/e2e_http.py
View file

@ -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."""

87
tests/e2e/models.py
View file

@ -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

30
tests/e2e/proxy_client.py
View file

@ -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]: