Skip to content

refactor(e2e/claude_code): dedupe tool_use cell body into shared _tool_use module #33625

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
mateo-berri wants to merge 1 commit into
base: litellm_internal_staging
Choose a base branch
from
Open

refactor(e2e/claude_code): dedupe tool_use cell body into shared _tool_use module #33625

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
168 changes: 168 additions & 0 deletions tests/e2e/claude_code/_tool_use.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,168 @@
"""Shared body for the `tool_use` / `tool_use_streaming` x <provider> compat cells.

Every tool_use cell follows the same skeleton:

1. Read the proxy base URL + API key from env, fail-early if missing.
2. Ask each model tier (fanned out via `run_claude_models_parallel`)
to invoke the built-in `Bash` tool.
3. Inspect each model's outcome and report one `compat_result` row per
model; `ClaudeCLIError`, non-zero exit, and a missing `tool_use`
content block are all per-model fails, everything else is a
per-model pass.
4. Surface a joined failure message via `pytest.fail(...)` so the
pytest run also goes red.

The streaming variant (`verify_streaming=True`) adds the
`--include-partial-messages` CLI flag and additionally requires that
`input_json_delta` stream events were observed; zero deltas means the
proxy collapsed the streamed tool input into a single complete block or
stripped fine-grained tool streaming.

Security rationale for `TOOL_USE_ARGS`: the prompt asks for the exact
command `echo pong`, the `--allowed-tools` rule restricts the Bash tool
to that command, and `--permission-mode dontAsk` auto-denies anything
else the model returns instead of executing it. `dontAsk` mode in
headless `--print` mode only runs tools matching an explicit `allow`
rule (plus the built-in read-only set), so a compromised provider
response cannot turn the `Bash` allowlist into arbitrary host execution
(which would expose `docker inspect compat-proxy` /
`/proc/<proxy_pid>/environ` and thereby provider credentials living in
the proxy container).

The conftest infers `(feature_id, provider)` purely from the test file
path, so each per-provider file just declares its model list and calls
`run_tool_use_cell(...)`. This keeps all cell logic in one place; a
future tweak to the outcome checks or the failure-loop shape now
propagates to every cell automatically.

The leading underscore in the filename is what keeps pytest from
collecting this module as a test file.
"""

from __future__ import annotations

from typing import Any, Callable, Mapping, Sequence

import pytest

from claude_code._env import require_proxy
from claude_code.cli_driver import (
ClaudeCLIError,
DriverResult,
failure_diagnostic,
run_claude_models_parallel,
)


ClaudeRunner = Callable[..., Mapping[str, DriverResult | ClaudeCLIError]]

TOOL_USE_PROMPT = "Use the Bash tool to run the command `echo pong` and report what it printed."

TOOL_USE_ARGS = (
"--allowed-tools",
"Bash(echo pong)",
"--permission-mode",
"dontAsk",
)

TOOL_USE_STREAMING_ARGS = TOOL_USE_ARGS + ("--include-partial-messages",)


def _has_tool_use_event(events: Sequence[Mapping[str, Any]]) -> bool:
"""Walk the stream-json events and return True if any assistant
message included a `tool_use` content block."""
for event in events:
if event.get("type") != "assistant":
continue
message = event.get("message") or {}
content = message.get("content")
if not isinstance(content, list):
continue
for block in content:
if isinstance(block, dict) and block.get("type") == "tool_use":
return True
return False


def _count_input_json_deltas(events: Sequence[Mapping[str, Any]]) -> int:
"""Count `input_json_delta` records among the `stream_event`
entries. Zero means the proxy collapsed the streamed tool input
into a single complete block instead of forwarding the incremental
deltas the upstream emitted."""
inner_events = (event.get("event") for event in events if event.get("type") == "stream_event")
return sum(
1
for inner in inner_events
if isinstance(inner, Mapping)
and inner.get("type") == "content_block_delta"
and isinstance(inner.get("delta"), Mapping)
and inner["delta"].get("type") == "input_json_delta"
)


