refactor(e2e): fold claude_code HTTP probes onto shared ProxyClient methods

[yassin-berriai]

Relevant issues

Linear ticket

Resolves LIT-4556

Pre-Submission checklist

Please complete all items before asking a LiteLLM maintainer to review your PR

• I have added meaningful tests
• My PR passes all CI/CD checks (e.g., lint, format, unit tests)
• My PR's scope is as isolated as possible; it only solves 1 specific problem
• I have received a Greptile Confidence Score of at least 4/5 before requesting a maintainer review (Greptile reviews automatically once the PR is opened; only comment @greptileai to re-request a review after pushing changes)

Screenshots / Proof of Fix

This is a harness-internal refactor of two existing HTTP-probe cells; it adds no new asserted product behavior and changes no assertions, so there is no new manual proxy scenario beyond the QA runbook below. The two migrated cells stay real end-to-end rows in the daily compat matrix; the safety argument here is that the request the proxy sees is byte-for-byte what it saw before, plus the harness typecheck and whole-tree collection. Live matrix execution needs the real claude matrix credentials, which are not available in this environment

Wire shape is byte-for-byte identical to the old httpx probes (captured at 78f10dc):

count_tokens body OK: {'model': 'claude-haiku-4-5', 'messages': [{'role': 'user', 'content': 'hello world'}]}
headers OK (case-insensitive names): {'authorization': 'Bearer KEY', 'anthropic-version': '2023-06-01'}
tool_search body OK
ALL WIRE-SHAPE CHECKS PASS

basedpyright over the changed shared harness files is clean:

$ basedpyright tests/e2e/models.py tests/e2e/proxy_client.py tests/e2e/e2e_http.py
0 errors, 0 warnings, 0 notes

The whole claude_code tree still collects, including the ten migrated cells:

$ pytest --collect-only -qq claude_code/count_tokens claude_code/tool_search
claude_code/count_tokens/test_anthropic.py: 1
... (10 cells, no import or fixture errors)

Type

🧹 Refactoring

✅ Test

Changes

tests/e2e/claude_code/ is a second, parallel e2e harness that ships its own driver, builder, publisher, rate limiter, HTTP probe, and version resolver, and its HTTP probes deliberately hand-rolled their own httpx client. Keeping that probe layer separate from the shared transport is part of the maintenance liability the ticket tracks. This folds the HTTP-probe surface onto the shared transport and goes one step further than a pure lift-and-shift: it promotes the two probed endpoints to first-class shared ProxyClient methods so any other suite gets them for free

http_probe.py used to build its own httpx client and a hand-rolled ProbeResult dataclass for the two cells that have no CLI surface (count_tokens and tool_search). It now takes an injected ProxyClient and issues each request through ProxyClient.count_tokens / ProxyClient.messages, reusing the shared split control/data-plane routing, timeout, and typed Result union the rest of tests/e2e/ uses. httpx is gone from the module

The two endpoints are now shared ProxyClient methods rather than inline probe plumbing. ProxyClient.count_tokens(key, body) posts to /v1/messages/count_tokens and ProxyClient.messages(key, body) posts to /v1/messages, both carrying the anthropic-version header through a small AnthropicHeaders model and returning a validated typed Result. Their request and response bodies live in the shared models.py (CountTokensBody / CountTokensResponse, and an extended AnthropicMessagesBody / AnthropicMessagesResponse that models the tool-search tools array and the content-or-choices response shape), so a future messages suite reuses them instead of re-deriving the wire shape

The wire shape is preserved exactly. The migrated pydantic bodies serialize byte-for-byte to what the old httpx probes sent (model + messages for count_tokens; model + max_tokens + messages + the tool_search_tool_regex_20251119 tools array for tool_search), and the header names are HTTP-case-insensitive equivalents of the old ones with Content-Type set by the transport. The tool_search tools array is preserved intact: the tool_search_tool_regex_20251119 discovery tool plus the trivial add_numbers user tool are now the _TOOL_SEARCH_TOOLS pydantic constant, so the azure / bedrock tool_search rows keep sending exactly the list that surfaced the earlier provider bug. The assert_count_tokens_shape / assert_tool_search_shape helpers keep their contract (return None on success, else a diagnostic string) but now match on the typed Result; the status 429: wording is preserved so the compat conftest still classifies a rate-limited cell via RATE_LIMIT_SHAPED_RE

