From 2187be08b59261c30f5176ab1975ae31ebe5ddb0 Mon Sep 17 00:00:00 2001 From: yassin Date: Fri, 17 Jul 2026 17:24:06 +0000 Subject: [PATCH] refactor(e2e): fold claude_code HTTP probes onto shared transport Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- tests/e2e/CLAUDE.md | 2 +- tests/e2e/claude_code/_env.py | 28 ++ .../count_tokens/test_anthropic.py | 8 +- .../claude_code/count_tokens/test_azure.py | 8 +- .../count_tokens/test_bedrock_converse.py | 8 +- .../count_tokens/test_bedrock_invoke.py | 8 +- .../count_tokens/test_vertex_ai.py | 8 +- tests/e2e/claude_code/http_probe.py | 436 ++++++++---------- .../claude_code/tool_search/test_anthropic.py | 6 +- .../e2e/claude_code/tool_search/test_azure.py | 6 +- .../tool_search/test_bedrock_converse.py | 6 +- .../tool_search/test_bedrock_invoke.py | 6 +- .../claude_code/tool_search/test_vertex_ai.py | 6 +- 13 files changed, 256 insertions(+), 280 deletions(-) diff --git a/tests/e2e/CLAUDE.md b/tests/e2e/CLAUDE.md index 0e1eafb5196..b4115734e82 100644 --- a/tests/e2e/CLAUDE.md +++ b/tests/e2e/CLAUDE.md @@ -17,7 +17,7 @@ Each subdirectory under `tests/e2e/` is one suite, scoped to an endpoint family - `security/` - secret handling and log-leak protection - `router/` - routing and reliability behavior (fallbacks, cooldowns) - `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. Its HTTP-probe cells (`count_tokens`, `tool_search`) ride the shared transport through an injected `Gateway`; the CLI-driven cells still run through its own driver ## 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..7bd5ea507af 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 e2e_gateway import Gateway, build_gateway + class ProxyConfig(NamedTuple): base_url: str @@ -72,3 +74,29 @@ def require_proxy( if cfg is None: _fail_missing_proxy_env(compat_result) return cfg + + +def gateway_from(cfg: ProxyConfig) -> Gateway: + """Build the shared e2e Gateway from a resolved proxy config. + + The compat suite talks to a single monolithic proxy, so the data and control + planes share ``cfg.base_url`` and the master key is ``cfg.api_key``. Probes + inject the returned Gateway and issue their requests through its transport, + reusing the shared routing, timeout, and error normalization.""" + return build_gateway( + base_url=cfg.base_url, + master_key=cfg.api_key, + control_plane_base_url=cfg.base_url, + ) + + +def require_gateway( + compat_result, + *, + env: Mapping[str, str] | None = None, +) -> Gateway: + """Return the shared Gateway for the resolved proxy, or hard-fail the test. + + Wraps ``require_proxy`` so HTTP-probe cells can depend on the shared transport + with the same missing-env failure behavior as the CLI-driven cells.""" + return gateway_from(require_proxy(compat_result, env=env)) diff --git a/tests/e2e/claude_code/count_tokens/test_anthropic.py b/tests/e2e/claude_code/count_tokens/test_anthropic.py index 2fbdd4212c4..57da51606ef 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_gateway from claude_code.http_probe import ( assert_count_tokens_shape, probe_count_tokens, @@ -57,13 +57,11 @@ 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) + gateway = require_gateway(compat_result) failures = [] for model in ANTHROPIC_MODELS: - result = probe_count_tokens( - base_url=base_url, api_key=api_key, model=model - ) + result = probe_count_tokens(gateway=gateway, model=model) shape_error = assert_count_tokens_shape(result) if shape_error is not None: error = f"[{model}] count_tokens probe failed: {shape_error}" diff --git a/tests/e2e/claude_code/count_tokens/test_azure.py b/tests/e2e/claude_code/count_tokens/test_azure.py index a9aa168ccea..a87d614add6 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_gateway from claude_code.http_probe import ( assert_count_tokens_shape, probe_count_tokens, @@ -57,13 +57,11 @@ 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) + gateway = require_gateway(compat_result) failures = [] for model in AZURE_MODELS: - result = probe_count_tokens( - base_url=base_url, api_key=api_key, model=model - ) + result = probe_count_tokens(gateway=gateway, model=model) shape_error = assert_count_tokens_shape(result) if shape_error is not None: error = f"[{model}] count_tokens probe failed: {shape_error}" 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..faab4fa705a 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_gateway from claude_code.http_probe import ( assert_count_tokens_shape, probe_count_tokens, @@ -57,13 +57,11 @@ 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) + gateway = require_gateway(compat_result) failures = [] for model in BEDROCK_CONVERSE_MODELS: - result = probe_count_tokens( - base_url=base_url, api_key=api_key, model=model - ) + result = probe_count_tokens(gateway=gateway, model=model) shape_error = assert_count_tokens_shape(result) if shape_error is not None: error = f"[{model}] count_tokens probe failed: {shape_error}" 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..d9ce8aa4381 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_gateway from claude_code.http_probe import ( assert_count_tokens_shape, probe_count_tokens, @@ -57,13 +57,11 @@ 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) + gateway = require_gateway(compat_result) failures = [] for model in BEDROCK_INVOKE_MODELS: - result = probe_count_tokens( - base_url=base_url, api_key=api_key, model=model - ) + result = probe_count_tokens(gateway=gateway, model=model) shape_error = assert_count_tokens_shape(result) if shape_error is not None: error = f"[{model}] count_tokens probe failed: {shape_error}" 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..3c549aab9b1 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_gateway from claude_code.http_probe import ( assert_count_tokens_shape, probe_count_tokens, @@ -58,13 +58,11 @@ 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) + gateway = require_gateway(compat_result) failures = [] for model in VERTEX_AI_MODELS: - result = probe_count_tokens( - base_url=base_url, api_key=api_key, model=model - ) + result = probe_count_tokens(gateway=gateway, model=model) shape_error = assert_count_tokens_shape(result) if shape_error is not None: error = f"[{model}] count_tokens probe failed: {shape_error}" diff --git a/tests/e2e/claude_code/http_probe.py b/tests/e2e/claude_code/http_probe.py index d95307db3b9..e024c24d985 100644 --- a/tests/e2e/claude_code/http_probe.py +++ b/tests/e2e/claude_code/http_probe.py @@ -10,27 +10,27 @@ through the proxy for any given provider. This module is the second test pattern the matrix supports: a plain HTTP POST against a LiteLLM proxy endpoint, parsed and shape-checked -in the test, with the same `compat_result` recording convention as the -CLI-driven cells. The goal is to keep this pattern *narrow* -- if a -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. +in the test. It rides the shared e2e transport rather than a bespoke +HTTP client: each probe takes an injected `Gateway` and issues the +request through `gateway.transport.send`, so the split control/data +plane routing, timeouts, and error normalization are the same ones the +rest of `tests/e2e/` uses. The claude-specific parts that stay are the +request bodies (the `tool_search_tool_regex_20251119` discovery tool +shape) and the per-provider rate-limit accounting the CLI rows share. -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 goal is to keep this pattern *narrow* -- if a 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. """ from __future__ import annotations -import json -from dataclasses import dataclass -from typing import Any, Mapping, Optional +from pydantic import BaseModel, Field, ValidationError -import httpx +from e2e_gateway import Gateway +from e2e_http import Headers, StreamingResponse +from models import ChatMessage from claude_code.rate_limiter import ( RateLimiter, @@ -39,254 +39,212 @@ from claude_code.rate_limiter import ( ) -DEFAULT_TIMEOUT_SECONDS = 30.0 +ANTHROPIC_VERSION = "2023-06-01" -@dataclass -class ProbeResult: - """Structured outcome of a single HTTP probe. +class AnthropicProbeHeaders(Headers): + """Bearer auth plus the `anthropic-version` header Claude Code sends on its + own internal `/v1/messages` and `count_tokens` calls. It is required by + Anthropic's native API and harmless on every other provider the proxy routes + to, so the probe mirrors it for wire fidelity.""" - `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). - """ + authorization: str | None = None + anthropic_version: str = Field(default=ANTHROPIC_VERSION, alias="anthropic-version") - status_code: int - body: str - payload: Optional[Mapping[str, Any]] = None - error: Optional[str] = None + +class CountTokensRequest(BaseModel): + model: str + messages: list[ChatMessage] + + +class ToolSearchDiscoveryTool(BaseModel): + """The tool_search discovery tool itself. `type` is the SDK-version-pinned + `_20251119` suffix LiteLLM keys its beta-header translation on; `name` is the + canonical `tool_search_tool_regex` (no suffix) Anthropic accepts.""" + + type: str = "tool_search_tool_regex_20251119" + name: str = "tool_search_tool_regex" + + +class ToolInputProperty(BaseModel): + type: str + + +class ToolInputSchema(BaseModel): + type: str = "object" + properties: dict[str, ToolInputProperty] + required: list[str] + + +class UserTool(BaseModel): + """A trivial user tool for the discovery tool to potentially surface, so the + wire shape mirrors what real Claude Code sends (>=1 non-search tool).""" + + name: str + description: str + input_schema: ToolInputSchema + + +ProbeTool = ToolSearchDiscoveryTool | UserTool + + +class ToolSearchRequest(BaseModel): + model: str + max_tokens: int + messages: list[ChatMessage] + tools: list[ProbeTool] + + +class CountTokensResult(BaseModel): + input_tokens: int + + +class ContentBlock(BaseModel): + type: str | None = None + + +class ResponseChoice(BaseModel): + index: int | None = None + + +class ToolSearchResponse(BaseModel): + """LiteLLM normalizes some provider responses to OpenAI shape (`choices`) and + passes others through Anthropic-shape (`content`). Either proves the proxy + round-tripped the request; the matrix does not assert the model actually + invoked tool_search.""" + + content: list[ContentBlock] | None = None + choices: list[ResponseChoice] | None = None + + +def _probe_headers(gateway: Gateway) -> AnthropicProbeHeaders: + return AnthropicProbeHeaders(authorization=gateway.transport.master.authorization) + + +def _acquire(model: str, rate_limiter: RateLimiter | None) -> None: + 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, - api_key: str, + gateway: Gateway, 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, +) -> StreamingResponse: + """POST to `/v1/messages/count_tokens` for `model` and return the raw outcome. - 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 Anthropic / LiteLLM `count_tokens` endpoint accepts a request body whose + shape mirrors `/v1/messages` (model + messages) and returns + `{"input_tokens": N}` on success; `assert_count_tokens_shape` checks that. - 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. - """ - 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, + 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. `rate_limiter` is an injection seam for unit tests; production callers + should leave it unset to use the process-wide default.""" + _acquire(model, rate_limiter) + return gateway.transport.send( + "/v1/messages/count_tokens", + headers=_probe_headers(gateway), + json=CountTokensRequest( + model=model, messages=[ChatMessage(role="user", content=message)] + ), ) def probe_tool_search( *, - base_url: str, - api_key: str, + gateway: Gateway, 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, +) -> StreamingResponse: + """POST to `/v1/messages` with a `tool_search_tool_regex_20251119` tool and + return the raw outcome. - 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. + The tools array is the one Claude Code emits when its MCP-tool-search beta is + active: a discovery tool plus at least one regular user tool. LiteLLM keys on + the `_20251119` type string to attach the provider-specific beta header + (`advanced-tool-use-2025-11-20` for Anthropic/Azure, + `tool-search-tool-2025-10-19` for Vertex/Bedrock); a regression in that + translation surfaces as a 400 from the upstream. - 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. - """ - 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'." + The prompt deliberately does not force a tool call -- the goal is to verify + the request round-trips without 400, not that the model chose to invoke + tool_search. Like `probe_count_tokens`, this acquires one rate-limit token.""" + _acquire(model, rate_limiter) + return gateway.transport.send( + "/v1/messages", + headers=_probe_headers(gateway), + json=ToolSearchRequest( + model=model, + max_tokens=64, + messages=[ + ChatMessage( + role="user", + content=( + "If you have a tool to discover other tools, use it to " + "find one. Otherwise reply with the word 'done'." + ), + ) + ], + tools=[ + ToolSearchDiscoveryTool(), + UserTool( + name="add_numbers", + description="Add two integers", + input_schema=ToolInputSchema( + properties={ + "a": ToolInputProperty(type="integer"), + "b": ToolInputProperty(type="integer"), + }, + required=["a", "b"], + ), ), - } - ], - "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, + ], + ), ) -def assert_tool_search_shape(result: ProbeResult) -> Optional[str]: - """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. - """ - 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 - - -def assert_count_tokens_shape(result: ProbeResult) -> Optional[str]: +def assert_count_tokens_shape(result: StreamingResponse) -> 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. - - 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. - """ - 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}" + Acceptance criteria are intentionally minimal: HTTP status 200, a valid JSON + body, and an `input_tokens` key whose value is a positive int. Anything beyond + that (cache token fields, server metadata) varies by provider/transport and + would make the cell flip red on neutral protocol drift.""" + if not result.ok: + return _status_error(result) + try: + parsed = CountTokensResult.model_validate_json(result.body) + except ValidationError as exc: + return f"invalid count_tokens body ({result.body[:200]}): {exc}" + if parsed.input_tokens <= 0: + return f"input_tokens must be positive; got {parsed.input_tokens}" return None + + +def assert_tool_search_shape(result: StreamingResponse) -> str | None: + """Return None on success, else describe the first violation. + + Acceptance: HTTP status 200 (no 400 rejecting the tool_search type or a + missing beta header), a valid JSON body, and either `content` (Anthropic + passthrough) or `choices` (LiteLLM normalized OpenAI shape). The cell goes red + when the upstream rejects the tool type, the proxy drops the beta header, or + the response shape is unusable; a model deciding not to call tool_search is + irrelevant.""" + if not result.ok: + return _status_error(result) + try: + parsed = ToolSearchResponse.model_validate_json(result.body) + except ValidationError as exc: + return f"non-JSON or unparseable body ({result.body[:200]}): {exc}" + if parsed.content is None and parsed.choices is None: + return f"response has neither `content` nor `choices`: {result.body[:400]}" + return None + + +def _status_error(result: StreamingResponse) -> str: + if result.status_code < 0: + return f"transport error: {result.body}" + return f"status {result.status_code}: {result.body[:400]}" diff --git a/tests/e2e/claude_code/tool_search/test_anthropic.py b/tests/e2e/claude_code/tool_search/test_anthropic.py index 7b8ea07aa07..06fbd47b391 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_gateway 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) + gateway = require_gateway(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(gateway=gateway, 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..76e29f12452 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_gateway 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) + gateway = require_gateway(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(gateway=gateway, 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..9de72d2a05f 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_gateway 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) + gateway = require_gateway(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(gateway=gateway, 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..651d73c3ac9 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_gateway 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) + gateway = require_gateway(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(gateway=gateway, 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..40306d7fe2f 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_gateway 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) + gateway = require_gateway(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(gateway=gateway, model=model) shape_error = assert_tool_search_shape(result) if shape_error is not None: error = f"[{model}] tool_search probe failed: {shape_error}"