def run_tool_use_cell(
*,
compat_result,
models: Sequence[str],
verify_streaming: bool = False,
env: Mapping[str, str] | None = None,
runner: ClaudeRunner = run_claude_models_parallel,
) -> None:
"""Run the shared `tool_use` / `tool_use_streaming` x <provider> cell body.

With ``verify_streaming=False`` the cell asserts that a `tool_use`
content block came back over the wire. With ``verify_streaming=True``
it additionally passes ``--include-partial-messages`` and asserts
that `input_json_delta` stream events were observed for the block,
which catches both known gateway regressions: buffering the streamed
tool input into one complete block and stripping the
fine-grained-tool-streaming beta so the upstream falls back to
non-streaming tool_use.
"""
base_url, api_key = require_proxy(compat_result, env=env)

extra_args = TOOL_USE_STREAMING_ARGS if verify_streaming else TOOL_USE_ARGS

outcomes = runner(
models=models,
prompt=TOOL_USE_PROMPT,
base_url=base_url,
api_key=api_key,
extra_args=extra_args,
)

failures = []
for model in models:
outcome = outcomes[model]
if isinstance(outcome, ClaudeCLIError):
error = f"[{model}] {outcome}"
compat_result.add({"status": "fail", "error": error})
failures.append(error)
continue

if outcome.exit_code != 0:
error = f"[{model}] claude CLI failed: {failure_diagnostic(outcome)}"
compat_result.add({"status": "fail", "error": error})
failures.append(error)
continue

if not _has_tool_use_event(outcome.events):
error = f"[{model}] no tool_use content block observed in stream-json events"
compat_result.add({"status": "fail", "error": error})
failures.append(error)
continue

if verify_streaming and _count_input_json_deltas(outcome.events) == 0:
error = (
f"[{model}] no input_json_delta stream events observed; proxy "
f"likely buffered the tool input into a complete block or "
f"stripped fine-grained tool streaming"
)
compat_result.add({"status": "fail", "error": error})
failures.append(error)
continue

compat_result.add({"status": "pass"})

if failures:
pytest.fail("; ".join(failures), pytrace=False)
88 changes: 6 additions & 82 deletions tests/e2e/claude_code/tool_use/test_anthropic.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -5,6 +5,10 @@
that the upstream returned a `tool_use` content block. This proves the
proxy preserves Claude Code's tool-call wire shape end-to-end.

Bash is restricted to the exact command `echo pong` plus
`--permission-mode dontAsk`; the security rationale lives in
`claude_code/_tool_use.py` alongside the shared cell body.

The (feature, provider) for this cell is inferred from the file path by
`tests/e2e/claude_code/conftest.py`:

Expand All @@ -15,16 +19,9 @@

from __future__ import annotations

from typing import Any, Mapping, Sequence

import pytest

from claude_code._env import require_proxy
from claude_code.cli_driver import (
ClaudeCLIError,
failure_diagnostic,
run_claude_models_parallel,
)
from claude_code._tool_use import run_tool_use_cell


ANTHROPIC_MODELS = [
Expand All @@ -33,82 +30,9 @@
"claude-opus-4-7",
]

# Built-in tool-use prompt: ask Claude to use the `Bash` tool. The CLI
# allow-lists the tool via `--allowed-tools` so the run completes without
# an interactive permission prompt.
TOOL_USE_PROMPT = (
"Use the Bash tool to run the command `echo pong` and report what it printed."
)
# Restrict the Bash tool to the exact command `echo pong` and put the
# CLI in `dontAsk` mode so anything else the model returns is auto-
# denied instead of executed. `dontAsk` mode in headless `--print` mode
# only runs tools matching an explicit `allow` rule (plus the built-in
# read-only set), so a compromised provider response cannot turn the
# `Bash` allowlist into arbitrary host execution (which would expose
# `docker inspect compat-proxy` / `/proc/<proxy_pid>/environ` and
# thereby provider credentials living in the proxy container).
TOOL_USE_ARGS = [
"--allowed-tools",
"Bash(echo pong)",
"--permission-mode",
"dontAsk",
]


def _has_tool_use_event(events: Sequence[Mapping[str, Any]]) -> bool:
"""Walk the stream-json events and return True if any assistant
message included a `tool_use` content block."""
for event in events:
if event.get("type") != "assistant":
continue
message = event.get("message") or {}
content = message.get("content")
if not isinstance(content, list):
continue
for block in content:
if isinstance(block, dict) and block.get("type") == "tool_use":
return True
return False


@pytest.mark.covers("llm.messages.anthropic.tool_use.nonstream.works")
def test_tool_use_anthropic(compat_result):
"""Drive the `claude` CLI against the LiteLLM proxy and assert a
tool call was emitted on the wire."""
base_url, api_key = require_proxy(compat_result)