Dependency injection is threaded through the ten count_tokens/* and tool_search/* cells: _env.py gains require_proxy_client, which resolves the same LITELLM_PROXY_URL / LITELLM_MASTER_KEY env and builds the shared ProxyClient, and the cells now grab a ProxyClient plus the master key and pass them to the probe instead of raw base_url / api_key strings. Their asserts and compat_result reporting are untouched

What I deliberately left claude-specific, since the ticket scopes those out: the claude CLI driver (cli_driver.py), the compat-matrix building (matrix_builder.py), the tagged-union compat_result reporting and its xdist merge in conftest.py, the cross-process rate_limiter.py (the probes still acquire one token per call so probe traffic counts against the same per-provider budget as CLI rows), and pr_gate_version_resolver.py. These have no shared equivalent, and http_probe.py was the only claude_code module doing raw HTTP

QA runbook

Both probed cells are existing HTTP-probe rows; the refactor changes how the request is issued, not what it asserts, and the bodies are byte-identical to before. A reviewer can reproduce each by hand against a live proxy that fronts a Claude-style deployment (needs the provider credentials on the proxy side, e.g. an anthropic/* model registered as a claude-* name)

• tests/e2e/claude_code/count_tokens/test_anthropic.py::test_count_tokens_anthropic - /v1/messages/count_tokens returns a positive input_tokens for each Claude tier

  • POST http://localhost:4000/v1/messages/count_tokens with -H "Authorization: Bearer $LITELLM_MASTER_KEY" -H "anthropic-version: 2023-06-01" and body {"model": "claude-haiku-4-5", "messages": [{"role": "user", "content": "hello world"}]}
  • Expect 200 with {"input_tokens": <positive int>}
  • Sanity check: this test makes sense to add and is not hand-wavey (it asserts a positive integer count, not merely a 200), and is not potentially flaky

• tests/e2e/claude_code/tool_search/test_anthropic.py::test_tool_search_anthropic - /v1/messages accepts a tool_search_tool_regex_20251119 tool and the proxy forwards it without a 400

  • POST http://localhost:4000/v1/messages with the same auth headers and body {"model": "claude-haiku-4-5", "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": [{"type": "tool_search_tool_regex_20251119", "name": "tool_search_tool_regex"}, {"name": "add_numbers", "description": "Add two integers", "input_schema": {"type": "object", "properties": {"a": {"type": "integer"}, "b": {"type": "integer"}}, "required": ["a", "b"]}}]}
  • Expect 200 whose body has either content (Anthropic passthrough) or choices (OpenAI-normalized), proving the proxy attached the per-provider tool-search beta header and forwarded it
  • Sanity check: this test makes sense to add and is not hand-wavey (it asserts the request round-trips without a 400, not model tool-call behavior), and is not potentially flaky

Final Attestation

• The tests check the right things, including the edge cases, and regressions in the respective real-world customer use-cases are not possible after this PR

[gemini-code-assist]

Caution

The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased.

[CLAassistant]

Thank you for your submission! We really appreciate it. Like many open source projects, we ask that you sign our Contributor License Agreement before we can accept your contribution.
_{You have signed the CLA already but the status is still pending? Let us recheck it.}

[greptile-apps]

Greptile Summary

This PR folds the claude_code HTTP probe layer onto the shared ProxyClient transport, removing the hand-rolled httpx client from http_probe.py and promoting count_tokens and messages as first-class ProxyClient methods. The ten migrated test cells in count_tokens/ and tool_search/ now receive an injected ProxyClient instead of raw URL/key strings, and their assertions and compat_result reporting are untouched.

• ProxyClient.count_tokens and ProxyClient.messages are added alongside a new AnthropicHeaders header model that carries anthropic-version: 2023-06-01; wire shape is preserved byte-for-byte (verified in the PR description).
• models.py gains CountTokensBody, CountTokensResponse, and an extended AnthropicMessagesBody/AnthropicMessagesResponse (with tool-schema types and extra="allow" on the response for triage diagnostics); the tool-search tools constant is lifted to a module-level _TOOL_SEARCH_TOOLS tuple.
• _failure_diagnostic now has a wildcard arm and assert_tool_search_shape includes the actual response-key list in the "neither content nor choices" diagnostic, addressing concerns raised in the previous review cycle.

Confidence Score: 5/5

Safe to merge — this is a pure test-harness refactor with no changes to production code paths; all assertions are preserved intact.

The refactor correctly threads dependency injection through all ten migrated cells, the wire shape is preserved (headers, body fields, and serialization all match the old httpx probes), and both previously-flagged gaps from the last review cycle — the missing wildcard arm in _failure_diagnostic and the missing key-list in the "neither content nor choices" diagnostic — are now resolved. No functional regressions are introduced.

No files require special attention.

Important Files Changed

Filename	Overview
tests/e2e/claude_code/http_probe.py	Core refactor: replaces hand-rolled httpx probes with injected ProxyClient methods; wildcard arm and diagnostic key-list both present; clean.
tests/e2e/proxy_client.py	Adds count_tokens and messages methods plus _anthropic_headers helper; header construction preserves Authorization and adds anthropic-version correctly.
tests/e2e/models.py	Promotes AnthropicMessagesBody/Response to full models with tools, content, choices; adds CountTokensBody/Response and supporting tool schema models; extra="allow" on response for triage diagnostics.
tests/e2e/e2e_http.py	Adds AnthropicHeaders extending AuthHeaders with a defaulted anthropic-version field; straightforward addition with no concerns.
tests/e2e/claude_code/_env.py	Adds require_proxy_client wrapping require_proxy and build_proxy_client; correctly sets both planes to the same base_url for the single-URL matrix case.
tests/e2e/claude_code/count_tokens/test_anthropic.py	Migrated from require_proxy/base_url to require_proxy_client/client; assertions and compat_result reporting unchanged.
tests/e2e/claude_code/tool_search/test_anthropic.py	Same mechanical migration as count_tokens cells; assertions and compat_result reporting unchanged.
tests/e2e/CLAUDE.md	Documentation updated to reflect that HTTP probes now use the shared transport while the CLI-driving path remains bespoke.

_{Reviews (4): Last reviewed commit: "refactor(e2e): fold claude_code HTTP pro..." | Re-trigger Greptile}

[yassin-berriai]

@greptileai please review the current head 7755bd5. Addressed both findings: _failure_diagnostic now has a case _: wildcard so a future Result member can't silently return None, and the tool_search shape-check restores the response-key list (AnthropicMessagesResponse now allows extra keys, and the diagnostic reports keys=[...]) for triage parity with the old probe

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.

📢 Thoughts on this report? Let us know!

[codspeed-hq]

Merging this PR will not alter performance

✅ 31 untouched benchmarks

---

_{Comparing litellm_e2e_claude_probe_gateway_methods (78f10dc) with litellm_internal_staging (010b200)¹}

Footnotes

1. No successful run was found on litellm_internal_staging (08fa250) during the generation of this report, so 010b200 was used instead as the comparison base. There might be some changes unrelated to this pull request in this report. ↩

[mubashir1osmani]

i'd rather not remove this, this tool list brought up an internally found bug on azure / bedrock. i assigned this to @mateo

[yassin-berriai]

Good call to double-check, but the tool list isn't removed; it's preserved byte-for-byte. The raw "tools": [...] dict you're pointing at moved to the _TOOL_SEARCH_TOOLS pydantic constant at the top of http_probe.py (the tool_search_tool_regex_20251119 discovery tool plus the add_numbers user tool), and probe_tool_search still sends that exact array. The tool_search/test_azure.py, test_bedrock_converse.py, and test_bedrock_invoke.py rows are untouched and keep exercising the same list that surfaced the internal azure / bedrock bug. The serialized request body is identical to what the old httpx probe sent, which the wire-shape check in the PR description confirms. Leaving this open for @mateo to sign off since you assigned it to him

[yassin-berriai]

@greptileai please review the current head 260bc32. Since your last review this was rebased onto litellm_internal_staging and squashed to a single commit; the unit-level harness test (test_http_probe.py) was dropped, so the two migrated e2e rows (count_tokens, tool_search) are now the only tests exercising the probe path, and the safety argument for the refactor rests on the byte-for-byte wire-shape equality documented in the PR description. The case _: wildcard on _failure_diagnostic and the keys=[...] diagnostic in assert_tool_search_shape from the previous round are still in place

[yassin-berriai]

@greptileai please review the current head 78f10dc. Rebased onto litellm_internal_staging, which since the last review renamed the shared Gateway class to ProxyClient and its module to proxy_client.py (#33750); this reintegrates the probe migration against ProxyClient.count_tokens / ProxyClient.messages and require_proxy_client. No behavior change from the version you scored 5/5; the wire shape is still byte-for-byte identical (verified by importing the models under 3.12 and dumping the serialized bodies) and all ten migrated cells still collect