litellm/tests/test_litellm/proxy/client/cli/test_agents.py

620 lines
22 KiB
Python

import os
import sys
from unittest.mock import patch
import click
import pytest
import requests
from click.testing import CliRunner
sys.path.insert(
0, os.path.abspath("../../..")
) # Adds the parent directory to the system path
from litellm.proxy.client.cli.commands.agents import (
AgentRunError,
agent_commands,
agent_launch_args,
agent_profile,
build_agent_env,
run_agent,
verify_proxy_key,
)
AGENTS_MODULE = "litellm.proxy.client.cli.commands.agents"
def _agent_command(name):
return next(c for c in agent_commands() if c.name == name)
class _FakeResponse:
def __init__(self, status_code):
self.status_code = status_code
class _FakeJsonResponse:
def __init__(self, status_code, payload=None):
self.status_code = status_code
self._payload = payload
def json(self):
return self._payload
class TestAgentProfile:
def test_claude_is_anthropic(self):
name, profiles = agent_profile("claude")
assert name == "Claude Code"
assert profiles == frozenset({"anthropic"})
def test_claude_full_path_uses_basename(self):
name, profiles = agent_profile("/usr/local/bin/claude")
assert name == "Claude Code"
assert profiles == frozenset({"anthropic"})
def test_codex_and_opencode_are_openai(self):
assert agent_profile("codex") == ("Codex", frozenset({"openai"}))
assert agent_profile("opencode") == ("OpenCode", frozenset({"openai"}))
def test_pi_is_litellm(self):
assert agent_profile("pi") == ("pi", frozenset({"litellm"}))
def test_unknown_command_gets_both_profiles(self):
name, profiles = agent_profile("mytool")
assert name == "mytool"
assert profiles == frozenset({"anthropic", "openai"})
class TestBuildAgentEnv:
def test_anthropic_profile_uses_bare_root_and_bearer(self):
env = build_agent_env(
{}, "http://localhost:4000/", "sk-key", frozenset({"anthropic"})
)
assert env["ANTHROPIC_BASE_URL"] == "http://localhost:4000"
assert env["ANTHROPIC_AUTH_TOKEN"] == "sk-key"
assert "OPENAI_BASE_URL" not in env
assert "OPENAI_API_KEY" not in env
def test_anthropic_profile_drops_existing_api_key(self):
env = build_agent_env(
{"ANTHROPIC_API_KEY": "real-key"},
"http://localhost:4000",
"sk-key",
frozenset({"anthropic"}),
)
assert "ANTHROPIC_API_KEY" not in env
def test_openai_profile_appends_v1(self):
env = build_agent_env(
{}, "http://localhost:4000/", "sk-key", frozenset({"openai"})
)
assert env["OPENAI_BASE_URL"] == "http://localhost:4000/v1"
assert env["OPENAI_API_KEY"] == "sk-key"
assert "ANTHROPIC_BASE_URL" not in env
def test_both_profiles_set_everything(self):
env = build_agent_env(
{}, "http://localhost:4000", "sk-key", frozenset({"anthropic", "openai"})
)
assert env["ANTHROPIC_BASE_URL"] == "http://localhost:4000"
assert env["OPENAI_BASE_URL"] == "http://localhost:4000/v1"
assert env["ANTHROPIC_AUTH_TOKEN"] == "sk-key"
assert env["OPENAI_API_KEY"] == "sk-key"
def test_litellm_profile_exports_only_the_proxy_key(self):
env = build_agent_env(
{}, "http://localhost:4000/", "sk-key", frozenset({"litellm"})
)
assert env["LITELLM_PROXY_API_KEY"] == "sk-key"
assert "ANTHROPIC_BASE_URL" not in env
assert "OPENAI_BASE_URL" not in env
assert "OPENAI_API_KEY" not in env
def test_preserves_unrelated_env_and_does_not_mutate_input(self):
base = {"PATH": "/usr/bin", "ANTHROPIC_API_KEY": "real-key"}
env = build_agent_env(
base, "http://localhost:4000", "sk-key", frozenset({"anthropic"})
)
assert env["PATH"] == "/usr/bin"
assert base == {"PATH": "/usr/bin", "ANTHROPIC_API_KEY": "real-key"}
class TestAgentLaunchArgs:
def test_claude_and_opencode_get_no_extra_args(self):
assert agent_launch_args("claude", "http://localhost:4000") == []
assert agent_launch_args("opencode", "http://localhost:4000") == []
def test_unknown_agent_gets_no_extra_args(self):
assert agent_launch_args("mytool", "http://localhost:4000") == []
def test_codex_points_provider_at_proxy_over_http(self):
args = agent_launch_args("codex", "http://localhost:4000/")
joined = " ".join(args)
assert 'model_provider="litellm"' in args
assert 'model_providers.litellm.base_url="http://localhost:4000/v1"' in args
assert 'model_providers.litellm.env_key="OPENAI_API_KEY"' in args
assert 'model_providers.litellm.wire_api="responses"' in args
assert "model_providers.litellm.supports_websockets=false" in args
assert joined.count("-c") == 6
def test_codex_uses_basename(self):
assert agent_launch_args("/usr/local/bin/codex", "http://localhost:4000") == (
agent_launch_args("codex", "http://localhost:4000")
)
def test_pi_gets_no_static_args(self):
assert agent_launch_args("pi", "http://localhost:4000") == []
class TestVerifyProxyKey:
def test_ok_status_passes_and_uses_models_endpoint(self):
captured = {}
def fake_get(url, headers, timeout):
captured["url"] = url
captured["headers"] = headers
return _FakeResponse(200)
verify_proxy_key("http://localhost:4000/", "sk-key", get=fake_get)
assert captured["url"] == "http://localhost:4000/v1/models"
assert captured["headers"] == {"Authorization": "Bearer sk-key"}
@pytest.mark.parametrize("status", [401, 403])
def test_rejected_key_raises(self, status):
with pytest.raises(AgentRunError, match="rejected your key"):
verify_proxy_key(
"http://localhost:4000",
"sk-key",
get=lambda *a, **k: _FakeResponse(status),
)
def test_unreachable_proxy_raises(self):
def boom(*a, **k):
raise requests.ConnectionError("refused")
with pytest.raises(AgentRunError, match="Could not reach"):
verify_proxy_key("http://localhost:4000", "sk-key", get=boom)
def test_other_non_2xx_is_tolerated(self):
verify_proxy_key(
"http://localhost:4000",
"sk-key",
get=lambda *a, **k: _FakeResponse(500),
)
class TestRunAgent:
def test_wires_env_and_launches_resolved_binary(self):
calls = {}
def fake_launcher(path, args, env):
calls["path"] = path
calls["args"] = tuple(args)
calls["env"] = dict(env)
run_agent(
"http://localhost:4000",
"sk-key",
["claude", "--resume"],
base_env={"PATH": "/usr/bin", "ANTHROPIC_API_KEY": "leaked"},
which=lambda name: "/usr/local/bin/claude",
verify=lambda *a: None,
launcher=fake_launcher,
)
assert calls["path"] == "/usr/local/bin/claude"
assert calls["args"] == ("claude", "--resume")
env = calls["env"]
assert env["ANTHROPIC_BASE_URL"] == "http://localhost:4000"
assert env["ANTHROPIC_AUTH_TOKEN"] == "sk-key"
assert "ANTHROPIC_API_KEY" not in env
assert "OPENAI_BASE_URL" not in env
def test_codex_gets_openai_env(self):
calls = {}
run_agent(
"http://localhost:4000",
"sk-key",
["codex"],
base_env={},
which=lambda name: "/usr/local/bin/codex",
verify=lambda *a: None,
launcher=lambda p, a, e: calls.update(env=dict(e)),
)
assert calls["env"]["OPENAI_BASE_URL"] == "http://localhost:4000/v1"
assert calls["env"]["OPENAI_API_KEY"] == "sk-key"
assert "ANTHROPIC_BASE_URL" not in calls["env"]
def test_codex_injects_proxy_provider_args_before_user_args(self):
calls = {}
run_agent(
"http://localhost:4000",
"sk-key",
["codex", "exec", "do a thing"],
base_env={},
which=lambda name: "/usr/local/bin/codex",
verify=lambda *a: None,
launcher=lambda p, a, e: calls.update(args=tuple(a)),
)
args = calls["args"]
assert args[0] == "codex"
assert args[-2:] == ("exec", "do a thing")
assert 'model_provider="litellm"' in args
assert 'model_providers.litellm.base_url="http://localhost:4000/v1"' in args
# overrides must precede the codex subcommand so codex parses them
assert args.index('model_provider="litellm"') < args.index("exec")
def test_pi_preparer_runs_after_verify_and_before_launch(self):
order = []
captured = {}
def fake_prepare(base_url, api_key, base_env):
order.append("prepare")
captured["args"] = (base_url, api_key, dict(base_env))
return []
run_agent(
"http://localhost:4000",
"sk-key",
["pi"],
base_env={"HOME": "/home/u"},
which=lambda name: "/usr/local/bin/pi",
verify=lambda *a: order.append("verify"),
launcher=lambda *a: order.append("launch"),
preparers={"pi": fake_prepare},
)
assert order == ["verify", "prepare", "launch"]
assert captured["args"] == (
"http://localhost:4000",
"sk-key",
{"HOME": "/home/u"},
)
def test_pi_prepared_args_precede_user_args_and_env_has_proxy_key(self):
calls = {}
run_agent(
"http://localhost:4000",
"sk-key",
["pi", "-p", "hello"],
base_env={},
which=lambda name: "/usr/local/bin/pi",
verify=lambda *a: None,
launcher=lambda p, a, e: calls.update(args=tuple(a), env=dict(e)),
preparers={"pi": lambda *a: ["--model", "litellm/m-1"]},
)
# user args come last so a user-supplied --model wins in pi's parser
assert calls["args"] == ("pi", "--model", "litellm/m-1", "-p", "hello")
assert calls["env"]["LITELLM_PROXY_API_KEY"] == "sk-key"
assert "OPENAI_API_KEY" not in calls["env"]
assert "ANTHROPIC_BASE_URL" not in calls["env"]
def test_failed_preparer_aborts_before_launch(self):
launched = []
def boom(*a):
raise AgentRunError("sync failed")
with pytest.raises(AgentRunError, match="sync failed"):
run_agent(
"http://localhost:4000",
"sk-key",
["pi"],
base_env={},
which=lambda name: "/usr/local/bin/pi",
verify=lambda *a: None,
launcher=lambda *a: launched.append(a),
preparers={"pi": boom},
)
assert launched == []
def test_prepare_pi_syncs_models_json_and_pins_first_model(self, tmp_path):
from litellm.proxy.client.cli.commands.agents import prepare_pi
def fake_get(url, headers, timeout):
if url.endswith("/model_group/info"):
return _FakeJsonResponse(
200,
{"data": [{"model_group": "m-first", "max_input_tokens": 131072, "max_output_tokens": 8192}]},
)
return _FakeJsonResponse(200, {"data": [{"id": "m-first"}, {"id": "m-second"}]})
pin = prepare_pi(
"http://localhost:4000",
"sk-key",
{"PI_CODING_AGENT_DIR": str(tmp_path)},
get=fake_get,
)
assert pin == ["--model", "litellm/m-first"]
import json
written = json.loads((tmp_path / "models.json").read_text())
assert written["providers"]["litellm"]["apiKey"] == "$LITELLM_PROXY_API_KEY"
assert written["providers"]["litellm"]["models"] == [
{"id": "m-first", "contextWindow": 131072, "maxTokens": 8192},
{"id": "m-second"},
]
def test_prepare_pi_surfaces_fetch_failure_as_agent_error(self, tmp_path):
from litellm.proxy.client.cli.commands.agents import prepare_pi
with pytest.raises(AgentRunError, match="HTTP 500"):
prepare_pi(
"http://localhost:4000",
"sk-key",
{"PI_CODING_AGENT_DIR": str(tmp_path)},
get=lambda *a, **k: _FakeJsonResponse(500),
)
def test_claude_has_no_preparer(self):
prepared = []
def fake_prepare(*a):
prepared.append(a)
return []
run_agent(
"http://localhost:4000",
"sk-key",
["claude"],
base_env={},
which=lambda name: "/usr/local/bin/claude",
verify=lambda *a: None,
launcher=lambda *a: None,
preparers={"pi": fake_prepare},
)
assert prepared == []
def test_claude_launches_without_injected_args(self):
calls = {}
run_agent(
"http://localhost:4000",
"sk-key",
["claude", "--resume"],
base_env={},
which=lambda name: "/usr/local/bin/claude",
verify=lambda *a: None,
launcher=lambda p, a, e: calls.update(args=tuple(a)),
)
assert calls["args"] == ("claude", "--resume")
def test_missing_binary_raises_with_install_hint(self):
with pytest.raises(AgentRunError, match="claude.*Install it first"):
run_agent(
"http://localhost:4000",
"sk-key",
["claude"],
base_env={},
which=lambda name: None,
verify=lambda *a: None,
launcher=lambda *a: None,
)
def test_skip_verify_does_not_call_verify(self):
verified = []
launched = []
run_agent(
"http://localhost:4000",
"sk-key",
["claude"],
skip_verify=True,
base_env={},
which=lambda name: "/usr/local/bin/claude",
verify=lambda *a: verified.append(a),
launcher=lambda *a: launched.append(a),
)
assert verified == []
assert len(launched) == 1
def test_verify_failure_aborts_before_launch(self):
launched = []
def boom(*a):
raise AgentRunError("rejected")
with pytest.raises(AgentRunError):
run_agent(
"http://localhost:4000",
"sk-key",
["claude"],
base_env={},
which=lambda name: "/usr/local/bin/claude",
verify=boom,
launcher=lambda *a: launched.append(a),
)
assert launched == []
def test_empty_command_raises(self):
with pytest.raises(AgentRunError):
run_agent("http://localhost:4000", "sk-key", [])
def test_reattach_terminal_runs_just_before_launch(self):
order = []
run_agent(
"http://localhost:4000",
"sk-key",
["claude"],
skip_verify=True,
base_env={},
which=lambda name: "/usr/local/bin/claude",
launcher=lambda *a: order.append("launch"),
reattach_terminal=lambda: order.append("reattach"),
)
assert order == ["reattach", "launch"]
def test_no_reattach_terminal_by_default(self):
order = []
run_agent(
"http://localhost:4000",
"sk-key",
["claude"],
skip_verify=True,
base_env={},
which=lambda name: "/usr/local/bin/claude",
launcher=lambda *a: order.append("launch"),
)
assert order == ["launch"]
class TestAgentCommands:
def setup_method(self):
self.runner = CliRunner()
def test_one_command_per_known_agent(self):
assert {c.name for c in agent_commands()} == {"claude", "codex", "opencode", "pi"}
def test_claude_launches_with_stored_key_and_forwards_args(self):
captured = {}
def fake_run_agent(base_url, api_key, command, **kwargs):
captured["base_url"] = base_url
captured["api_key"] = api_key
captured["command"] = list(command)
captured["skip_verify"] = kwargs.get("skip_verify")
with patch(f"{AGENTS_MODULE}.run_agent", side_effect=fake_run_agent):
result = self.runner.invoke(
_agent_command("claude"),
["--resume", "-p", "hi"],
obj={"base_url": "http://localhost:4000", "api_key": "sk-key"},
)
assert result.exit_code == 0, result.output
assert captured["api_key"] == "sk-key"
assert captured["command"] == ["claude", "--resume", "-p", "hi"]
assert captured["skip_verify"] is False
assert (
"routing Claude Code through proxy at http://localhost:4000"
in result.output
)
def test_codex_shows_friendly_name(self):
captured = {}
with patch(
f"{AGENTS_MODULE}.run_agent",
side_effect=lambda b, k, c, **kw: captured.update(command=list(c)),
):
result = self.runner.invoke(
_agent_command("codex"),
["exec", "do a thing"],
obj={"base_url": "http://localhost:4000", "api_key": "sk-key"},
)
assert result.exit_code == 0, result.output
assert captured["command"] == ["codex", "exec", "do a thing"]
assert "routing Codex through proxy" in result.output
def test_skip_verify_is_consumed_not_forwarded(self):
captured = {}
def fake_run_agent(base_url, api_key, command, **kwargs):
captured["command"] = list(command)
captured["skip_verify"] = kwargs.get("skip_verify")
with patch(f"{AGENTS_MODULE}.run_agent", side_effect=fake_run_agent):
result = self.runner.invoke(
_agent_command("claude"),
["--skip-verify", "--resume"],
obj={"base_url": "http://localhost:4000", "api_key": "sk-key"},
)
assert result.exit_code == 0, result.output
assert captured["skip_verify"] is True
assert captured["command"] == ["claude", "--resume"]
def test_non_interactive_without_key_errors_clearly(self):
with (
patch(f"{AGENTS_MODULE}._is_interactive", return_value=False),
patch(f"{AGENTS_MODULE}.run_agent") as mock_run,
):
result = self.runner.invoke(
_agent_command("claude"),
[],
obj={"base_url": "http://localhost:4000", "api_key": None},
)
assert result.exit_code != 0
assert "LITELLM_PROXY_API_KEY" in result.output
mock_run.assert_not_called()
def test_interactive_without_key_logs_in_then_launches(self):
captured = {}
@click.command()
def fake_login():
pass
with (
patch(f"{AGENTS_MODULE}._is_interactive", return_value=True),
patch(f"{AGENTS_MODULE}.login", fake_login),
patch(
f"{AGENTS_MODULE}.get_stored_api_key", return_value="sk-after-login"
) as mock_get,
patch(
f"{AGENTS_MODULE}.run_agent",
side_effect=lambda base_url, api_key, command, **k: captured.update(
api_key=api_key
),
),
):
result = self.runner.invoke(
_agent_command("claude"),
[],
obj={"base_url": "http://localhost:4000", "api_key": None},
)
assert result.exit_code == 0, result.output
assert captured["api_key"] == "sk-after-login"
mock_get.assert_called_once_with(expected_base_url="http://localhost:4000")
def test_agent_run_error_becomes_click_error(self):
with patch(
f"{AGENTS_MODULE}.run_agent",
side_effect=AgentRunError("could not reach proxy"),
):
result = self.runner.invoke(
_agent_command("claude"),
[],
obj={"base_url": "http://localhost:4000", "api_key": "sk-key"},
)
assert result.exit_code != 0
assert "could not reach proxy" in result.output
def test_interactive_session_reattaches_terminal_before_handoff(self):
from litellm.proxy.client.cli.commands.agents import (
_restore_controlling_terminal,
)
captured = {}
with (
patch(f"{AGENTS_MODULE}._is_interactive", return_value=True),
patch(
f"{AGENTS_MODULE}.run_agent",
side_effect=lambda b, k, c, **kw: captured.update(kw),
),
):
result = self.runner.invoke(
_agent_command("claude"),
[],
obj={"base_url": "http://localhost:4000", "api_key": "sk-key"},
)
assert result.exit_code == 0, result.output
assert captured["reattach_terminal"] is _restore_controlling_terminal
def test_non_interactive_agent_mode_leaves_stdin_alone(self):
captured = {}
with (
patch(f"{AGENTS_MODULE}._is_interactive", return_value=False),
patch(
f"{AGENTS_MODULE}.run_agent",
side_effect=lambda b, k, c, **kw: captured.update(kw),
),
):
result = self.runner.invoke(
_agent_command("claude"),
[],
obj={"base_url": "http://localhost:4000", "api_key": "sk-key"},
)
assert result.exit_code == 0, result.output
assert captured["reattach_terminal"] is None