outcomes = run_claude_models_parallel(
models=ANTHROPIC_MODELS,
prompt=TOOL_USE_PROMPT,
base_url=base_url,
api_key=api_key,
extra_args=TOOL_USE_ARGS,
)

failures = []
for model in ANTHROPIC_MODELS:
outcome = outcomes[model]
if isinstance(outcome, ClaudeCLIError):
error = f"[{model}] {outcome}"
compat_result.add({"status": "fail", "error": error})
failures.append(error)
continue

if outcome.exit_code != 0:
error = f"[{model}] claude CLI failed: {failure_diagnostic(outcome)}"
compat_result.add({"status": "fail", "error": error})
failures.append(error)
continue

if not _has_tool_use_event(outcome.events):
error = (
f"[{model}] no tool_use content block observed in stream-json events"
)
compat_result.add({"status": "fail", "error": error})
failures.append(error)
continue

compat_result.add({"status": "pass"})

if failures:
pytest.fail("; ".join(failures), pytrace=False)
run_tool_use_cell(compat_result=compat_result, models=ANTHROPIC_MODELS)
78 changes: 6 additions & 72 deletions tests/e2e/claude_code/tool_use/test_azure.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -9,6 +9,10 @@
identically to anthropic.com; LiteLLM's `azure_ai/claude-*` route
inherits the full Anthropic tool-use transformation.

Bash is restricted to the exact command `echo pong` plus
`--permission-mode dontAsk`; see `claude_code/_tool_use.py` for the
security rationale.

The (feature, provider) for this cell is inferred from the file path by
`tests/e2e/claude_code/conftest.py`:

Expand All @@ -19,16 +23,9 @@

from __future__ import annotations

from typing import Any, Mapping, Sequence

import pytest

from claude_code._env import require_proxy
from claude_code.cli_driver import (
ClaudeCLIError,
failure_diagnostic,
run_claude_models_parallel,
)
from claude_code._tool_use import run_tool_use_cell


AZURE_MODELS = [
Expand All @@ -37,72 +34,9 @@
"claude-opus-4-7-azure",
]

TOOL_USE_PROMPT = (
"Use the Bash tool to run the command `echo pong` and report what it printed."
)
# Bash is restricted to the exact command `echo pong` + `dontAsk`
# permission mode; see `tool_use/test_anthropic.py` for the security
# rationale.
TOOL_USE_ARGS = [
"--allowed-tools",
"Bash(echo pong)",
"--permission-mode",
"dontAsk",
]


def _has_tool_use_event(events: Sequence[Mapping[str, Any]]) -> bool:
for event in events:
if event.get("type") != "assistant":
continue
message = event.get("message") or {}
content = message.get("content")
if not isinstance(content, list):
continue
for block in content:
if isinstance(block, dict) and block.get("type") == "tool_use":
return True
return False


@pytest.mark.covers("llm.messages.azure_foundry.tool_use.nonstream.works")
def test_tool_use_azure(compat_result):
"""Drive the `claude` CLI against the LiteLLM proxy and assert a
tool call was emitted on the wire."""
base_url, api_key = require_proxy(compat_result)

outcomes = run_claude_models_parallel(
models=AZURE_MODELS,
prompt=TOOL_USE_PROMPT,
base_url=base_url,
api_key=api_key,
extra_args=TOOL_USE_ARGS,
)

failures = []
for model in AZURE_MODELS:
outcome = outcomes[model]
if isinstance(outcome, ClaudeCLIError):
error = f"[{model}] {outcome}"
compat_result.add({"status": "fail", "error": error})
failures.append(error)
continue

if outcome.exit_code != 0:
error = f"[{model}] claude CLI failed: {failure_diagnostic(outcome)}"
compat_result.add({"status": "fail", "error": error})
failures.append(error)
continue

if not _has_tool_use_event(outcome.events):
error = (
f"[{model}] no tool_use content block observed in stream-json events"
)
compat_result.add({"status": "fail", "error": error})
failures.append(error)
continue

compat_result.add({"status": "pass"})

if failures:
pytest.fail("; ".join(failures), pytrace=False)
run_tool_use_cell(compat_result=compat_result, models=AZURE_MODELS)
Loading
Loading