mirror of
https://github.com/BerriAI/litellm.git
synced 2026-08-28 05:25:59 +00:00
Merge branch 'litellm_internal_staging' of https://github.com/BerriAI/litellm into litellm_lit_5785_vertex_regional_pricing
# Conflicts: # type-discipline-budget.json
This commit is contained in:
mateo-berri
commit
7e27e211a1
26 changed files with 2069 additions and 233 deletions
|
|
@ -57,7 +57,7 @@
|
|||
"limit": 5663
|
||||
},
|
||||
"reportMissingTypeArgument": {
|
||||
"limit": 15556
|
||||
"limit": 15555
|
||||
},
|
||||
"reportMissingTypeStubs": {
|
||||
"limit": 40
|
||||
|
|
@ -105,13 +105,13 @@
|
|||
"limit": 109
|
||||
},
|
||||
"reportUnknownMemberType": {
|
||||
"limit": 39030
|
||||
"limit": 39017
|
||||
},
|
||||
"reportUnknownParameterType": {
|
||||
"limit": 19886
|
||||
"limit": 19885
|
||||
},
|
||||
"reportUnknownVariableType": {
|
||||
"limit": 30573
|
||||
"limit": 30572
|
||||
},
|
||||
"reportUnnecessaryCast": {
|
||||
"limit": 117
|
||||
|
|
|
|||
|
|
@ -1,5 +1,5 @@
|
|||
import json
|
||||
from collections.abc import Iterable, Iterator
|
||||
from collections.abc import Iterable, Iterator, Mapping
|
||||
from dataclasses import dataclass
|
||||
from typing import Any, Final, Literal
|
||||
|
||||
|
|
@ -87,7 +87,7 @@ async def _handle_completed_batch(
|
|||
return batch_cost, batch_usage, [model_name]
|
||||
|
||||
return _aggregate_batch_cost_usage_models(
|
||||
entries=_iter_batch_input_entries(file_content),
|
||||
entries=_iter_batch_output_entries(file_content),
|
||||
custom_llm_provider=custom_llm_provider,
|
||||
model_name=model_name,
|
||||
model_info=model_info,
|
||||
|
|
@ -111,43 +111,91 @@ def _iter_successful_output_line_stats(
|
|||
model_name: str | None,
|
||||
model_info: ModelInfo | None,
|
||||
) -> Iterator[_BatchOutputLineStats]:
|
||||
for entry in entries:
|
||||
stats = _safe_output_line_stats(entry, custom_llm_provider, model_name, model_info)
|
||||
if stats is not None:
|
||||
yield stats
|
||||
|
||||
|
||||
def _safe_output_line_stats(
|
||||
entry: Mapping[str, Any],
|
||||
custom_llm_provider: Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic", "bedrock"],
|
||||
model_name: str | None,
|
||||
model_info: ModelInfo | None,
|
||||
) -> _BatchOutputLineStats | None:
|
||||
"""Return the stats for one batch output line, or None for a line that is
|
||||
unsuccessful or cannot be costed, so a single bad line never aborts the
|
||||
whole batch's cost accounting."""
|
||||
custom_id: Final = entry.get("custom_id") if isinstance(entry, dict) else None
|
||||
try:
|
||||
if not _batch_response_was_successful(entry, custom_llm_provider):
|
||||
return None
|
||||
return _compute_output_line_stats(entry, custom_llm_provider, model_name, model_info)
|
||||
except Exception as e: # noqa: BLE001 # any single line's costing failure must not abort the whole batch
|
||||
verbose_logger.warning(
|
||||
"batch output line could not be costed, so it is billed at $0 and the rest of the batch "
|
||||
"is still billed. custom_id=%s error=%s",
|
||||
custom_id,
|
||||
str(e),
|
||||
)
|
||||
return None
|
||||
|
||||
|
||||
def _compute_output_line_stats(
|
||||
entry: Mapping[str, Any],
|
||||
custom_llm_provider: Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic", "bedrock"],
|
||||
model_name: str | None,
|
||||
model_info: ModelInfo | None,
|
||||
) -> _BatchOutputLineStats:
|
||||
response_body: Final = _get_response_from_batch_job_output_file(entry, custom_llm_provider)
|
||||
usage: Final = _get_batch_job_usage_from_response_body(response_body, custom_llm_provider)
|
||||
prompt_details: Final = parse_prompt_tokens_details(usage)
|
||||
raw_model: Final = response_body.get("model")
|
||||
response_model: Final = raw_model if isinstance(raw_model, str) and raw_model else None
|
||||
return _BatchOutputLineStats(
|
||||
cost=_output_line_cost(
|
||||
response_body=response_body,
|
||||
usage=usage,
|
||||
custom_llm_provider=custom_llm_provider,
|
||||
model_name=model_name,
|
||||
response_model=response_model,
|
||||
model_info=model_info,
|
||||
),
|
||||
prompt_tokens=usage.prompt_tokens,
|
||||
completion_tokens=usage.completion_tokens,
|
||||
total_tokens=usage.total_tokens,
|
||||
cache_read_tokens=prompt_details["cache_hit_tokens"],
|
||||
cache_creation_tokens=prompt_details["cache_creation_tokens"],
|
||||
model=response_model,
|
||||
)
|
||||
|
||||
|
||||
def _output_line_cost(
|
||||
response_body: Mapping[str, Any],
|
||||
usage: Usage,
|
||||
custom_llm_provider: Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic", "bedrock"],
|
||||
model_name: str | None,
|
||||
response_model: str | None,
|
||||
model_info: ModelInfo | None,
|
||||
) -> float:
|
||||
from litellm.cost_calculator import batch_cost_calculator
|
||||
|
||||
for entry in entries:
|
||||
if not _batch_response_was_successful(entry, custom_llm_provider):
|
||||
continue
|
||||
response_body = _get_response_from_batch_job_output_file(entry, custom_llm_provider)
|
||||
usage = _get_batch_job_usage_from_response_body(response_body, custom_llm_provider)
|
||||
prompt_details = parse_prompt_tokens_details(usage)
|
||||
raw_model = response_body.get("model")
|
||||
response_model = raw_model if isinstance(raw_model, str) and raw_model else None
|
||||
if model_info is not None or custom_llm_provider in ("anthropic", "bedrock"):
|
||||
if custom_llm_provider == "bedrock" and model_name:
|
||||
cost_model = model_name
|
||||
else:
|
||||
cost_model = response_model or model_name or ""
|
||||
prompt_cost, completion_cost = batch_cost_calculator(
|
||||
usage=usage,
|
||||
model=cost_model,
|
||||
custom_llm_provider=custom_llm_provider,
|
||||
model_info=model_info,
|
||||
)
|
||||
line_cost = prompt_cost + completion_cost
|
||||
else:
|
||||
line_cost = litellm.completion_cost(
|
||||
completion_response=response_body,
|
||||
custom_llm_provider=custom_llm_provider,
|
||||
call_type=CallTypes.aretrieve_batch.value,
|
||||
)
|
||||
yield _BatchOutputLineStats(
|
||||
cost=line_cost,
|
||||
prompt_tokens=usage.prompt_tokens,
|
||||
completion_tokens=usage.completion_tokens,
|
||||
total_tokens=usage.total_tokens,
|
||||
cache_read_tokens=prompt_details["cache_hit_tokens"],
|
||||
cache_creation_tokens=prompt_details["cache_creation_tokens"],
|
||||
model=response_model,
|
||||
if model_info is None and custom_llm_provider not in ("anthropic", "bedrock"):
|
||||
return litellm.completion_cost(
|
||||
completion_response=response_body,
|
||||
custom_llm_provider=custom_llm_provider,
|
||||
call_type=CallTypes.aretrieve_batch.value,
|
||||
)
|
||||
cost_model: Final = (
|
||||
model_name if custom_llm_provider == "bedrock" and model_name else response_model or model_name or ""
|
||||
)
|
||||
prompt_cost, completion_cost = batch_cost_calculator(
|
||||
usage=usage,
|
||||
model=cost_model,
|
||||
custom_llm_provider=custom_llm_provider,
|
||||
model_info=model_info,
|
||||
)
|
||||
return prompt_cost + completion_cost
|
||||
|
||||
|
||||
def _aggregate_batch_cost_usage_models(
|
||||
|
|
@ -338,9 +386,10 @@ def _extract_file_access_credentials(litellm_params: dict | None) -> dict:
|
|||
|
||||
def _get_file_content_as_dictionary(file_content: bytes) -> list[dict]:
|
||||
"""
|
||||
Get the file content as a list of dictionaries from JSON Lines format
|
||||
Get the file content as a list of dictionaries from JSON Lines format,
|
||||
skipping malformed lines
|
||||
"""
|
||||
return list(_iter_batch_input_entries(file_content))
|
||||
return list(_iter_batch_output_entries(file_content))
|
||||
|
||||
|
||||
def _iter_batch_input_lines(file_content: bytes) -> Iterator[bytes]:
|
||||
|
|
@ -361,15 +410,29 @@ def _iter_batch_input_lines(file_content: bytes) -> Iterator[bytes]:
|
|||
yield line
|
||||
|
||||
|
||||
def _iter_batch_input_entries(file_content: bytes) -> Iterator[dict]:
|
||||
def _iter_batch_output_entries(file_content: bytes) -> Iterator[dict]:
|
||||
"""
|
||||
Yield parsed batch input JSONL entries one at a time without materializing the
|
||||
whole file as a list, so peak memory stays bounded. Raises on a malformed line;
|
||||
callers that must survive bad rows should iterate ``_iter_batch_input_lines``
|
||||
and parse per-row instead.
|
||||
Yield parsed batch output JSONL entries one at a time without materializing
|
||||
the whole file as a list, so peak memory stays bounded. A malformed or
|
||||
non-object line is skipped with a warning so one bad line never aborts the
|
||||
whole batch's cost accounting.
|
||||
"""
|
||||
for line in _iter_batch_input_lines(file_content):
|
||||
yield json.loads(line)
|
||||
entry = _parse_batch_output_line(line)
|
||||
if entry is not None:
|
||||
yield entry
|
||||
|
||||
|
||||
def _parse_batch_output_line(line: bytes) -> dict | None:
|
||||
try:
|
||||
parsed: Final = json.loads(line)
|
||||
except ValueError as e:
|
||||
verbose_logger.warning("skipping malformed batch output line: %s", str(e))
|
||||
return None
|
||||
if isinstance(parsed, dict):
|
||||
return parsed
|
||||
verbose_logger.warning("skipping non-object batch output line of type %s", type(parsed).__name__)
|
||||
return None
|
||||
|
||||
|
||||
# A batch request's input tokens scale roughly with its serialized size, so this
|
||||
|
|
@ -440,7 +503,9 @@ def _count_prompt_or_input_tokens(model: str, value: Any) -> int:
|
|||
return 0
|
||||
|
||||
|
||||
def _get_batch_job_usage_from_response_body(response_body: dict, custom_llm_provider: str = "openai") -> Usage:
|
||||
def _get_batch_job_usage_from_response_body(
|
||||
response_body: Mapping[str, Any], custom_llm_provider: str = "openai"
|
||||
) -> Usage:
|
||||
"""
|
||||
Get the tokens of a batch job from the response body
|
||||
"""
|
||||
|
|
@ -472,7 +537,7 @@ def _get_batch_job_usage_from_response_body(response_body: dict, custom_llm_prov
|
|||
return usage
|
||||
|
||||
|
||||
def _get_anthropic_result_from_batch_results_line(batch_results_line: dict) -> dict:
|
||||
def _get_anthropic_result_from_batch_results_line(batch_results_line: Mapping[str, Any]) -> dict:
|
||||
"""
|
||||
Get the ``result`` object from a line of an Anthropic message batch results JSONL file.
|
||||
|
||||
|
|
@ -482,7 +547,9 @@ def _get_anthropic_result_from_batch_results_line(batch_results_line: dict) -> d
|
|||
return batch_results_line.get("result", None) or {}
|
||||
|
||||
|
||||
def _get_response_from_batch_job_output_file(batch_job_output_file: dict, custom_llm_provider: str = "openai") -> Any:
|
||||
def _get_response_from_batch_job_output_file(
|
||||
batch_job_output_file: Mapping[str, Any], custom_llm_provider: str = "openai"
|
||||
) -> Any:
|
||||
"""
|
||||
Get the response from the batch job output file
|
||||
"""
|
||||
|
|
@ -495,7 +562,9 @@ def _get_response_from_batch_job_output_file(batch_job_output_file: dict, custom
|
|||
return _response_body
|
||||
|
||||
|
||||
def _batch_response_was_successful(batch_job_output_file: dict, custom_llm_provider: str = "openai") -> bool:
|
||||
def _batch_response_was_successful(
|
||||
batch_job_output_file: Mapping[str, Any], custom_llm_provider: str = "openai"
|
||||
) -> bool:
|
||||
"""
|
||||
Check if the batch job response was successful
|
||||
|
||||
|
|
|
|||
|
|
@ -521,6 +521,20 @@ This is a one-time file patch and restore, not a live traffic interceptor. A Cla
|
|||
|
||||
Cursor is not supported: it has no equivalent file-based config to hot-patch this way, since its model routing lives in its own app storage and is configured through its GUI.
|
||||
|
||||
#### Making It Permanent at Login
|
||||
|
||||
`lite up` holds the patch only for as long as it runs. To wire Claude Code up once and leave it that way, pass `--config-claude` to `lite login`:
|
||||
|
||||
```bash
|
||||
lite --base-url https://your-proxy.example.com login --config-claude
|
||||
```
|
||||
|
||||
It writes the same two settings `lite up` does, `env.ANTHROPIC_BASE_URL` and `apiKeyHelper`, but persistently: there is no backup, nothing to restore, and no foreground process to keep alive. Every other key in `~/.claude/settings.json` is preserved, the file is created if it does not exist, and it is written atomically with owner-only permissions. Plain `lite login` is unchanged; nothing happens to your Claude Code config unless you pass the flag.
|
||||
|
||||
Because the credential is reached through `apiKeyHelper` rather than copied into the file, a later `lite login` refreshes it with no further action: Claude Code re-runs the helper on every request and picks up whatever token the most recent login stored. Nothing secret is written to `settings.json`.
|
||||
|
||||
Run it again to point Claude Code at a different proxy; the base URL and the helper are both rewritten. `lite up` and `--config-claude` manage the same file, so the flag refuses to run while a `lite up` session holds a backup, and tells you to run `lite down` first, rather than writing settings that `lite up` would silently revert when it stops.
|
||||
|
||||
### QA Complexity-Based Auto-Routing Against Your Real Proxy
|
||||
|
||||
`lite autoroute` lets you try LiteLLM's complexity-based auto-routing -- picking a cheaper or more expensive model depending on how complex a prompt looks -- against models your key already has access to on your real, running proxy, without editing that proxy's `config.yaml` and without any real request ever bypassing it. It builds a second, throwaway proxy locally that forwards every request back to your real proxy, and points Claude Code at that local proxy for the duration of the session.
|
||||
|
|
|
|||
|
|
@ -16,6 +16,12 @@ from typing_extensions import NotRequired, TypedDict
|
|||
from litellm.constants import CLI_JWT_EXPIRATION_HOURS
|
||||
from litellm.litellm_core_utils.cli_token_utils import is_cli_token_fresh
|
||||
|
||||
from .claude_settings import (
|
||||
CLAUDE_SETTINGS_PATH,
|
||||
SETTINGS_FILE_OWNERS,
|
||||
ClaudeSettingsError,
|
||||
write_claude_settings,
|
||||
)
|
||||
from .private_json import write_private_json
|
||||
|
||||
|
||||
|
|
@ -629,9 +635,28 @@ def _render_and_prompt_for_team_selection(teams: list[CliTeam]) -> str | None:
|
|||
return None
|
||||
|
||||
|
||||
def _configure_claude_code(base_url: str) -> None:
|
||||
"""Point Claude Code at base_url by patching ~/.claude/settings.json."""
|
||||
try:
|
||||
write_claude_settings(base_url, CLAUDE_SETTINGS_PATH, SETTINGS_FILE_OWNERS)
|
||||
except ClaudeSettingsError as e:
|
||||
raise click.ClickException(f"Logged in, but could not configure Claude Code: {e}")
|
||||
click.echo(f"\nConfigured Claude Code: {CLAUDE_SETTINGS_PATH} now routes through {base_url.rstrip('/')}.")
|
||||
click.echo("Your other Claude Code settings were left untouched. Restart Claude Code to pick this up.")
|
||||
|
||||
|
||||
@click.command(name="login")
|
||||
@click.option(
|
||||
"--config-claude",
|
||||
is_flag=True,
|
||||
default=False,
|
||||
help=(
|
||||
"After logging in, update ~/.claude/settings.json so Claude Code routes through this proxy. "
|
||||
"Unrelated settings are preserved."
|
||||
),
|
||||
)
|
||||
@click.pass_context
|
||||
def login(ctx: click.Context):
|
||||
def login(ctx: click.Context, config_claude: bool):
|
||||
"""Login to LiteLLM proxy using SSO authentication"""
|
||||
from litellm.constants import LITELLM_CLI_SOURCE_IDENTIFIER
|
||||
from litellm.proxy.client.cli.interface import show_commands
|
||||
|
|
@ -683,6 +708,9 @@ def login(ctx: click.Context):
|
|||
click.echo(f"JWT Token: {api_key[:20]}...")
|
||||
click.echo("You can now use the CLI without specifying --api-key")
|
||||
|
||||
if config_claude:
|
||||
_configure_claude_code(base_url)
|
||||
|
||||
# Show available commands after successful login
|
||||
click.echo("\n" + "=" * 60)
|
||||
show_commands()
|
||||
|
|
@ -698,6 +726,10 @@ def login(ctx: click.Context):
|
|||
except KeyboardInterrupt:
|
||||
click.echo("\nAuthentication cancelled by user.")
|
||||
return
|
||||
except click.ClickException:
|
||||
# Login itself already succeeded; only the post-login step failed, so this
|
||||
# must not be relabelled as an authentication failure by the handler below.
|
||||
raise
|
||||
except Exception as e:
|
||||
click.echo(f"Authentication failed: {e}")
|
||||
return
|
||||
|
|
|
|||
|
|
@ -10,11 +10,16 @@ import click
|
|||
import yaml
|
||||
from pydantic import JsonValue, TypeAdapter, ValidationError
|
||||
|
||||
from ..up import CLAUDE_SETTINGS_PATH, UpError, load_json_or_empty, restore_claude_settings, write_backup
|
||||
from ..claude_settings import (
|
||||
AUTOROUTE_BACKUP_PATH,
|
||||
CLAUDE_SETTINGS_PATH,
|
||||
ClaudeSettingsError,
|
||||
load_json_or_empty,
|
||||
)
|
||||
from ..up import BackupRecord as ClaudeBackupRecord
|
||||
from ..up import restore_claude_settings, write_backup
|
||||
from .config import master_key_from_config
|
||||
from .process import (
|
||||
AUTOROUTE_DIR,
|
||||
CONFIG_PATH,
|
||||
DEFAULT_AUTOROUTE_PORT,
|
||||
LOG_PATH,
|
||||
|
|
@ -35,8 +40,6 @@ from .process import (
|
|||
from .settings import merge_claude_settings_static_token
|
||||
from .wizard import run_configure_wizard
|
||||
|
||||
AUTOROUTE_BACKUP_PATH: Final = AUTOROUTE_DIR / "claude_settings_backup.json"
|
||||
|
||||
_GENERATED_CONFIG_ADAPTER: Final = TypeAdapter(dict[str, JsonValue])
|
||||
|
||||
|
||||
|
|
@ -108,7 +111,7 @@ def up(port: int) -> None:
|
|||
|
||||
try:
|
||||
existing_pid: Final = read_pid_record()
|
||||
except UpError as e:
|
||||
except ClaudeSettingsError as e:
|
||||
raise click.ClickException(str(e))
|
||||
if existing_pid is not None and is_running(existing_pid.pid):
|
||||
raise click.ClickException(
|
||||
|
|
@ -157,7 +160,7 @@ def up(port: int) -> None:
|
|||
CLAUDE_SETTINGS_PATH.parent.mkdir(parents=True, exist_ok=True)
|
||||
with secure_create(CLAUDE_SETTINGS_PATH) as f:
|
||||
json.dump(merged, f, indent=2)
|
||||
except UpError as e:
|
||||
except ClaudeSettingsError as e:
|
||||
terminate(process.pid)
|
||||
clear_pid_record()
|
||||
raise click.ClickException(str(e))
|
||||
|
|
@ -175,7 +178,7 @@ def up(port: int) -> None:
|
|||
clear_pid_record()
|
||||
try:
|
||||
restore_claude_settings(CLAUDE_SETTINGS_PATH, AUTOROUTE_BACKUP_PATH)
|
||||
except UpError as e:
|
||||
except ClaudeSettingsError as e:
|
||||
# Runs from atexit/a signal handler too, outside Click's own exception
|
||||
# handling -- raising here would only produce an unhandled-exception
|
||||
# warning on stderr, not a clean message.
|
||||
|
|
@ -207,7 +210,7 @@ def down() -> None:
|
|||
"""Restore Claude Code settings and stop a leftover ephemeral proxy, if any"""
|
||||
try:
|
||||
record: PidRecord | None = read_pid_record()
|
||||
except UpError as e:
|
||||
except ClaudeSettingsError as e:
|
||||
# down is the crash-recovery path -- a corrupt pid record must not block it; clear the
|
||||
# unusable record and keep going rather than leaving the user with no way to clean up.
|
||||
click.echo(f"{e} Clearing it and continuing cleanup.", err=True)
|
||||
|
|
@ -219,7 +222,7 @@ def down() -> None:
|
|||
|
||||
try:
|
||||
restored: Final = restore_claude_settings(CLAUDE_SETTINGS_PATH, AUTOROUTE_BACKUP_PATH)
|
||||
except UpError as e:
|
||||
except ClaudeSettingsError as e:
|
||||
raise click.ClickException(str(e))
|
||||
if restored is None:
|
||||
click.echo("Nothing to restore.")
|
||||
|
|
|
|||
155
litellm/proxy/client/cli/commands/claude_settings.py
Normal file
155
litellm/proxy/client/cli/commands/claude_settings.py
Normal file
|
|
@ -0,0 +1,155 @@
|
|||
"""Shared handling of Claude Code's ~/.claude/settings.json.
|
||||
|
||||
`lite up` patches this file temporarily and restores it on exit; `lite login
|
||||
--config-claude` patches it persistently. Both need the same merge and the same
|
||||
apiKeyHelper command, and `up` already imports from `auth`, so the shared parts
|
||||
live here rather than in either command module.
|
||||
"""
|
||||
|
||||
import shlex
|
||||
import shutil
|
||||
from collections.abc import Mapping, Sequence
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
from typing import Final
|
||||
|
||||
from pydantic import JsonValue, TypeAdapter, ValidationError
|
||||
|
||||
from .private_json import write_private_json
|
||||
|
||||
ENV_KEY: Final = "env"
|
||||
API_KEY_HELPER_KEY: Final = "apiKeyHelper"
|
||||
ANTHROPIC_BASE_URL_KEY: Final = "ANTHROPIC_BASE_URL"
|
||||
ANTHROPIC_API_KEY_KEY: Final = "ANTHROPIC_API_KEY"
|
||||
|
||||
CLAUDE_SETTINGS_PATH: Final = Path.home() / ".claude" / "settings.json"
|
||||
BACKUP_PATH: Final = Path.home() / ".litellm" / "claude_settings_backup.json"
|
||||
AUTOROUTE_BACKUP_PATH: Final = Path.home() / ".litellm" / "autorouter" / "claude_settings_backup.json"
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class SettingsFileOwner:
|
||||
"""A command that takes temporary ownership of CLAUDE_SETTINGS_PATH and restores it later."""
|
||||
|
||||
backup_path: Path
|
||||
start_command: str
|
||||
stop_command: str
|
||||
|
||||
|
||||
SETTINGS_FILE_OWNERS: Final = (
|
||||
SettingsFileOwner(BACKUP_PATH, "lite up", "lite down"),
|
||||
SettingsFileOwner(AUTOROUTE_BACKUP_PATH, "lite autoroute up", "lite autoroute down"),
|
||||
)
|
||||
|
||||
_SETTINGS_ADAPTER: Final = TypeAdapter(dict[str, JsonValue])
|
||||
|
||||
|
||||
class ClaudeSettingsError(Exception):
|
||||
"""Raised for any user-actionable failure while reading or writing Claude Code settings."""
|
||||
|
||||
|
||||
def load_json_or_empty(path: Path) -> dict[str, JsonValue]:
|
||||
try:
|
||||
content: Final = path.read_bytes() if path.exists() else b""
|
||||
except OSError as e:
|
||||
raise ClaudeSettingsError(f"Could not read {path}: {e}") from e
|
||||
if not content.strip():
|
||||
return {}
|
||||
try:
|
||||
return _SETTINGS_ADAPTER.validate_json(content)
|
||||
except ValidationError:
|
||||
raise ClaudeSettingsError(
|
||||
f"{path} contains invalid JSON (or its root is not an object); cannot proceed safely."
|
||||
)
|
||||
|
||||
|
||||
def merge_claude_settings(
|
||||
settings: Mapping[str, JsonValue], base_url: str, api_key_helper: str
|
||||
) -> dict[str, JsonValue]:
|
||||
"""Return a new settings dict wired to route Claude Code through the proxy.
|
||||
|
||||
Only env.ANTHROPIC_BASE_URL and the top-level apiKeyHelper are overridden; a
|
||||
stray env.ANTHROPIC_API_KEY is dropped so it cannot outrank the helper-issued
|
||||
token (same reasoning as build_agent_env in agents.py). Every other key is
|
||||
preserved untouched.
|
||||
"""
|
||||
raw_env: Final = settings.get(ENV_KEY, {})
|
||||
base_env: Final = raw_env if isinstance(raw_env, dict) else {}
|
||||
env: Final = {
|
||||
**{key: value for key, value in base_env.items() if key != ANTHROPIC_API_KEY_KEY},
|
||||
ANTHROPIC_BASE_URL_KEY: base_url.rstrip("/"),
|
||||
}
|
||||
return {**settings, ENV_KEY: env, API_KEY_HELPER_KEY: api_key_helper}
|
||||
|
||||
|
||||
def resolve_api_key_helper(base_url: str) -> str:
|
||||
"""Build the shell command Claude Code should run for its apiKeyHelper.
|
||||
|
||||
Resolves `lite` to an absolute path so the helper works regardless of the
|
||||
PATH visible to whatever subprocess Claude Code spawns it from. Passing
|
||||
--base-url explicitly (rather than relying on the bare invocation Claude
|
||||
Code would otherwise use) makes `print-token` enforce that the cached
|
||||
token was actually issued for this proxy -- without it, a token minted
|
||||
for a different, previously-logged-into proxy would be handed to
|
||||
whichever server the settings currently point at.
|
||||
|
||||
--base-url belongs to the top-level `lite` group, so it has to precede the
|
||||
subcommand; click rejects it outright after `print-token`.
|
||||
"""
|
||||
lite_path: Final = shutil.which("lite")
|
||||
if lite_path is None:
|
||||
raise ClaudeSettingsError(
|
||||
"Could not find `lite` on your PATH. Claude Code's apiKeyHelper needs an absolute path to it."
|
||||
)
|
||||
return f"{shlex.quote(lite_path)} --base-url {shlex.quote(base_url)} auth print-token"
|
||||
|
||||
|
||||
def write_claude_settings(base_url: str, settings_path: Path, owners: Sequence[SettingsFileOwner]) -> None:
|
||||
"""Persistently point Claude Code at base_url, preserving every unrelated setting.
|
||||
|
||||
Refuses while any owner holds a backup: each restores its backup when it
|
||||
stops, which would silently undo this write.
|
||||
"""
|
||||
for owner in owners:
|
||||
if owner.backup_path.exists():
|
||||
raise ClaudeSettingsError(
|
||||
f"`{owner.start_command}` is currently managing {settings_path} (backup at "
|
||||
f"{owner.backup_path}) and will restore it when it stops. "
|
||||
f"Run `{owner.stop_command}` first, then retry."
|
||||
)
|
||||
normalized_base_url: Final = base_url.rstrip("/")
|
||||
api_key_helper: Final = resolve_api_key_helper(normalized_base_url)
|
||||
existing: Final = load_json_or_empty(settings_path)
|
||||
raw_env: Final = existing.get(ENV_KEY)
|
||||
if raw_env is not None and not isinstance(raw_env, dict):
|
||||
raise ClaudeSettingsError(
|
||||
f'{settings_path} has a non-object "{ENV_KEY}" value, which this would discard. '
|
||||
"Fix or remove it, then retry."
|
||||
)
|
||||
merged: Final = merge_claude_settings(existing, normalized_base_url, api_key_helper)
|
||||
# os.replace() swaps the symlink itself for a regular file, silently detaching a
|
||||
# settings.json that is symlinked into a dotfiles repo. There is no backup to undo
|
||||
# that here, unlike `lite up`, so write through to the link's target instead.
|
||||
target: Final = settings_path.resolve() if settings_path.is_symlink() else settings_path
|
||||
try:
|
||||
write_private_json(str(target), merged)
|
||||
except OSError as e:
|
||||
raise ClaudeSettingsError(f"Could not write {target}: {e}") from e
|
||||
|
||||
|
||||
__all__ = (
|
||||
"ANTHROPIC_API_KEY_KEY",
|
||||
"ANTHROPIC_BASE_URL_KEY",
|
||||
"API_KEY_HELPER_KEY",
|
||||
"AUTOROUTE_BACKUP_PATH",
|
||||
"BACKUP_PATH",
|
||||
"CLAUDE_SETTINGS_PATH",
|
||||
"ENV_KEY",
|
||||
"SETTINGS_FILE_OWNERS",
|
||||
"ClaudeSettingsError",
|
||||
"SettingsFileOwner",
|
||||
"load_json_or_empty",
|
||||
"merge_claude_settings",
|
||||
"resolve_api_key_helper",
|
||||
"write_claude_settings",
|
||||
)
|
||||
|
|
@ -2,12 +2,10 @@ import atexit
|
|||
import contextlib
|
||||
import json
|
||||
import os
|
||||
import shlex
|
||||
import shutil
|
||||
import signal
|
||||
import sys
|
||||
import threading
|
||||
from collections.abc import Iterator, Mapping
|
||||
from collections.abc import Iterator
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
from types import FrameType
|
||||
|
|
@ -20,17 +18,17 @@ from litellm.litellm_core_utils.cli_token_utils import is_cli_token_fresh
|
|||
|
||||
from .agents import AgentRunError, resolve_api_key, verify_proxy_key
|
||||
from .auth import load_token, login
|
||||
|
||||
ENV_KEY: Final = "env"
|
||||
API_KEY_HELPER_KEY: Final = "apiKeyHelper"
|
||||
ANTHROPIC_BASE_URL_KEY: Final = "ANTHROPIC_BASE_URL"
|
||||
ANTHROPIC_API_KEY_KEY: Final = "ANTHROPIC_API_KEY"
|
||||
|
||||
CLAUDE_SETTINGS_PATH: Final = Path.home() / ".claude" / "settings.json"
|
||||
BACKUP_PATH: Final = Path.home() / ".litellm" / "claude_settings_backup.json"
|
||||
from .claude_settings import (
|
||||
BACKUP_PATH,
|
||||
CLAUDE_SETTINGS_PATH,
|
||||
ClaudeSettingsError,
|
||||
load_json_or_empty,
|
||||
merge_claude_settings,
|
||||
resolve_api_key_helper,
|
||||
)
|
||||
|
||||
|
||||
class UpError(Exception):
|
||||
class UpError(ClaudeSettingsError):
|
||||
"""Raised for any user-actionable failure while starting/stopping interception."""
|
||||
|
||||
|
||||
|
|
@ -42,40 +40,9 @@ class BackupRecord:
|
|||
content: dict[str, JsonValue] | None
|
||||
|
||||
|
||||
_SETTINGS_ADAPTER: Final = TypeAdapter(dict[str, JsonValue])
|
||||
_BACKUP_RECORD_ADAPTER: Final = TypeAdapter(BackupRecord)
|
||||
|
||||
|
||||
def load_json_or_empty(path: Path) -> dict[str, JsonValue]:
|
||||
if not path.exists():
|
||||
return {}
|
||||
with open(path, "r") as f:
|
||||
content: Final = f.read()
|
||||
if not content.strip():
|
||||
return {}
|
||||
try:
|
||||
return _SETTINGS_ADAPTER.validate_json(content)
|
||||
except ValidationError:
|
||||
raise UpError(f"{path} contains invalid JSON (or its root is not an object); cannot proceed safely.")
|
||||
|
||||
|
||||
def merge_claude_settings(
|
||||
settings: Mapping[str, JsonValue], base_url: str, api_key_helper: str
|
||||
) -> dict[str, JsonValue]:
|
||||
"""Return a new settings dict wired to route Claude Code through the proxy.
|
||||
|
||||
Only env.ANTHROPIC_BASE_URL and the top-level apiKeyHelper are overridden; a
|
||||
stray env.ANTHROPIC_API_KEY is dropped so it cannot outrank the helper-issued
|
||||
token (same reasoning as build_agent_env in agents.py). Every other key is
|
||||
preserved untouched.
|
||||
"""
|
||||
raw_env: Final = settings.get(ENV_KEY, {})
|
||||
base_env: Final = raw_env if isinstance(raw_env, dict) else {}
|
||||
env: Final = {**base_env, ANTHROPIC_BASE_URL_KEY: base_url.rstrip("/")}
|
||||
env.pop(ANTHROPIC_API_KEY_KEY, None)
|
||||
return {**settings, ENV_KEY: env, API_KEY_HELPER_KEY: api_key_helper}
|
||||
|
||||
|
||||
@contextlib.contextmanager
|
||||
def secure_create(path: Path) -> Iterator[IO[str]]:
|
||||
"""Open path for writing with mode 0600 fixed up before any content is written.
|
||||
|
|
@ -136,26 +103,6 @@ def restore_claude_settings(settings_path: Path | None = None, backup_path: Path
|
|||
return record
|
||||
|
||||
|
||||
def resolve_api_key_helper(base_url: str) -> str:
|
||||
"""Build the shell command Claude Code should run for its apiKeyHelper.
|
||||
|
||||
Resolves `lite` to an absolute path so the helper works regardless of the
|
||||
PATH visible to whatever subprocess Claude Code spawns it from. Passing
|
||||
--base-url explicitly (rather than relying on the bare invocation Claude
|
||||
Code would otherwise use) makes `print-token` enforce that the cached
|
||||
token was actually issued for this proxy -- without it, a token minted
|
||||
for a different, previously-logged-into proxy would be handed to
|
||||
whichever server `up` currently points at.
|
||||
"""
|
||||
lite_path: Final = shutil.which("lite")
|
||||
if lite_path is None:
|
||||
raise UpError(
|
||||
"Could not find `lite` on your PATH. Claude Code's apiKeyHelper needs "
|
||||
"an absolute path to it, so `lite up` cannot continue."
|
||||
)
|
||||
return f"{shlex.quote(lite_path)} auth print-token --base-url {shlex.quote(base_url)}"
|
||||
|
||||
|
||||
def _ensure_fresh_login(ctx: click.Context) -> None:
|
||||
base_url: Final = ctx.obj["base_url"].rstrip("/")
|
||||
token_data = load_token()
|
||||
|
|
@ -224,7 +171,7 @@ def up(ctx: click.Context) -> None:
|
|||
merged: Final = merge_claude_settings(original_settings, base_url, api_key_helper)
|
||||
with open(CLAUDE_SETTINGS_PATH, "w") as f:
|
||||
json.dump(merged, f, indent=2)
|
||||
except (AgentRunError, UpError) as e:
|
||||
except (AgentRunError, ClaudeSettingsError) as e:
|
||||
raise click.ClickException(str(e))
|
||||
|
||||
click.echo(f"litellm: routing Claude Code through proxy at {base_url.rstrip('/')}")
|
||||
|
|
@ -241,7 +188,7 @@ def up(ctx: click.Context) -> None:
|
|||
return
|
||||
try:
|
||||
_restore_and_report()
|
||||
except UpError as e:
|
||||
except ClaudeSettingsError as e:
|
||||
# Runs from atexit/a signal handler, outside Click's own exception
|
||||
# handling -- raising here would only produce an unhandled-exception
|
||||
# warning on stderr, not a clean message.
|
||||
|
|
@ -264,7 +211,7 @@ def down() -> None:
|
|||
"""
|
||||
try:
|
||||
_restore_and_report()
|
||||
except UpError as e:
|
||||
except ClaudeSettingsError as e:
|
||||
raise click.ClickException(str(e))
|
||||
|
||||
|
||||
|
|
@ -272,6 +219,7 @@ __all__ = [
|
|||
"BACKUP_PATH",
|
||||
"CLAUDE_SETTINGS_PATH",
|
||||
"BackupRecord",
|
||||
"ClaudeSettingsError",
|
||||
"UpError",
|
||||
"down",
|
||||
"load_json_or_empty",
|
||||
|
|
|
|||
|
|
@ -10253,11 +10253,13 @@ class Router:
|
|||
returned_models.extend(self.get_model_list_from_routing_groups(model_name=model_name))
|
||||
|
||||
if len(returned_models) == 0: # check if wildcard route
|
||||
potential_wildcard_models: Final = self.pattern_router.route(model_name) or []
|
||||
potential_wildcard_models: Final = self.pattern_router.get_deployments_by_pattern(model=model_name or "")
|
||||
|
||||
## check for team-specific wildcard models
|
||||
if team_id is not None and team_id in self.team_pattern_routers:
|
||||
potential_team_only_wildcard_models: Final = self.team_pattern_routers[team_id].route(model_name) or []
|
||||
potential_team_only_wildcard_models: Final = self.team_pattern_routers[
|
||||
team_id
|
||||
].get_deployments_by_pattern(model=model_name or "")
|
||||
potential_wildcard_models.extend(potential_team_only_wildcard_models)
|
||||
|
||||
if model_name is not None and potential_wildcard_models is not None:
|
||||
|
|
|
|||
|
|
@ -75,11 +75,11 @@ Mark live tests with `@pytest.mark.e2e` (on the class or the module). Pure cover
|
|||
|
||||
`E2E_FIXTURE_MODE` selects the transport every client is built on: `live` (the default, and what an unset variable means: nothing changes), `record` (run against the live proxy and write every interaction to a fixture bundle), or `replay` (serve every interaction back from the bundle with no HTTP at all, so a replay run needs no proxy and cannot bill a provider). The seam is `select_transport` in `fixture_transport.py`, applied inside `build_proxy_client`; both transports fulfil the same `Transport` protocol, so no test or client changes shape in any mode
|
||||
|
||||
A bundle (default `tests/e2e/.fixtures`, override with `E2E_FIXTURE_DIR`) is a directory: `manifest.json` carries the record timestamp, harness git version, and format version, and each test gets a subdirectory holding one JSON file per transport call in call order (`0000-post-chat-completions.json`). Auth header values are redacted on write, and file uploads store a sha256 digest instead of the bytes; response bodies are stored verbatim (a /key/generate response keeps the ephemeral virtual key it minted), which is part of why bundles are gitignored. `fixture_bundle.py` owns the format
|
||||
A bundle (default `tests/e2e/.fixtures`, override with `E2E_FIXTURE_DIR`) is a directory: `manifest.json` carries the record timestamp, harness git version, and format version, and each test gets a subdirectory holding one JSON file per transport call in call order (`0000-post-chat-completions.json`). Auth header values and credential request fields (`api_key`, `*_secret_key`, `static_headers`, and the like; the list is `fixture_canonical.py`'s) are redacted on write, and file uploads store a sha256 digest instead of the bytes; response bodies are stored verbatim (a /key/generate response keeps the ephemeral virtual key it minted), which is part of why bundles are gitignored. `fixture_bundle.py` owns the format
|
||||
|
||||
Replay matches calls per test by transport verb and path in recorded order and raises `ReplayMiss` on any drift, naming the recorded and the actual call; a passed test must also consume its whole recording, or teardown fails it naming the first leftover interaction. Either way the fix is always to re-record with `E2E_FIXTURE_MODE=record`. Record starts fresh every time: it wipes the previous bundle (refusing to wipe a directory that is not a bundle) and never reads it. A replay bundle whose manifest is older than seven days hard-fails at collection time naming the bundle's age, so replay can never certify against fixtures that have drifted more than a week from the live proxy
|
||||
Replay matches calls per test by canonical key: `fixture_canonical.py` canonicalizes the recorded request (volatile headers and credential fields out, unique markers, generated ids, uuids, and timestamps replaced with fixed placeholders, object keys sorted) and the key is the method, path, and a content hash, so identity survives re-records and machine changes while any real content drift is a `ReplayMiss` that names the computed key, the closest recorded key with its file, and a content diff, and never falls through to a live call. Matching is order-independent across distinct keys (concurrent calls may interleave) and FIFO within one key (a poll loop replays its responses in recorded order); a passed test must also consume its whole recording, or teardown fails it naming a leftover key. Either way the fix is always to re-record with `E2E_FIXTURE_MODE=record`. Every rewrite rule lives in `fixture_canonical.py`, so a new volatile header, credential field name, or generated-id shape is one edit there. Record starts fresh every time: it wipes the previous bundle (refusing to wipe a directory that is not a bundle) and never reads it. A replay bundle whose manifest is older than seven days hard-fails at collection time naming the bundle's age, so replay can never certify against fixtures that have drifted more than a week from the live proxy
|
||||
|
||||
Deliberately not here yet: canonical content-based match keys (LIT-5741), streaming chunk fidelity (LIT-5742), and scoping record/replay to provider-bound traffic (LIT-5745)
|
||||
Deliberately not here yet: streaming chunk fidelity (LIT-5742) and scoping record/replay to provider-bound traffic (LIT-5745)
|
||||
|
||||
## Typing
|
||||
|
||||
|
|
|
|||
|
|
@ -2,12 +2,13 @@
|
|||
|
||||
from __future__ import annotations
|
||||
|
||||
import time
|
||||
from dataclasses import dataclass
|
||||
|
||||
from pydantic import BaseModel, ValidationError
|
||||
|
||||
from proxy_client import ProxyClient
|
||||
from e2e_http import StreamingResponse
|
||||
from e2e_http import NoBody, StreamingResponse, is_ok, unwrap
|
||||
from models import (
|
||||
ChatBody,
|
||||
ChatMessage,
|
||||
|
|
@ -15,9 +16,16 @@ from models import (
|
|||
LiteLLMParamsBody,
|
||||
ModelInfoBody,
|
||||
ModelNewBody,
|
||||
TeamDeleteBody,
|
||||
TeamInfoParams,
|
||||
TeamInfoResponse,
|
||||
TeamNewBody,
|
||||
TeamNewResponse,
|
||||
TeamUpdateBody,
|
||||
)
|
||||
|
||||
MODEL_ACCESS_DENIED_MARKER = "key_model_access_denied"
|
||||
TEAM_MODEL_ACCESS_DENIED_MARKER = "team_model_access_denied"
|
||||
ROUTE_NOT_ALLOWED_MARKER = "not allowed to call this route"
|
||||
|
||||
|
||||
|
|
@ -31,6 +39,14 @@ class ApiErrorEnvelope(BaseModel):
|
|||
error: ApiErrorDetail
|
||||
|
||||
|
||||
class AccessGroupInfoResponse(BaseModel):
|
||||
"""GET /access_group/{name}/info: the deployments a model access group grants."""
|
||||
|
||||
access_group: str
|
||||
model_names: list[str]
|
||||
deployment_count: int
|
||||
|
||||
|
||||
def error_envelope(body: str) -> ApiErrorEnvelope | None:
|
||||
"""The OpenAI-shaped `{"error": {...}}` a client parses, or None if absent."""
|
||||
try:
|
||||
|
|
@ -51,15 +67,75 @@ class AccessControlClient:
|
|||
def delete_key(self, key: str) -> None:
|
||||
self.proxy.delete_key(key)
|
||||
|
||||
def chat_status(self, key: str, model: str, content: str) -> StreamingResponse:
|
||||
def chat_status(
|
||||
self, key: str, model: str, content: str, max_completion_tokens: int | None = None
|
||||
) -> StreamingResponse:
|
||||
return self.proxy.transport.send(
|
||||
"/chat/completions",
|
||||
headers=self.proxy.transport.bearer(key),
|
||||
json=ChatBody(
|
||||
model=model, messages=[ChatMessage(role="user", content=content)]
|
||||
model=model,
|
||||
messages=[ChatMessage(role="user", content=content)],
|
||||
max_completion_tokens=max_completion_tokens,
|
||||
),
|
||||
)
|
||||
|
||||
def create_team(self, team_alias: str, models: list[str]) -> str:
|
||||
team_id = unwrap(
|
||||
self.proxy.transport.post(
|
||||
"/team/new",
|
||||
headers=self.proxy.transport.master,
|
||||
json=TeamNewBody(team_alias=team_alias, models=models),
|
||||
response_type=TeamNewResponse,
|
||||
)
|
||||
).team_id
|
||||
self._await_team(team_id)
|
||||
return team_id
|
||||
|
||||
def set_team_models(self, team_id: str, team_alias: str, models: list[str]) -> None:
|
||||
"""Replace the team's allow-list. /model/new appends a team-scoped deployment's
|
||||
public name to it, so a test that means to grant only an access group has to
|
||||
put the allow-list back afterwards."""
|
||||
_ = unwrap(
|
||||
self.proxy.transport.post(
|
||||
"/team/update",
|
||||
headers=self.proxy.transport.master,
|
||||
json=TeamUpdateBody(team_id=team_id, team_alias=team_alias, models=models),
|
||||
response_type=NoBody,
|
||||
)
|
||||
)
|
||||
|
||||
def delete_team(self, team_id: str) -> None:
|
||||
_ = self.proxy.transport.post(
|
||||
"/team/delete",
|
||||
headers=self.proxy.transport.master,
|
||||
json=TeamDeleteBody(team_ids=[team_id]),
|
||||
response_type=NoBody,
|
||||
)
|
||||
|
||||
def access_group_info(self, access_group: str) -> AccessGroupInfoResponse | None:
|
||||
result = self.proxy.transport.get(
|
||||
f"/access_group/{access_group}/info",
|
||||
headers=self.proxy.transport.master,
|
||||
params=NoBody(),
|
||||
response_type=AccessGroupInfoResponse,
|
||||
)
|
||||
return unwrap(result) if is_ok(result) else None
|
||||
|
||||
def _await_team(self, team_id: str) -> None:
|
||||
deadline = time.monotonic() + self.proxy.poll_timeout
|
||||
while time.monotonic() < deadline:
|
||||
result = self.proxy.transport.get(
|
||||
"/team/info",
|
||||
headers=self.proxy.transport.master,
|
||||
params=TeamInfoParams(team_id=team_id),
|
||||
response_type=TeamInfoResponse,
|
||||
)
|
||||
if is_ok(result):
|
||||
return
|
||||
time.sleep(self.proxy.poll_interval)
|
||||
raise AssertionError(f"/team/info never resolved team {team_id!r} created by /team/new")
|
||||
|
||||
def create_model_status(self, key: str, model_name: str) -> StreamingResponse:
|
||||
return self.proxy.transport.send(
|
||||
"/model/new",
|
||||
|
|
|
|||
277
tests/e2e/access_control/test_model_access_group_e2e.py
Normal file
277
tests/e2e/access_control/test_model_access_group_e2e.py
Normal file
|
|
@ -0,0 +1,277 @@
|
|||
"""Live e2e: a model access group as the grant on a key and on a team.
|
||||
|
||||
Whoever holds the group can call every deployment in it and nothing else, whether
|
||||
the request names a deployment exactly, names a model that a wildcard deployment
|
||||
in the group covers, or spells that model with its provider prefix. The bare-name
|
||||
spelling is the LIT-5813 regression: the group-membership lookup skipped the
|
||||
provider-prefix retry every other model-resolution path performs, so a group
|
||||
holding `openai/gpt-5.4*` denied `gpt-5.4-nano` while allowing `openai/gpt-5.4-nano`.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
import time
|
||||
from collections.abc import Callable, Iterator
|
||||
from dataclasses import dataclass
|
||||
from typing import Final
|
||||
|
||||
import pytest
|
||||
|
||||
from access_control_client import (
|
||||
AccessControlClient,
|
||||
MODEL_ACCESS_DENIED_MARKER,
|
||||
TEAM_MODEL_ACCESS_DENIED_MARKER,
|
||||
)
|
||||
from e2e_config import unique_marker
|
||||
from lifecycle import ResourceManager
|
||||
from models import (
|
||||
ChatResponse,
|
||||
KeyGenerateBody,
|
||||
LiteLLMParamsBody,
|
||||
ModelInfoBody,
|
||||
ModelNewBody,
|
||||
)
|
||||
|
||||
pytestmark = pytest.mark.e2e
|
||||
|
||||
WILDCARD_PATTERN: Final = "openai/gpt-5.4*"
|
||||
WILDCARD_BARE_MODEL: Final = "gpt-5.4-nano"
|
||||
WILDCARD_PREFIXED_MODEL: Final = "openai/gpt-5.4-nano"
|
||||
GROUP_BACKEND: Final = "openai/gpt-5.4-nano"
|
||||
UNCOVERED_OPENAI_MODEL: Final = "gpt-5.2"
|
||||
|
||||
TEAM_WILDCARD_PATTERN: Final = "openai/gpt-5.6*"
|
||||
TEAM_WILDCARD_BARE_MODEL: Final = "gpt-5.6-luna"
|
||||
|
||||
MAX_COMPLETION_TOKENS: Final = 16
|
||||
PROMPT: Final = "Reply with exactly: OK"
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class GroupedDeployments:
|
||||
"""A wildcard deployment and an exactly-named one inside `access_group`, plus a
|
||||
deployment left out of it."""
|
||||
|
||||
access_group: str
|
||||
member_model: str
|
||||
outsider_model: str
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class TeamGrant:
|
||||
"""A team whose whole allow-list is `access_group`, holding one team-scoped
|
||||
wildcard deployment, and a key that belongs to it."""
|
||||
|
||||
access_group: str
|
||||
team_id: str
|
||||
key: str
|
||||
|
||||
|
||||
ModelSelector = Callable[[GroupedDeployments], str]
|
||||
|
||||
ALLOWED: Final[tuple[tuple[str, ModelSelector], ...]] = (
|
||||
("bare name the group's wildcard covers", lambda grouped: WILDCARD_BARE_MODEL),
|
||||
("provider-prefixed name the group's wildcard covers", lambda grouped: WILDCARD_PREFIXED_MODEL),
|
||||
("exactly-named deployment in the group", lambda grouped: grouped.member_model),
|
||||
)
|
||||
|
||||
DENIED: Final[tuple[tuple[str, ModelSelector], ...]] = (
|
||||
("deployment outside the group", lambda grouped: grouped.outsider_model),
|
||||
("provider model outside the group's wildcard", lambda grouped: UNCOVERED_OPENAI_MODEL),
|
||||
("name no provider claims", lambda grouped: f"e2e-ag-unknown-{unique_marker()}"),
|
||||
)
|
||||
|
||||
|
||||
def _provider_key(env_var: str) -> str:
|
||||
return os.environ.get(env_var) or f"os.environ/{env_var}"
|
||||
|
||||
|
||||
def _grouped_model(model_name: str, backend: str, access_groups: list[str] | None) -> ModelNewBody:
|
||||
return ModelNewBody(
|
||||
model_name=model_name,
|
||||
litellm_params=LiteLLMParamsBody(model=backend, api_key=_provider_key("OPENAI_API_KEY")),
|
||||
model_info=ModelInfoBody(access_groups=access_groups),
|
||||
)
|
||||
|
||||
|
||||
def _await_group_members(client: AccessControlClient, access_group: str, expected: frozenset[str]) -> None:
|
||||
"""The grant under test is the group's membership, so prove the proxy recorded it
|
||||
before asserting on what the group lets through."""
|
||||
deadline = time.monotonic() + client.proxy.poll_timeout
|
||||
listed: list[str] = []
|
||||
while time.monotonic() < deadline:
|
||||
info = client.access_group_info(access_group)
|
||||
listed = info.model_names if info is not None else []
|
||||
if expected.issubset(listed):
|
||||
return
|
||||
time.sleep(client.proxy.poll_interval)
|
||||
pytest.fail(
|
||||
f"/access_group/{access_group}/info never listed {sorted(expected)} as members; last read {listed}"
|
||||
)
|
||||
|
||||
|
||||
def _await_team_allowlist(client: AccessControlClient, grant_key: str, access_group: str) -> None:
|
||||
"""Registering a team-scoped deployment appends its public name to the team's
|
||||
allow-list, and a wildcard sitting there directly would grant the model under test
|
||||
on its own. Poll a denial until the message enumerates the allow-list the test
|
||||
means to exercise: the group, and nothing else."""
|
||||
allowlist: Final = f"models=['{access_group}']"
|
||||
deadline = time.monotonic() + client.proxy.poll_timeout
|
||||
body = ""
|
||||
while time.monotonic() < deadline:
|
||||
body = client.chat_status(
|
||||
grant_key, UNCOVERED_OPENAI_MODEL, f"{PROMPT} {unique_marker()}", MAX_COMPLETION_TOKENS
|
||||
).body
|
||||
if allowlist in body:
|
||||
return
|
||||
time.sleep(client.proxy.poll_interval)
|
||||
pytest.fail(f"the team's allow-list never settled to {allowlist}; last denial read {body[:300]}")
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def grouped(client: AccessControlClient) -> Iterator[GroupedDeployments]:
|
||||
marker: Final = unique_marker()
|
||||
deployments: Final = GroupedDeployments(
|
||||
access_group=f"e2e-ag-{marker}",
|
||||
member_model=f"e2e-ag-member-{marker}",
|
||||
outsider_model=f"e2e-ag-outsider-{marker}",
|
||||
)
|
||||
registrations: Final = (
|
||||
_grouped_model(WILDCARD_PATTERN, WILDCARD_PATTERN, [deployments.access_group]),
|
||||
_grouped_model(deployments.member_model, GROUP_BACKEND, [deployments.access_group]),
|
||||
_grouped_model(deployments.outsider_model, GROUP_BACKEND, None),
|
||||
)
|
||||
created: Final = tuple(client.proxy.register_model(body) for body in registrations)
|
||||
try:
|
||||
_await_group_members(
|
||||
client,
|
||||
deployments.access_group,
|
||||
frozenset({WILDCARD_PATTERN, deployments.member_model}),
|
||||
)
|
||||
yield deployments
|
||||
finally:
|
||||
for model_id in created:
|
||||
client.proxy.delete_model(model_id)
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def team_grant(client: AccessControlClient) -> Iterator[TeamGrant]:
|
||||
marker: Final = unique_marker()
|
||||
access_group: Final = f"e2e-agt-{marker}"
|
||||
team_alias: Final = f"e2e-ag-team-{marker}"
|
||||
team_id: Final = client.create_team(team_alias, [access_group])
|
||||
key: Final = client.proxy.generate_key(KeyGenerateBody(models=[], team_id=team_id))
|
||||
model_id: Final = client.proxy.register_model(
|
||||
ModelNewBody(
|
||||
model_name=TEAM_WILDCARD_PATTERN,
|
||||
litellm_params=LiteLLMParamsBody(
|
||||
model=TEAM_WILDCARD_PATTERN, api_key=_provider_key("OPENAI_API_KEY")
|
||||
),
|
||||
model_info=ModelInfoBody(team_id=team_id, access_groups=[access_group]),
|
||||
),
|
||||
listed_for=key,
|
||||
)
|
||||
client.set_team_models(team_id, team_alias, [access_group])
|
||||
try:
|
||||
_await_team_allowlist(client, key, access_group)
|
||||
yield TeamGrant(access_group=access_group, team_id=team_id, key=key)
|
||||
finally:
|
||||
client.proxy.delete_model(model_id)
|
||||
client.proxy.delete_key(key)
|
||||
client.delete_team(team_id)
|
||||
|
||||
|
||||
class TestKeyScopedToAccessGroup:
|
||||
@pytest.mark.covers(
|
||||
"other.auth.model_access_group.wildcard_bare_name_allowed",
|
||||
"other.auth.model_access_group.member_allowed",
|
||||
)
|
||||
@pytest.mark.parametrize(("case", "select_model"), ALLOWED)
|
||||
def test_group_grants_every_deployment_in_it(
|
||||
self,
|
||||
case: str,
|
||||
select_model: ModelSelector,
|
||||
client: AccessControlClient,
|
||||
resources: ResourceManager,
|
||||
grouped: GroupedDeployments,
|
||||
) -> None:
|
||||
key = resources.key(models=[grouped.access_group])
|
||||
model = select_model(grouped)
|
||||
|
||||
result = client.chat_status(
|
||||
key, model, f"{PROMPT} {unique_marker()}", MAX_COMPLETION_TOKENS
|
||||
)
|
||||
|
||||
assert result.status_code == 200, (
|
||||
f"a key holding access group {grouped.access_group!r} must be able to call "
|
||||
f"{model!r} ({case}), got {result.status_code}: {result.body[:300]}"
|
||||
)
|
||||
assert ChatResponse.model_validate_json(result.body).choices, (
|
||||
f"200 must carry a real completion, not an error envelope: {result.body[:300]}"
|
||||
)
|
||||
|
||||
@pytest.mark.covers("other.auth.model_access_group.non_member_denied")
|
||||
@pytest.mark.parametrize(("case", "select_model"), DENIED)
|
||||
def test_group_grants_nothing_outside_it(
|
||||
self,
|
||||
case: str,
|
||||
select_model: ModelSelector,
|
||||
client: AccessControlClient,
|
||||
resources: ResourceManager,
|
||||
grouped: GroupedDeployments,
|
||||
) -> None:
|
||||
key = resources.key(models=[grouped.access_group])
|
||||
model = select_model(grouped)
|
||||
|
||||
result = client.chat_status(
|
||||
key, model, f"{PROMPT} {unique_marker()}", MAX_COMPLETION_TOKENS
|
||||
)
|
||||
|
||||
assert result.status_code == 403, (
|
||||
f"a key holding only access group {grouped.access_group!r} must be denied 403 on "
|
||||
f"{model!r} ({case}), got {result.status_code}: {result.body[:300]}"
|
||||
)
|
||||
assert MODEL_ACCESS_DENIED_MARKER in result.body, (
|
||||
f"403 body must be a key model-access denial, got: {result.body[:300]}"
|
||||
)
|
||||
|
||||
|
||||
class TestTeamScopedToAccessGroup:
|
||||
@pytest.mark.covers("other.auth.model_access_group.team_wildcard_bare_name_allowed")
|
||||
def test_group_grants_the_teams_own_wildcard(
|
||||
self, client: AccessControlClient, team_grant: TeamGrant
|
||||
) -> None:
|
||||
result = client.chat_status(
|
||||
team_grant.key,
|
||||
TEAM_WILDCARD_BARE_MODEL,
|
||||
f"{PROMPT} {unique_marker()}",
|
||||
MAX_COMPLETION_TOKENS,
|
||||
)
|
||||
|
||||
assert result.status_code == 200, (
|
||||
f"a team whose allow-list is access group {team_grant.access_group!r} must be able to "
|
||||
f"call {TEAM_WILDCARD_BARE_MODEL!r} through its team-scoped {TEAM_WILDCARD_PATTERN!r} "
|
||||
f"deployment, got {result.status_code}: {result.body[:300]}"
|
||||
)
|
||||
assert ChatResponse.model_validate_json(result.body).choices, (
|
||||
f"200 must carry a real completion, not an error envelope: {result.body[:300]}"
|
||||
)
|
||||
|
||||
@pytest.mark.covers("other.auth.model_access_group.team_non_member_denied")
|
||||
def test_group_grants_the_team_nothing_outside_it(
|
||||
self, client: AccessControlClient, team_grant: TeamGrant
|
||||
) -> None:
|
||||
model = f"e2e-ag-unknown-{unique_marker()}"
|
||||
|
||||
result = client.chat_status(
|
||||
team_grant.key, model, f"{PROMPT} {unique_marker()}", MAX_COMPLETION_TOKENS
|
||||
)
|
||||
|
||||
assert result.status_code == 403, (
|
||||
f"a team holding only access group {team_grant.access_group!r} must be denied 403 on "
|
||||
f"{model!r}, got {result.status_code}: {result.body[:300]}"
|
||||
)
|
||||
assert TEAM_MODEL_ACCESS_DENIED_MARKER in result.body, (
|
||||
f"403 body must be a team model-access denial, got: {result.body[:300]}"
|
||||
)
|
||||
|
|
@ -12,6 +12,11 @@
|
|||
- {id: other.auth.jwt.valid_token_allows, module: other, tier: P0, area: auth, assertions: [valid_token_allows], source: "handle_jwt.py:77-150", rationale: "Valid JWT with correct issuer + claims grants access"}
|
||||
- {id: other.auth.jwt.expired_denied, module: other, tier: P0, area: auth, assertions: [expired_denied], source: "handle_jwt.py:125-135", rationale: "Expired JWT rejected even with valid signature"}
|
||||
- {id: other.auth.jwt.invalid_signature_denied, module: other, tier: P0, area: auth, assertions: [invalid_signature_denied], source: "handle_jwt.py:145-150", rationale: "Bad/missing signature fails verification"}
|
||||
- {id: other.auth.model_access_group.wildcard_bare_name_allowed, module: other, tier: P0, area: auth, assertions: [wildcard_bare_name_allowed], source: "auth_checks.py:3232 / LIT-5813", fail_before_fix: proven, rationale: "A grant of a group holding a wildcard deployment covers the bare model names callers actually send, not only the provider-prefixed spelling"}
|
||||
- {id: other.auth.model_access_group.member_allowed, module: other, tier: P0, area: auth, assertions: [member_allowed], source: "auth_checks.py:3232", rationale: "A key whose allow-list is a model access group can call the deployments in that group"}
|
||||
- {id: other.auth.model_access_group.non_member_denied, module: other, tier: P0, area: auth, assertions: [non_member_denied], source: "auth_checks.py:3232", rationale: "That same grant reaches nothing outside the group, including provider models the group's wildcard does not cover"}
|
||||
- {id: other.auth.model_access_group.team_wildcard_bare_name_allowed, module: other, tier: P1, area: auth, assertions: [team_wildcard_bare_name_allowed], source: "auth_checks.py:3232 / LIT-5813", fail_before_fix: proven, rationale: "The same bare-name grant holds when the wildcard deployment is team-scoped and the team's allow-list is the group"}
|
||||
- {id: other.auth.model_access_group.team_non_member_denied, module: other, tier: P1, area: auth, assertions: [team_non_member_denied], source: "auth_checks.py:3232", rationale: "A team-level group grant reaches nothing outside the group"}
|
||||
- {id: other.auth.virtual_key.route_permission_enforced, module: other, tier: P0, area: auth, assertions: [route_permission_enforced], source: "route_checks.py:89-151", rationale: "allowed_routes whitelist denies disallowed routes"}
|
||||
- {id: other.auth.virtual_key.route_group_allowed, module: other, tier: P1, area: auth, assertions: [route_group_allowed], source: "route_checks.py:106-128", rationale: "allowed_routes=[llm_api_routes] grants all LLM endpoints"}
|
||||
- {id: other.auth.passthrough.model_allowlist_enforced, module: other, tier: P1, area: auth, assertions: [model_allowlist_enforced], source: "route_checks.py:135-151", rationale: "Passthrough enforces per-key model allow-lists"}
|
||||
|
|
|
|||
|
|
@ -8,10 +8,10 @@ green replay run can never certify against fixtures that have drifted more than
|
|||
a week from the live proxy.
|
||||
|
||||
This module owns the format only. The transports that produce and consume it
|
||||
live in fixture_transport.py; canonical request matching, streaming chunk
|
||||
fidelity, and provider-scoping are follow-ups (LIT-5741/5742/5745) and are
|
||||
deliberately absent here, which is why every interaction file stores the full
|
||||
redacted request even though replay today matches by call order.
|
||||
live in fixture_transport.py and the canonical match keys they compute live in
|
||||
fixture_canonical.py (LIT-5741); streaming chunk fidelity and provider-scoping
|
||||
are follow-ups (LIT-5742/5745). Every interaction file stores the full redacted
|
||||
request because replay matches on its canonicalized content.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
|
@ -54,12 +54,13 @@ class Manifest(BaseModel):
|
|||
|
||||
|
||||
class RecordedRequest(BaseModel):
|
||||
"""The request as the transport saw it, auth header values redacted.
|
||||
"""The request as the transport saw it, auth header values and credential
|
||||
body/form fields redacted.
|
||||
|
||||
Replay today only matches ``method`` (the transport verb, not the HTTP verb)
|
||||
and ``path`` in call order; the rest is stored so LIT-5741 can move to
|
||||
content-based match keys without re-recording. File uploads store a content
|
||||
digest instead of the bytes."""
|
||||
Replay matches on the canonical content key fixture_canonical.py computes
|
||||
over ``method`` (the transport verb, not the HTTP verb), ``path``, and the
|
||||
canonicalized headers, params, body, form, and file identity. File uploads
|
||||
store a content digest instead of the bytes."""
|
||||
|
||||
method: str
|
||||
path: str
|
||||
|
|
|
|||
150
tests/e2e/fixture_canonical.py
Normal file
150
tests/e2e/fixture_canonical.py
Normal file
|
|
@ -0,0 +1,150 @@
|
|||
"""Canonical request identity for replay matching (LIT-5741).
|
||||
|
||||
Matching a replayed call against the raw recorded request never hits: unique
|
||||
markers salt prompts, model names, and tags; every run mints fresh virtual
|
||||
keys; request ids and timestamps differ on every call. Matching on transport
|
||||
verb + path alone collides: two different requests to the same route silently
|
||||
swap responses, which passes when it should miss. The canonicalizer strips
|
||||
exactly the volatile material (volatile headers, credential fields, markers,
|
||||
generated ids, timestamps) and hashes what remains with sorted object keys, so
|
||||
identity is content-based and stable across runs and machines.
|
||||
|
||||
Every rewrite rule lives in this module, next to the transports that apply it:
|
||||
a new volatile header, credential field name, or generated-id shape is one
|
||||
edit here, never a per-suite change.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import hashlib
|
||||
import json
|
||||
import re
|
||||
from dataclasses import dataclass
|
||||
from functools import reduce
|
||||
from typing import Final
|
||||
|
||||
from pydantic import JsonValue
|
||||
|
||||
from fixture_bundle import RecordedRequest
|
||||
|
||||
VOLATILE_HEADER_NAMES: Final[frozenset[str]] = frozenset(
|
||||
{
|
||||
"authorization",
|
||||
"x-litellm-api-key",
|
||||
"x-api-key",
|
||||
"x-goog-api-key",
|
||||
"x-request-id",
|
||||
"traceparent",
|
||||
"tracestate",
|
||||
}
|
||||
)
|
||||
|
||||
SECRET_FIELD_NAMES: Final[frozenset[str]] = frozenset(
|
||||
{"api_key", "aws_access_key_id", "static_headers", "vertex_credentials"}
|
||||
)
|
||||
SECRET_FIELD_SUFFIXES: Final[tuple[str, ...]] = (
|
||||
"_api_key",
|
||||
"_secret_key",
|
||||
"_secret_access_key",
|
||||
"_session_token",
|
||||
"_credentials",
|
||||
"_password",
|
||||
)
|
||||
SECRET_PLACEHOLDER: Final = "<secret>"
|
||||
|
||||
PLACEHOLDER_RULES: Final[tuple[tuple[re.Pattern[str], str], ...]] = (
|
||||
(re.compile(r"(?<![0-9a-fA-F])[0-9a-f]{64}(?![0-9a-fA-F])"), "<sha256>"),
|
||||
(
|
||||
re.compile(r"[0-9a-fA-F]{8}-[0-9a-fA-F]{4}-[0-9a-fA-F]{4}-[0-9a-fA-F]{4}-[0-9a-fA-F]{12}"),
|
||||
"<uuid>",
|
||||
),
|
||||
(re.compile(r"sk-[A-Za-z0-9_-]{16,}"), "<key>"),
|
||||
(
|
||||
re.compile(r"\d{4}-\d{2}-\d{2}[T ]\d{2}:\d{2}:\d{2}(?:\.\d+)?(?:Z|[+-]\d{2}:?\d{2})?"),
|
||||
"<timestamp>",
|
||||
),
|
||||
(re.compile(r"(?<!\d)\d{4}-\d{2}-\d{2}(?!\d)"), "<date>"),
|
||||
(
|
||||
re.compile(r"\b(?:chatcmpl|msgbatch|msg|resp|batch|call|req|ftjob|gen|file)[-_][A-Za-z0-9]{8,}\b"),
|
||||
"<id>",
|
||||
),
|
||||
(re.compile(r"(?<![0-9a-fA-F])[0-9a-f]{12}(?![0-9a-fA-F])"), "<marker>"),
|
||||
)
|
||||
|
||||
|
||||
def is_secret_field(name: str) -> bool:
|
||||
lowered: Final = name.lower()
|
||||
return lowered in SECRET_FIELD_NAMES or lowered.endswith(SECRET_FIELD_SUFFIXES)
|
||||
|
||||
|
||||
def canonical_string(value: str) -> str:
|
||||
return reduce(lambda acc, rule: rule[0].sub(rule[1], acc), PLACEHOLDER_RULES, value)
|
||||
|
||||
|
||||
def _canonical_flat(fields: dict[str, str]) -> dict[str, JsonValue]:
|
||||
return {
|
||||
key: SECRET_PLACEHOLDER if is_secret_field(key) else canonical_string(value)
|
||||
for key, value in fields.items()
|
||||
}
|
||||
|
||||
|
||||
def _canonical_value(value: JsonValue) -> JsonValue:
|
||||
match value:
|
||||
case str():
|
||||
return canonical_string(value)
|
||||
case dict():
|
||||
return {
|
||||
key: SECRET_PLACEHOLDER
|
||||
if is_secret_field(key) and item is not None
|
||||
else _canonical_value(item)
|
||||
for key, item in value.items()
|
||||
}
|
||||
case list():
|
||||
return [_canonical_value(item) for item in value]
|
||||
case _:
|
||||
return value
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class CanonicalRequest:
|
||||
method: str
|
||||
path: str
|
||||
content: str
|
||||
|
||||
@property
|
||||
def key(self) -> str:
|
||||
digest: Final = hashlib.sha256(
|
||||
f"{self.method} {self.path}\n{self.content}".encode()
|
||||
).hexdigest()[:16]
|
||||
return f"{self.method} {self.path} #{digest}"
|
||||
|
||||
def pretty_content(self) -> str:
|
||||
return json.dumps(json.loads(self.content), indent=2, sort_keys=True)
|
||||
|
||||
|
||||
def canonicalize(request: RecordedRequest) -> CanonicalRequest:
|
||||
file_identity: Final[JsonValue | None] = (
|
||||
None
|
||||
if request.file_name is None and request.file_sha256 is None
|
||||
else {
|
||||
"name": None if request.file_name is None else canonical_string(request.file_name),
|
||||
"sha256": request.file_sha256,
|
||||
"bytes": request.file_bytes,
|
||||
}
|
||||
)
|
||||
content: Final[dict[str, JsonValue]] = {
|
||||
"headers": {
|
||||
name.lower(): canonical_string(value)
|
||||
for name, value in request.headers.items()
|
||||
if name.lower() not in VOLATILE_HEADER_NAMES
|
||||
},
|
||||
"params": _canonical_flat(request.params),
|
||||
"body": _canonical_value(request.body),
|
||||
"form": None if request.form is None else _canonical_flat(request.form),
|
||||
"file": file_identity,
|
||||
}
|
||||
return CanonicalRequest(
|
||||
method=request.method,
|
||||
path=canonical_string(request.path),
|
||||
content=json.dumps(content, sort_keys=True, separators=(",", ":")),
|
||||
)
|
||||
|
|
@ -7,23 +7,30 @@ HTTP, no proxy, no provider spend. Because both fulfil ``Transport``, no test
|
|||
or client changes shape; ``build_proxy_client`` picks the transport from
|
||||
``E2E_FIXTURE_MODE`` (live | record | replay, default live).
|
||||
|
||||
Replay matches each call by test node id and call order, verifying transport
|
||||
verb + path and failing hard on any drift (``ReplayMiss``). Canonical
|
||||
content-based match keys are LIT-5741; streaming chunk fidelity is LIT-5742;
|
||||
scoping record/replay to provider-bound traffic is LIT-5745.
|
||||
Replay matches each call by test node id and canonical content key
|
||||
(fixture_canonical.py, LIT-5741): volatile headers, credential fields, unique
|
||||
markers, generated ids, and timestamps are canonicalized out before hashing, so
|
||||
matching is order-independent across distinct keys, FIFO within a key, and a
|
||||
miss fails hard (``ReplayMiss``) printing the computed key and the closest
|
||||
recorded key without ever falling through to a live call. Streaming chunk
|
||||
fidelity is LIT-5742; scoping record/replay to provider-bound traffic is
|
||||
LIT-5745.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import difflib
|
||||
import functools
|
||||
import hashlib
|
||||
import os
|
||||
from collections import deque
|
||||
from dataclasses import dataclass, field
|
||||
from datetime import datetime
|
||||
from itertools import islice
|
||||
from pathlib import Path
|
||||
from typing import Final, Literal, assert_never
|
||||
|
||||
from pydantic import BaseModel
|
||||
from pydantic import BaseModel, JsonValue
|
||||
|
||||
from e2e_http import AuthHeaders, BinaryStream, ProbeResult, Result, StreamingResponse
|
||||
from fixture_bundle import (
|
||||
|
|
@ -43,12 +50,14 @@ from fixture_bundle import (
|
|||
check_freshness,
|
||||
format_age,
|
||||
from_result,
|
||||
interaction_filename,
|
||||
load_bundle,
|
||||
prepare_bundle,
|
||||
slug_for_test,
|
||||
to_json_value,
|
||||
to_result,
|
||||
)
|
||||
from fixture_canonical import CanonicalRequest, canonicalize, is_secret_field
|
||||
from transport import Transport
|
||||
|
||||
type FixtureMode = Literal["live", "record", "replay"]
|
||||
|
|
@ -118,6 +127,27 @@ def _redact(headers: dict[str, str]) -> dict[str, str]:
|
|||
}
|
||||
|
||||
|
||||
def _redact_secret_fields(value: JsonValue) -> JsonValue:
|
||||
match value:
|
||||
case dict():
|
||||
return {
|
||||
key: REDACTED_VALUE
|
||||
if is_secret_field(key) and item is not None
|
||||
else _redact_secret_fields(item)
|
||||
for key, item in value.items()
|
||||
}
|
||||
case list():
|
||||
return [_redact_secret_fields(item) for item in value]
|
||||
case _:
|
||||
return value
|
||||
|
||||
|
||||
def _redact_flat(fields: dict[str, str]) -> dict[str, str]:
|
||||
return {
|
||||
key: REDACTED_VALUE if is_secret_field(key) else value for key, value in fields.items()
|
||||
}
|
||||
|
||||
|
||||
def recorded_request(
|
||||
method: str,
|
||||
path: str,
|
||||
|
|
@ -133,9 +163,9 @@ def recorded_request(
|
|||
method=method,
|
||||
path=path,
|
||||
headers=_redact(_dump_flat(headers)),
|
||||
params=_dump_flat(params),
|
||||
body=None if body is None else to_json_value(body),
|
||||
form=None if form is None else _dump_flat(form),
|
||||
params=_redact_flat(_dump_flat(params)),
|
||||
body=None if body is None else _redact_secret_fields(to_json_value(body)),
|
||||
form=None if form is None else _redact_flat(_dump_flat(form)),
|
||||
file_name=file_name,
|
||||
file_sha256=None if file_content is None else hashlib.sha256(file_content).hexdigest(),
|
||||
file_bytes=None if file_content is None else len(file_content),
|
||||
|
|
@ -304,46 +334,112 @@ class RecordingTransport:
|
|||
return response
|
||||
|
||||
|
||||
def _build_pool(recorded: tuple[Interaction, ...]) -> dict[str, deque[Interaction]]:
|
||||
keys: Final = tuple(canonicalize(interaction.request).key for interaction in recorded)
|
||||
return {
|
||||
key: deque(
|
||||
interaction
|
||||
for candidate_key, interaction in zip(keys, recorded, strict=True)
|
||||
if candidate_key == key
|
||||
)
|
||||
for key in dict.fromkeys(keys)
|
||||
}
|
||||
|
||||
|
||||
def _closest_recorded(
|
||||
canonical: CanonicalRequest, recorded: tuple[Interaction, ...]
|
||||
) -> tuple[CanonicalRequest, str]:
|
||||
candidates: Final = tuple(canonicalize(interaction.request) for interaction in recorded)
|
||||
ratios: Final = tuple(
|
||||
difflib.SequenceMatcher(
|
||||
None, f"{canonical.method} {canonical.path}\n{canonical.content}",
|
||||
f"{candidate.method} {candidate.path}\n{candidate.content}",
|
||||
).ratio()
|
||||
for candidate in candidates
|
||||
)
|
||||
best: Final = max(range(len(candidates)), key=lambda index: ratios[index])
|
||||
return candidates[best], interaction_filename(best, recorded[best].request)
|
||||
|
||||
|
||||
def _miss_message(test_key: str, slug: str, canonical: CanonicalRequest, bundle: LoadedBundle) -> str:
|
||||
recorded: Final = bundle.interactions.get(slug, ())
|
||||
if not recorded:
|
||||
return (
|
||||
f"replay miss for {test_key}: computed key {canonical.key} but nothing is recorded "
|
||||
f"under {slug}; re-record with E2E_FIXTURE_MODE=record"
|
||||
)
|
||||
closest, closest_file = _closest_recorded(canonical, recorded)
|
||||
diff: Final = "\n".join(
|
||||
islice(
|
||||
difflib.unified_diff(
|
||||
closest.pretty_content().splitlines(),
|
||||
canonical.pretty_content().splitlines(),
|
||||
fromfile=f"closest recorded ({closest_file})",
|
||||
tofile="test made",
|
||||
lineterm="",
|
||||
),
|
||||
60,
|
||||
)
|
||||
)
|
||||
return (
|
||||
f"replay miss for {test_key}: no recorded interaction matches key {canonical.key}; "
|
||||
f"closest recorded key is {closest.key} ({closest_file})\n{diff}\n"
|
||||
"re-record with E2E_FIXTURE_MODE=record"
|
||||
)
|
||||
|
||||
|
||||
@dataclass(slots=True)
|
||||
class ReplaySource:
|
||||
"""One shared cursor set over a loaded bundle, so every client built in the
|
||||
session consumes the same recorded sequence per test."""
|
||||
"""One shared pool per test over a loaded bundle, so every client built in
|
||||
the session consumes the same recorded interactions. Every pool is built
|
||||
once at construction and per-key consumption is a single atomic deque pop,
|
||||
so concurrent replay calls never race. Calls match by canonical content
|
||||
key: order-independent across distinct keys (concurrent tests interleave
|
||||
calls nondeterministically), FIFO within one key (a poll loop replays its
|
||||
recorded responses in recorded order)."""
|
||||
|
||||
bundle: LoadedBundle
|
||||
_cursors: dict[str, int] = field(default_factory=dict)
|
||||
_pools: dict[str, dict[str, deque[Interaction]]] = field(init=False)
|
||||
|
||||
def next_interaction(self, method: str, path: str) -> Interaction:
|
||||
test_key = current_test_key()
|
||||
slug = slug_for_test(test_key)
|
||||
recorded = self.bundle.interactions.get(slug, ())
|
||||
index = self._cursors.get(slug, 0)
|
||||
if index >= len(recorded):
|
||||
def __post_init__(self) -> None:
|
||||
self._pools = {
|
||||
slug: _build_pool(recorded) for slug, recorded in self.bundle.interactions.items()
|
||||
}
|
||||
|
||||
def _pool(self, slug: str) -> dict[str, deque[Interaction]]:
|
||||
return self._pools.get(slug, {})
|
||||
|
||||
def next_interaction(self, request: RecordedRequest) -> Interaction:
|
||||
test_key: Final = current_test_key()
|
||||
slug: Final = slug_for_test(test_key)
|
||||
pool: Final = self._pool(slug)
|
||||
canonical: Final = canonicalize(request)
|
||||
queue: Final = pool.get(canonical.key)
|
||||
if queue is None:
|
||||
raise ReplayMiss(_miss_message(test_key, slug, canonical, self.bundle))
|
||||
try:
|
||||
return queue.popleft()
|
||||
except IndexError:
|
||||
raise ReplayMiss(
|
||||
f"replay exhausted for {test_key}: call #{index + 1} ({method} {path}) has no recorded "
|
||||
f"interaction ({len(recorded)} recorded under {slug}); re-record with E2E_FIXTURE_MODE=record"
|
||||
)
|
||||
interaction = recorded[index]
|
||||
if interaction.request.method != method or interaction.request.path != path:
|
||||
raise ReplayMiss(
|
||||
f"replay mismatch for {test_key} at call #{index + 1}: recorded "
|
||||
f"{interaction.request.method} {interaction.request.path}, test made {method} {path}; "
|
||||
"re-record with E2E_FIXTURE_MODE=record"
|
||||
)
|
||||
self._cursors[slug] = index + 1
|
||||
return interaction
|
||||
f"replay exhausted for {test_key}: every recorded interaction for key "
|
||||
f"{canonical.key} is already consumed; re-record with E2E_FIXTURE_MODE=record"
|
||||
) from None
|
||||
|
||||
def leftover_error(self, test_key: str) -> str | None:
|
||||
"""Non-None when the test consumed fewer interactions than were recorded,
|
||||
meaning a passing replay proved less than the bundle claims."""
|
||||
slug = slug_for_test(test_key)
|
||||
recorded = self.bundle.interactions.get(slug, ())
|
||||
consumed = self._cursors.get(slug, 0)
|
||||
if consumed >= len(recorded):
|
||||
slug: Final = slug_for_test(test_key)
|
||||
recorded: Final = self.bundle.interactions.get(slug, ())
|
||||
if not recorded:
|
||||
return None
|
||||
leftover: Final = tuple(
|
||||
interaction for queue in self._pool(slug).values() for interaction in queue
|
||||
)
|
||||
if not leftover:
|
||||
return None
|
||||
pending = recorded[consumed]
|
||||
return (
|
||||
f"replay incomplete for {test_key}: {len(recorded) - consumed} of {len(recorded)} recorded "
|
||||
f"interactions never consumed, next is {pending.request.method} {pending.request.path}; "
|
||||
f"replay incomplete for {test_key}: {len(leftover)} of {len(recorded)} recorded "
|
||||
f"interactions never consumed, e.g. {canonicalize(leftover[0].request).key}; "
|
||||
"re-record with E2E_FIXTURE_MODE=record"
|
||||
)
|
||||
|
||||
|
|
@ -386,7 +482,12 @@ class ReplayTransport:
|
|||
def post[R: BaseModel](
|
||||
self, path: str, *, headers: BaseModel, json: BaseModel, response_type: type[R]
|
||||
) -> Result[R]:
|
||||
return to_result(_expect_result(self.source.next_interaction("post", path)), response_type)
|
||||
return to_result(
|
||||
_expect_result(
|
||||
self.source.next_interaction(recorded_request("post", path, headers=headers, body=json))
|
||||
),
|
||||
response_type,
|
||||
)
|
||||
|
||||
def get[R: BaseModel](
|
||||
self,
|
||||
|
|
@ -397,7 +498,12 @@ class ReplayTransport:
|
|||
response_type: type[R],
|
||||
timeout: float | None = None,
|
||||
) -> Result[R]:
|
||||
return to_result(_expect_result(self.source.next_interaction("get", path)), response_type)
|
||||
return to_result(
|
||||
_expect_result(
|
||||
self.source.next_interaction(recorded_request("get", path, headers=headers, params=params))
|
||||
),
|
||||
response_type,
|
||||
)
|
||||
|
||||
def delete[R: BaseModel](
|
||||
self,
|
||||
|
|
@ -408,25 +514,46 @@ class ReplayTransport:
|
|||
response_type: type[R],
|
||||
params: BaseModel | None = None,
|
||||
) -> Result[R]:
|
||||
return to_result(_expect_result(self.source.next_interaction("delete", path)), response_type)
|
||||
return to_result(
|
||||
_expect_result(
|
||||
self.source.next_interaction(
|
||||
recorded_request("delete", path, headers=headers, body=json, params=params)
|
||||
)
|
||||
),
|
||||
response_type,
|
||||
)
|
||||
|
||||
def patch[R: BaseModel](
|
||||
self, path: str, *, headers: BaseModel, json: BaseModel, response_type: type[R]
|
||||
) -> Result[R]:
|
||||
return to_result(_expect_result(self.source.next_interaction("patch", path)), response_type)
|
||||
return to_result(
|
||||
_expect_result(
|
||||
self.source.next_interaction(recorded_request("patch", path, headers=headers, body=json))
|
||||
),
|
||||
response_type,
|
||||
)
|
||||
|
||||
def put[R: BaseModel](
|
||||
self, path: str, *, headers: BaseModel, json: BaseModel, response_type: type[R]
|
||||
) -> Result[R]:
|
||||
return to_result(_expect_result(self.source.next_interaction("put", path)), response_type)
|
||||
return to_result(
|
||||
_expect_result(
|
||||
self.source.next_interaction(recorded_request("put", path, headers=headers, body=json))
|
||||
),
|
||||
response_type,
|
||||
)
|
||||
|
||||
def stream(self, path: str, *, headers: BaseModel, json: BaseModel) -> StreamingResponse:
|
||||
return _expect_streaming(self.source.next_interaction("stream", path))
|
||||
return _expect_streaming(
|
||||
self.source.next_interaction(recorded_request("stream", path, headers=headers, body=json))
|
||||
)
|
||||
|
||||
def stream_binary(
|
||||
self, path: str, *, headers: BaseModel, json: BaseModel, chunk_size: int = 8192
|
||||
) -> BinaryStream:
|
||||
interaction = self.source.next_interaction("stream_binary", path)
|
||||
interaction = self.source.next_interaction(
|
||||
recorded_request("stream_binary", path, headers=headers, body=json)
|
||||
)
|
||||
match interaction.response:
|
||||
case RecordedBinary(payload=payload):
|
||||
return payload
|
||||
|
|
@ -444,10 +571,16 @@ class ReplayTransport:
|
|||
params: BaseModel | None = None,
|
||||
stream: bool = False,
|
||||
) -> StreamingResponse:
|
||||
return _expect_streaming(self.source.next_interaction("send", path))
|
||||
return _expect_streaming(
|
||||
self.source.next_interaction(
|
||||
recorded_request("send", path, headers=headers, body=json, params=params)
|
||||
)
|
||||
)
|
||||
|
||||
def probe(self, path: str, *, params: BaseModel) -> ProbeResult:
|
||||
interaction = self.source.next_interaction("probe", path)
|
||||
interaction = self.source.next_interaction(
|
||||
recorded_request("probe", path, headers=self.master, params=params)
|
||||
)
|
||||
match interaction.response:
|
||||
case RecordedProbe(payload=payload):
|
||||
return payload
|
||||
|
|
@ -467,10 +600,27 @@ class ReplayTransport:
|
|||
params: BaseModel | None = None,
|
||||
response_type: type[R],
|
||||
) -> Result[R]:
|
||||
return to_result(_expect_result(self.source.next_interaction("upload", path)), response_type)
|
||||
return to_result(
|
||||
_expect_result(
|
||||
self.source.next_interaction(
|
||||
recorded_request(
|
||||
"upload",
|
||||
path,
|
||||
headers=headers,
|
||||
params=params,
|
||||
form=form,
|
||||
file_name=filename,
|
||||
file_content=content,
|
||||
)
|
||||
)
|
||||
),
|
||||
response_type,
|
||||
)
|
||||
|
||||
def download(self, path: str, *, headers: BaseModel) -> StreamingResponse:
|
||||
return _expect_streaming(self.source.next_interaction("download", path))
|
||||
return _expect_streaming(
|
||||
self.source.next_interaction(recorded_request("download", path, headers=headers))
|
||||
)
|
||||
|
||||
|
||||
@functools.lru_cache(maxsize=8)
|
||||
|
|
|
|||
|
|
@ -766,6 +766,8 @@ class ModelInfoBody(BaseModel):
|
|||
# constraint when a prior run's teardown had not removed the row.
|
||||
id: str | None = None
|
||||
mode: ModelMode | None = None
|
||||
access_groups: list[str] | None = None
|
||||
team_id: str | None = None
|
||||
|
||||
|
||||
class ModelNewBody(BaseModel):
|
||||
|
|
@ -861,6 +863,7 @@ class TeamNewResponse(BaseModel):
|
|||
class TeamUpdateBody(BaseModel):
|
||||
team_id: str
|
||||
team_alias: str
|
||||
models: list[str] | None = None
|
||||
|
||||
|
||||
class TeamInfoParams(BaseModel):
|
||||
|
|
|
|||
|
|
@ -278,7 +278,21 @@ class ProxyClient:
|
|||
mode: ModelMode | None = None,
|
||||
) -> str:
|
||||
"""Register a deployment under `model_name` and return its proxy-assigned
|
||||
model_id, once the model is actually servable on the data plane.
|
||||
model_id, once the model is actually servable on the data plane."""
|
||||
return self.register_model(
|
||||
ModelNewBody(
|
||||
model_name=model_name,
|
||||
litellm_params=litellm_params,
|
||||
model_info=ModelInfoBody(mode=mode),
|
||||
)
|
||||
)
|
||||
|
||||
def register_model(self, body: ModelNewBody, listed_for: str | None = None) -> str:
|
||||
"""`create_model` for deployments that carry more than a mode: access groups,
|
||||
team scoping, a pinned id. `listed_for` is the virtual key whose /v1/models
|
||||
view must list the deployment before it counts as servable, because a
|
||||
team-scoped deployment is listed to its own team and to nobody else, master
|
||||
key included; leave it unset for a proxy-wide model.
|
||||
|
||||
/model/new is a control-plane route; the data plane (which serves /chat,
|
||||
/ocr, ...) only picks the new model up on its next DB reload, so a call
|
||||
|
|
@ -296,25 +310,22 @@ class ProxyClient:
|
|||
self.transport.post(
|
||||
"/model/new",
|
||||
headers=self.transport.master,
|
||||
json=ModelNewBody(
|
||||
model_name=model_name,
|
||||
litellm_params=litellm_params,
|
||||
model_info=ModelInfoBody(mode=mode),
|
||||
),
|
||||
json=body,
|
||||
response_type=ModelNewResponse,
|
||||
)
|
||||
).model_id
|
||||
written_at = time.monotonic()
|
||||
self._await_model_servable(model_name)
|
||||
self._await_model_servable(body.model_name, listed_for)
|
||||
settle_propagation(written_at)
|
||||
return model_id
|
||||
|
||||
def _await_model_servable(self, model_name: str) -> None:
|
||||
def _await_model_servable(self, model_name: str, listed_for: str | None = None) -> None:
|
||||
"""Block until the data plane lists `model_name`, or fail at model_servable_timeout."""
|
||||
headers = self.transport.master if listed_for is None else self.transport.bearer(listed_for)
|
||||
outcome = await_servable(
|
||||
lambda poll_timeout: self.transport.get(
|
||||
"/v1/models",
|
||||
headers=self.transport.master,
|
||||
headers=headers,
|
||||
params=NoBody(),
|
||||
response_type=ModelsListResponse,
|
||||
timeout=poll_timeout,
|
||||
|
|
|
|||
163
tests/e2e/test_fixture_canonical.py
Normal file
163
tests/e2e/test_fixture_canonical.py
Normal file
|
|
@ -0,0 +1,163 @@
|
|||
"""Harness coverage for canonical request identity (LIT-5741).
|
||||
|
||||
No proxy and no ``e2e`` marker: pure functions over ``RecordedRequest``. Pins
|
||||
the two failure modes match keys must avoid: keying on volatile material so
|
||||
nothing ever matches (markers, virtual keys, ids, timestamps, volatile
|
||||
headers), and keying on too little so different requests collide and a test
|
||||
silently asserts against another request's response.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import pytest
|
||||
from pydantic import JsonValue
|
||||
|
||||
from fixture_bundle import RecordedRequest
|
||||
from fixture_canonical import CanonicalRequest, canonical_string, canonicalize, is_secret_field
|
||||
|
||||
|
||||
def request(
|
||||
method: str = "post",
|
||||
path: str = "/chat/completions",
|
||||
*,
|
||||
headers: dict[str, str] | None = None,
|
||||
params: dict[str, str] | None = None,
|
||||
body: JsonValue | None = None,
|
||||
form: dict[str, str] | None = None,
|
||||
file_name: str | None = None,
|
||||
file_sha256: str | None = None,
|
||||
file_bytes: int | None = None,
|
||||
) -> RecordedRequest:
|
||||
return RecordedRequest(
|
||||
method=method,
|
||||
path=path,
|
||||
headers=headers or {},
|
||||
params=params or {},
|
||||
body=body,
|
||||
form=form,
|
||||
file_name=file_name,
|
||||
file_sha256=file_sha256,
|
||||
file_bytes=file_bytes,
|
||||
)
|
||||
|
||||
|
||||
class TestPlaceholders:
|
||||
@pytest.mark.parametrize(
|
||||
("raw", "expected"),
|
||||
[
|
||||
("Reply ok. 4d5152a995b7", "Reply ok. <marker>"),
|
||||
("e2e-chat-stream-4d5152a995b7", "e2e-chat-stream-<marker>"),
|
||||
("sk-3mCXCTGmYuEEIU2i2qmVE3Xq6tSK1O0X6ZIRP1Lpw8ZlbNjt", "<key>"),
|
||||
("9f1c8a2e-4b3d-4f6a-8f2f-0a1b2c3d4e5f", "<uuid>"),
|
||||
("z" * 64, "z" * 64),
|
||||
("0123456789abcdef" * 4, "<sha256>"),
|
||||
("2026-08-19T20:57:13.363499+00:00", "<timestamp>"),
|
||||
("2026-08-19", "<date>"),
|
||||
("chatcmpl-C0LO6rRkfJlpJ2mqW9BHYo4Sm8FWl", "<id>"),
|
||||
("batch_688a8b7f9a08819096e0f7c88fcd07c5", "<id>"),
|
||||
("file-XyZ12345abc", "<id>"),
|
||||
("gpt-4o-mini", "gpt-4o-mini"),
|
||||
("max_tokens", "max_tokens"),
|
||||
("sk-1234", "sk-1234"),
|
||||
],
|
||||
)
|
||||
def test_rewrites_exactly_the_volatile_shapes(self, raw: str, expected: str) -> None:
|
||||
assert canonical_string(raw) == expected
|
||||
|
||||
|
||||
class TestSecretFields:
|
||||
@pytest.mark.parametrize(
|
||||
("name", "secret"),
|
||||
[
|
||||
("api_key", True),
|
||||
("openai_api_key", True),
|
||||
("aws_secret_access_key", True),
|
||||
("aws_session_token", True),
|
||||
("vertex_credentials", True),
|
||||
("static_headers", True),
|
||||
("langfuse_secret_key", True),
|
||||
("model", False),
|
||||
("max_completion_tokens", False),
|
||||
("api_base", False),
|
||||
],
|
||||
)
|
||||
def test_names_that_carry_credentials(self, name: str, secret: bool) -> None:
|
||||
assert is_secret_field(name) is secret
|
||||
|
||||
|
||||
class TestKeyStability:
|
||||
def test_volatile_material_does_not_change_the_key(self) -> None:
|
||||
"""Acceptance: a suite recorded on one machine (fresh keys, that day's
|
||||
dates, that run's markers) replays on another with no misses."""
|
||||
first = request(
|
||||
headers={"authorization": "Bearer sk-run-one-aaaaaaaaaaaaaaaa", "x-request-id": "req-1"},
|
||||
params={"start_date": "2026-08-18"},
|
||||
body={
|
||||
"model": "e2e-chat-4d5152a995b7",
|
||||
"messages": [{"role": "user", "content": "Reply ok. 4d5152a995b7"}],
|
||||
"api_key": "sk-live-one-aaaaaaaaaaaaaaaa",
|
||||
},
|
||||
)
|
||||
second = request(
|
||||
headers={"authorization": "Bearer sk-run-two-bbbbbbbbbbbbbbbb", "x-request-id": "req-2"},
|
||||
params={"start_date": "2026-08-19"},
|
||||
body={
|
||||
"model": "e2e-chat-1a2b3c4d5e6f",
|
||||
"messages": [{"role": "user", "content": "Reply ok. 1a2b3c4d5e6f"}],
|
||||
"api_key": "os.environ/OPENAI_API_KEY",
|
||||
},
|
||||
)
|
||||
assert canonicalize(first).key == canonicalize(second).key
|
||||
|
||||
def test_serialization_order_is_not_identity(self) -> None:
|
||||
ordered = request(body={"model": "m", "stream": True})
|
||||
reversed_order = request(body={"stream": True, "model": "m"})
|
||||
assert canonicalize(ordered).key == canonicalize(reversed_order).key
|
||||
|
||||
def test_generated_ids_in_the_path_do_not_change_the_key(self) -> None:
|
||||
first = request("get", "/v1/batches/batch_688a8b7f9a08819096e0f7c88fcd07c5")
|
||||
second = request("get", "/v1/batches/batch_770b9c8f0b19920107f1f8d99fde18d6")
|
||||
assert canonicalize(first).key == canonicalize(second).key
|
||||
|
||||
|
||||
class TestKeyDistinctness:
|
||||
def test_requests_differing_only_inside_canonicalized_fields_stay_distinct(self) -> None:
|
||||
"""Acceptance: a naive verb+path hash collides these; the content key
|
||||
must not, or one test silently asserts against the other's response."""
|
||||
first = request(body={"messages": [{"content": "Reply ok. 4d5152a995b7"}]})
|
||||
second = request(body={"messages": [{"content": "Count to three. 4d5152a995b7"}]})
|
||||
naive = (first.method, first.path)
|
||||
assert naive == (second.method, second.path)
|
||||
assert canonicalize(first).key != canonicalize(second).key
|
||||
|
||||
def test_a_kept_header_is_identity(self) -> None:
|
||||
first = request(headers={"x-litellm-tags": "prod"})
|
||||
second = request(headers={"x-litellm-tags": "shadow"})
|
||||
assert canonicalize(first).key != canonicalize(second).key
|
||||
|
||||
def test_a_volatile_header_is_not_identity(self) -> None:
|
||||
first = request(headers={"traceparent": "00-aa-bb-01", "x-api-key": "one"})
|
||||
second = request(headers={"traceparent": "00-cc-dd-01", "x-api-key": "two"})
|
||||
assert canonicalize(first).key == canonicalize(second).key
|
||||
|
||||
def test_secret_set_versus_unset_stays_distinct(self) -> None:
|
||||
with_key = request(body={"api_key": "sk-live-aaaaaaaaaaaaaaaa"})
|
||||
without_key = request(body={"api_key": None})
|
||||
assert canonicalize(with_key).key != canonicalize(without_key).key
|
||||
|
||||
def test_file_content_is_identity(self) -> None:
|
||||
first = request(
|
||||
"upload", "/v1/files", file_name="batch.jsonl", file_sha256="a" * 64, file_bytes=10
|
||||
)
|
||||
second = request(
|
||||
"upload", "/v1/files", file_name="batch.jsonl", file_sha256="b" * 64, file_bytes=10
|
||||
)
|
||||
assert canonicalize(first).key != canonicalize(second).key
|
||||
|
||||
|
||||
class TestKeyShape:
|
||||
def test_key_names_method_path_and_digest(self) -> None:
|
||||
canonical = canonicalize(request("post", "/model/new", body={"model_name": "m"}))
|
||||
assert isinstance(canonical, CanonicalRequest)
|
||||
assert canonical.key.startswith("post /model/new #")
|
||||
assert len(canonical.key.rsplit("#", 1)[1]) == 16
|
||||
|
|
@ -5,17 +5,22 @@ the live one (dependency injection, no monkeypatching): recording must pass
|
|||
every value through unchanged while writing one redacted interaction file per
|
||||
call, and replay must serve identical values from the bundle alone - the
|
||||
fake's call log proves nothing reaches the inner transport - failing hard
|
||||
(``ReplayMiss``) on any drift in order, verb, or path. The collection-time
|
||||
gate and report header are pinned here too, including the stale message that
|
||||
names the bundle's age.
|
||||
(``ReplayMiss``) on any content drift, printing the computed canonical key and
|
||||
the closest recorded key (LIT-5741; the pure canonicalizer is pinned in
|
||||
test_fixture_canonical.py). The collection-time gate and report header are
|
||||
pinned here too, including the stale message that names the bundle's age.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import hashlib
|
||||
import sys
|
||||
import threading
|
||||
from concurrent.futures import ThreadPoolExecutor
|
||||
from dataclasses import dataclass, field
|
||||
from datetime import datetime, timedelta, timezone
|
||||
from pathlib import Path
|
||||
from uuid import uuid4
|
||||
|
||||
import pytest
|
||||
from pydantic import BaseModel
|
||||
|
|
@ -35,10 +40,12 @@ from fixture_bundle import (
|
|||
Interaction,
|
||||
LoadedBundle,
|
||||
Manifest,
|
||||
RecordedResult,
|
||||
load_bundle,
|
||||
prepare_bundle,
|
||||
slug_for_test,
|
||||
)
|
||||
from fixture_canonical import canonicalize
|
||||
from fixture_transport import (
|
||||
InvalidFixtureMode,
|
||||
RecordingTransport,
|
||||
|
|
@ -50,6 +57,7 @@ from fixture_transport import (
|
|||
fixture_mode_collection_error,
|
||||
fixture_report_lines,
|
||||
parse_fixture_mode,
|
||||
recorded_request,
|
||||
replay_leftover_error,
|
||||
select_transport,
|
||||
)
|
||||
|
|
@ -70,6 +78,17 @@ class Query(BaseModel):
|
|||
q: str
|
||||
|
||||
|
||||
class DeployParams(BaseModel):
|
||||
model: str
|
||||
api_key: str | None = None
|
||||
aws_secret_access_key: str | None = None
|
||||
|
||||
|
||||
class DeployBody(BaseModel):
|
||||
model_name: str
|
||||
litellm_params: DeployParams
|
||||
|
||||
|
||||
STREAMING = StreamingResponse(
|
||||
status_code=200,
|
||||
body="",
|
||||
|
|
@ -272,6 +291,28 @@ class TestRecordingTransport:
|
|||
}
|
||||
assert "sk-secret" not in this_tests_files(root)[0].read_text(encoding="utf-8")
|
||||
|
||||
def test_redacts_credential_body_fields_in_the_recorded_request(self, tmp_path: Path) -> None:
|
||||
fake = FakeTransport()
|
||||
root = tmp_path / "bundle"
|
||||
recording: Transport = RecordingTransport(inner=fake, recorder=make_recorder(root))
|
||||
recording.post(
|
||||
"/model/new",
|
||||
headers=fake.master,
|
||||
json=DeployBody(
|
||||
model_name="m",
|
||||
litellm_params=DeployParams(model="openai/gpt", api_key="sk-live-provider-secret-123456"),
|
||||
),
|
||||
response_type=Payload,
|
||||
)
|
||||
raw = this_tests_files(root)[0].read_text(encoding="utf-8")
|
||||
interaction = Interaction.model_validate_json(raw)
|
||||
assert "sk-live-provider-secret-123456" not in raw
|
||||
assert isinstance(interaction.request.body, dict)
|
||||
params = interaction.request.body["litellm_params"]
|
||||
assert isinstance(params, dict)
|
||||
assert params["api_key"] == "<redacted>"
|
||||
assert params["aws_secret_access_key"] is None
|
||||
|
||||
def test_upload_records_a_content_digest_not_the_bytes(self, tmp_path: Path) -> None:
|
||||
fake = FakeTransport()
|
||||
root = tmp_path / "bundle"
|
||||
|
|
@ -335,25 +376,175 @@ class TestReplayTransport:
|
|||
)
|
||||
assert fake.calls == calls_after_record
|
||||
|
||||
def test_mismatched_call_names_recorded_and_actual(self, tmp_path: Path) -> None:
|
||||
def test_miss_names_the_computed_key_and_the_closest_recorded_key(self, tmp_path: Path) -> None:
|
||||
fake = FakeTransport()
|
||||
root = tmp_path / "bundle"
|
||||
recording: Transport = RecordingTransport(inner=fake, recorder=make_recorder(root))
|
||||
recording.post("/model/new", headers=fake.master, json=Body(prompt="x"), response_type=Payload)
|
||||
replay: Transport = ReplayTransport(source=replay_source(root), master_key="sk-1234")
|
||||
with pytest.raises(ReplayMiss, match=r"recorded post /model/new, test made get /v1/models"):
|
||||
with pytest.raises(ReplayMiss) as excinfo:
|
||||
replay.get("/v1/models", headers=replay.master, params=Query(q="all"), response_type=Payload)
|
||||
message = str(excinfo.value)
|
||||
assert "no recorded interaction matches key get /v1/models #" in message
|
||||
assert "closest recorded key is post /model/new #" in message
|
||||
assert "0000-post-model-new.json" in message
|
||||
assert "re-record with E2E_FIXTURE_MODE=record" in message
|
||||
|
||||
def test_exhausted_recording_names_the_call_count(self, tmp_path: Path) -> None:
|
||||
def test_content_drift_on_the_same_route_misses_with_no_live_call(self, tmp_path: Path) -> None:
|
||||
"""The naive verb+path match replayed a stale response for a request
|
||||
whose content had changed, silently passing; a content key must miss,
|
||||
print both canonical forms' diff, and never reach the inner transport."""
|
||||
fake = FakeTransport()
|
||||
root = tmp_path / "bundle"
|
||||
recording: Transport = RecordingTransport(inner=fake, recorder=make_recorder(root))
|
||||
recording.post("/model/new", headers=fake.master, json=Body(prompt="x"), response_type=Payload)
|
||||
calls_after_record = list(fake.calls)
|
||||
replay: Transport = ReplayTransport(source=replay_source(root), master_key="sk-1234")
|
||||
with pytest.raises(ReplayMiss) as excinfo:
|
||||
replay.post("/model/new", headers=replay.master, json=Body(prompt="y"), response_type=Payload)
|
||||
message = str(excinfo.value)
|
||||
assert "no recorded interaction matches key post /model/new #" in message
|
||||
assert "closest recorded key is post /model/new #" in message
|
||||
assert '- "prompt": "x"' in message
|
||||
assert '+ "prompt": "y"' in message
|
||||
assert fake.calls == calls_after_record
|
||||
|
||||
def test_exhausted_key_names_the_key(self, tmp_path: Path) -> None:
|
||||
fake = FakeTransport()
|
||||
root = tmp_path / "bundle"
|
||||
recording: Transport = RecordingTransport(inner=fake, recorder=make_recorder(root))
|
||||
recording.post("/model/new", headers=fake.master, json=Body(prompt="x"), response_type=Payload)
|
||||
replay: Transport = ReplayTransport(source=replay_source(root), master_key="sk-1234")
|
||||
replay.post("/model/new", headers=replay.master, json=Body(prompt="x"), response_type=Payload)
|
||||
with pytest.raises(ReplayMiss, match=r"call #2 \(post /model/new\) has no recorded interaction \(1 recorded"):
|
||||
with pytest.raises(
|
||||
ReplayMiss, match=r"every recorded interaction for key post /model/new #\w{16} is already consumed"
|
||||
):
|
||||
replay.post("/model/new", headers=replay.master, json=Body(prompt="x"), response_type=Payload)
|
||||
|
||||
def test_replays_out_of_recorded_order_across_distinct_keys(self, tmp_path: Path) -> None:
|
||||
"""Concurrent tests interleave independent calls nondeterministically
|
||||
(e.g. a burst of parallel chat calls), so replay matches by content,
|
||||
never by recorded position."""
|
||||
fake = FakeTransport()
|
||||
root = tmp_path / "bundle"
|
||||
recording: Transport = RecordingTransport(inner=fake, recorder=make_recorder(root))
|
||||
recording.post("/model/new", headers=fake.master, json=Body(prompt="x"), response_type=Payload)
|
||||
recording.post("/key/generate", headers=fake.master, json=Body(prompt="k"), response_type=Payload)
|
||||
source = replay_source(root)
|
||||
replay: Transport = ReplayTransport(source=source, master_key="sk-1234")
|
||||
replay.post("/key/generate", headers=replay.master, json=Body(prompt="k"), response_type=Payload)
|
||||
replay.post("/model/new", headers=replay.master, json=Body(prompt="x"), response_type=Payload)
|
||||
assert source.leftover_error(current_test_key()) is None
|
||||
|
||||
def test_identical_requests_replay_their_responses_in_recorded_order(self, tmp_path: Path) -> None:
|
||||
"""A poll loop makes the same request repeatedly and asserts on the
|
||||
progression, so duplicates under one key stay FIFO."""
|
||||
root = tmp_path / "bundle"
|
||||
recorder = make_recorder(root)
|
||||
recorder.record(
|
||||
test_key=current_test_key(),
|
||||
request=recorded_request(
|
||||
"get", "/v1/models", headers=AuthHeaders(authorization="Bearer sk-x"), params=Query(q="all")
|
||||
),
|
||||
response=RecordedResult(kind="success", status_code=200, data={"value": "first"}),
|
||||
)
|
||||
recorder.record(
|
||||
test_key=current_test_key(),
|
||||
request=recorded_request(
|
||||
"get", "/v1/models", headers=AuthHeaders(authorization="Bearer sk-x"), params=Query(q="all")
|
||||
),
|
||||
response=RecordedResult(kind="success", status_code=200, data={"value": "second"}),
|
||||
)
|
||||
replay: Transport = ReplayTransport(source=replay_source(root), master_key="sk-1234")
|
||||
first = replay.get("/v1/models", headers=replay.master, params=Query(q="all"), response_type=Payload)
|
||||
second = replay.get("/v1/models", headers=replay.master, params=Query(q="all"), response_type=Payload)
|
||||
assert first == Success(status_code=200, data=Payload(value="first"))
|
||||
assert second == Success(status_code=200, data=Payload(value="second"))
|
||||
|
||||
def test_concurrent_replays_of_one_key_serve_each_recording_exactly_once(self, tmp_path: Path) -> None:
|
||||
"""A burst of parallel identical calls consumes one shared pool: no
|
||||
response duplicated, none forgotten, nothing left over at teardown.
|
||||
The tiny switch interval forces thread preemption inside pool setup
|
||||
and consumption, so a non-atomic pool build or pop fails this test."""
|
||||
root = tmp_path / "bundle"
|
||||
recorder = make_recorder(root)
|
||||
for ordinal in range(32):
|
||||
recorder.record(
|
||||
test_key=current_test_key(),
|
||||
request=recorded_request(
|
||||
"get", "/v1/models", headers=AuthHeaders(authorization="Bearer sk-x"), params=Query(q="all")
|
||||
),
|
||||
response=RecordedResult(kind="success", status_code=200, data={"value": f"v{ordinal:02d}"}),
|
||||
)
|
||||
source = replay_source(root)
|
||||
replay: Transport = ReplayTransport(source=source, master_key="sk-1234")
|
||||
barrier = threading.Barrier(8)
|
||||
|
||||
def consume_one() -> str:
|
||||
result = replay.get(
|
||||
"/v1/models", headers=replay.master, params=Query(q="all"), response_type=Payload
|
||||
)
|
||||
assert isinstance(result, Success)
|
||||
return result.data.value
|
||||
|
||||
def consume(_: int) -> tuple[str, ...]:
|
||||
barrier.wait()
|
||||
return tuple(consume_one() for _call in range(4))
|
||||
|
||||
previous_interval = sys.getswitchinterval()
|
||||
sys.setswitchinterval(1e-6)
|
||||
try:
|
||||
with ThreadPoolExecutor(max_workers=8) as executor:
|
||||
served = sorted(value for values in executor.map(consume, range(8)) for value in values)
|
||||
finally:
|
||||
sys.setswitchinterval(previous_interval)
|
||||
assert served == [f"v{ordinal:02d}" for ordinal in range(32)]
|
||||
assert source.leftover_error(current_test_key()) is None
|
||||
|
||||
|
||||
class TestRecordedKeySets:
|
||||
def test_two_separate_recordings_of_one_flow_produce_identical_key_sets(
|
||||
self, tmp_path: Path
|
||||
) -> None:
|
||||
"""Everything a run randomizes (markers, virtual keys, dates) must
|
||||
canonicalize out, so separately recorded runs of the same suite agree
|
||||
on every match key and a bundle recorded elsewhere replays here."""
|
||||
|
||||
def record_flow(root: Path, run_date: str) -> list[str]:
|
||||
fake = FakeTransport()
|
||||
recording: Transport = RecordingTransport(inner=fake, recorder=make_recorder(root))
|
||||
marker = deterministic_marker()
|
||||
recording.post(
|
||||
"/model/new",
|
||||
headers=fake.master,
|
||||
json=DeployBody(
|
||||
model_name=f"e2e-chat-{marker}",
|
||||
litellm_params=DeployParams(model="openai/gpt", api_key=f"sk-live-{uuid4().hex}"),
|
||||
),
|
||||
response_type=Payload,
|
||||
)
|
||||
recording.post(
|
||||
"/chat/completions",
|
||||
headers=recording.bearer(f"sk-{uuid4().hex}"),
|
||||
json=Body(prompt=f"Reply with the single word ok. {marker}"),
|
||||
response_type=Payload,
|
||||
)
|
||||
recording.get(
|
||||
"/spend/logs", headers=fake.master, params=Query(q=run_date), response_type=Payload
|
||||
)
|
||||
loaded = load_bundle(root)
|
||||
assert isinstance(loaded, LoadedBundle)
|
||||
return sorted(
|
||||
canonicalize(interaction.request).key
|
||||
for interactions in loaded.interactions.values()
|
||||
for interaction in interactions
|
||||
)
|
||||
|
||||
first_keys = record_flow(tmp_path / "one", "2026-08-18")
|
||||
second_keys = record_flow(tmp_path / "two", "2026-08-19")
|
||||
assert first_keys == second_keys
|
||||
assert len(first_keys) == 3
|
||||
|
||||
|
||||
class TestReplayLeftover:
|
||||
def test_fully_consumed_recording_leaves_nothing(self, tmp_path: Path) -> None:
|
||||
|
|
@ -378,7 +569,7 @@ class TestReplayLeftover:
|
|||
error = source.leftover_error(current_test_key())
|
||||
assert error is not None
|
||||
assert "1 of 2 recorded interactions never consumed" in error
|
||||
assert "next is probe /health/liveliness" in error
|
||||
assert "e.g. probe /health/liveliness #" in error
|
||||
assert "re-record with E2E_FIXTURE_MODE=record" in error
|
||||
|
||||
def test_test_without_recordings_has_no_leftover(self, tmp_path: Path) -> None:
|
||||
|
|
|
|||
|
|
@ -150,13 +150,13 @@ def test_parse_jsonl_empty_content_is_empty_list():
|
|||
assert bu._get_file_content_as_dictionary(b"") == []
|
||||
|
||||
|
||||
def test_parse_jsonl_malformed_raises():
|
||||
with pytest.raises(Exception):
|
||||
bu._get_file_content_as_dictionary(b"not valid json")
|
||||
def test_parse_jsonl_malformed_lines_skipped():
|
||||
content = b'{"a": 1}\nnot valid json\n{"b": 2}\n'
|
||||
assert bu._get_file_content_as_dictionary(content) == [{"a": 1}, {"b": 2}]
|
||||
|
||||
|
||||
# =========================================================================== #
|
||||
# _iter_batch_input_lines / _iter_batch_input_entries (JSONL parsing)
|
||||
# _iter_batch_input_lines / _iter_batch_output_entries (JSONL parsing)
|
||||
# =========================================================================== #
|
||||
|
||||
|
||||
|
|
@ -173,19 +173,22 @@ def test_iter_input_lines_empty():
|
|||
assert list(bu._iter_batch_input_lines(b"")) == []
|
||||
|
||||
|
||||
def test_iter_input_entries_parses_each_row():
|
||||
def test_iter_output_entries_parses_each_row():
|
||||
content = b'{"body": {"model": "gpt-4o"}}\n{"body": {"model": "claude-3"}}\n'
|
||||
assert list(bu._iter_batch_input_entries(content)) == [
|
||||
assert list(bu._iter_batch_output_entries(content)) == [
|
||||
{"body": {"model": "gpt-4o"}},
|
||||
{"body": {"model": "claude-3"}},
|
||||
]
|
||||
|
||||
|
||||
def test_iter_input_entries_raises_on_malformed_line():
|
||||
# _iter_batch_input_entries raises on a bad row; callers that must survive
|
||||
# bad rows iterate _iter_batch_input_lines and parse per-row instead.
|
||||
with pytest.raises(Exception):
|
||||
list(bu._iter_batch_input_entries(b'{"ok":1}\nnot-json\n'))
|
||||
def test_iter_output_entries_skips_malformed_and_non_object_lines():
|
||||
content = b'{"ok": 1}\nnot-json\n[1, 2]\n{"ok": 2}\n'
|
||||
assert list(bu._iter_batch_output_entries(content)) == [{"ok": 1}, {"ok": 2}]
|
||||
|
||||
|
||||
def test_iter_output_entries_skips_undecodable_line():
|
||||
content = b'{"ok": 1}\n{"note": "\xff-bad"}\n{"ok": 2}\n'
|
||||
assert list(bu._iter_batch_output_entries(content)) == [{"ok": 1}, {"ok": 2}]
|
||||
|
||||
|
||||
# =========================================================================== #
|
||||
|
|
@ -471,6 +474,25 @@ def test_cost_from_content_completion_cost_path(monkeypatch):
|
|||
assert len(calls) == 2 # failed row not costed
|
||||
|
||||
|
||||
def test_empty_body_line_does_not_zero_whole_batch():
|
||||
"""A status-200 row with an empty body makes litellm.completion_cost raise;
|
||||
that line must be skipped instead of zeroing the whole batch."""
|
||||
rows = [
|
||||
_success_row(usage=_usage(10, 5)),
|
||||
{
|
||||
"custom_id": "request-poison-empty",
|
||||
"response": {"status_code": 200, "request_id": "inject-empty-body", "body": {}},
|
||||
},
|
||||
_success_row(usage=_usage(20, 10)),
|
||||
]
|
||||
|
||||
cost, usage, models = bu._aggregate_batch_cost_usage_models(entries=rows, custom_llm_provider="openai")
|
||||
|
||||
assert cost > 0.0
|
||||
assert (usage.prompt_tokens, usage.completion_tokens, usage.total_tokens) == (30, 15, 45)
|
||||
assert models == ["gpt-4o", "gpt-4o"]
|
||||
|
||||
|
||||
def test_cost_from_content_model_info_path(monkeypatch):
|
||||
# model_info set -> batch_cost_calculator(prompt_cost, completion_cost).
|
||||
import litellm.cost_calculator as cc
|
||||
|
|
|
|||
|
|
@ -6392,3 +6392,168 @@ def test_is_user_proxy_admin_rejects_view_only_admin():
|
|||
assert _is_user_proxy_admin(user_obj=viewer) is False
|
||||
assert _is_user_proxy_admin(user_obj=admin) is True
|
||||
assert _is_user_proxy_admin(user_obj=None) is False
|
||||
|
||||
|
||||
def _make_wildcard_access_group_router():
|
||||
"""
|
||||
`openai/*` tagged into an access group, plus an untagged `azure/*`, mirroring a
|
||||
proxy that fronts a whole provider behind one wildcard deployment.
|
||||
"""
|
||||
from litellm import Router
|
||||
|
||||
return Router(
|
||||
model_list=[
|
||||
{
|
||||
"model_name": "openai/*",
|
||||
"litellm_params": {"model": "openai/*", "api_key": "fake"},
|
||||
"model_info": {
|
||||
"id": "wildcard-openai",
|
||||
"access_groups": ["default-models"],
|
||||
},
|
||||
},
|
||||
{
|
||||
"model_name": "azure/*",
|
||||
"litellm_params": {"model": "azure/*", "api_key": "fake"},
|
||||
"model_info": {"id": "wildcard-azure"},
|
||||
},
|
||||
]
|
||||
)
|
||||
|
||||
|
||||
def test_can_object_call_model_access_group_wildcard_accepts_bare_model_name():
|
||||
"""
|
||||
Regression: a key holding only the access group name was denied for `gpt-4o`
|
||||
while `openai/gpt-4o` was allowed, because group membership resolved through the
|
||||
pattern router's raw regex and skipped the `{provider}/{model}` retry that both
|
||||
routing and the direct-wildcard grant already perform.
|
||||
"""
|
||||
from litellm.proxy.auth.auth_checks import _can_object_call_model
|
||||
|
||||
router = _make_wildcard_access_group_router()
|
||||
|
||||
assert (
|
||||
_can_object_call_model(
|
||||
model="gpt-4o",
|
||||
llm_router=router,
|
||||
models=["default-models"],
|
||||
object_type="key",
|
||||
)
|
||||
is True
|
||||
)
|
||||
|
||||
|
||||
def test_can_object_call_model_access_group_wildcard_accepts_prefixed_model_name():
|
||||
from litellm.proxy.auth.auth_checks import _can_object_call_model
|
||||
|
||||
router = _make_wildcard_access_group_router()
|
||||
|
||||
assert (
|
||||
_can_object_call_model(
|
||||
model="openai/gpt-4o",
|
||||
llm_router=router,
|
||||
models=["default-models"],
|
||||
object_type="key",
|
||||
)
|
||||
is True
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"model",
|
||||
[
|
||||
"totally-made-up-model-zzz", # no provider can be inferred
|
||||
"azure/some-deployment", # wildcard exists but carries no access group
|
||||
],
|
||||
)
|
||||
def test_can_object_call_model_access_group_wildcard_does_not_over_grant(model):
|
||||
"""The bare-name retry must not turn an access group into a blanket grant."""
|
||||
from litellm.proxy._types import ProxyException
|
||||
from litellm.proxy.auth.auth_checks import _can_object_call_model
|
||||
|
||||
router = _make_wildcard_access_group_router()
|
||||
|
||||
with pytest.raises(ProxyException):
|
||||
_can_object_call_model(
|
||||
model=model,
|
||||
llm_router=router,
|
||||
models=["default-models"],
|
||||
object_type="key",
|
||||
)
|
||||
|
||||
|
||||
def test_can_object_call_model_access_group_rejects_unconsumed_namespace():
|
||||
"""
|
||||
`bedrockz/...` infers provider `bedrock` from a fragment of the name, so
|
||||
re-prefixing would smuggle an unrecognized namespace through a `bedrock/*` group.
|
||||
"""
|
||||
from litellm import Router
|
||||
from litellm.proxy._types import ProxyException
|
||||
from litellm.proxy.auth.auth_checks import _can_object_call_model
|
||||
|
||||
router = Router(
|
||||
model_list=[
|
||||
{
|
||||
"model_name": "bedrock/*",
|
||||
"litellm_params": {"model": "bedrock/*"},
|
||||
"model_info": {
|
||||
"id": "wildcard-bedrock",
|
||||
"access_groups": ["bedrock-models"],
|
||||
},
|
||||
}
|
||||
]
|
||||
)
|
||||
|
||||
assert (
|
||||
_can_object_call_model(
|
||||
model="anthropic.claude-3-5-sonnet-20240620-v1:0",
|
||||
llm_router=router,
|
||||
models=["bedrock-models"],
|
||||
object_type="key",
|
||||
)
|
||||
is True
|
||||
)
|
||||
|
||||
with pytest.raises(ProxyException):
|
||||
_can_object_call_model(
|
||||
model="bedrockz/anthropic.claude-3-5-sonnet-20240620-v1:0",
|
||||
llm_router=router,
|
||||
models=["bedrock-models"],
|
||||
object_type="key",
|
||||
)
|
||||
|
||||
|
||||
def test_can_object_call_model_team_scoped_wildcard_accepts_bare_model_name():
|
||||
"""
|
||||
Same regression as the proxy-wide wildcard, but for a team-scoped deployment
|
||||
whose public name is a wildcard: those live in a separate per-team pattern
|
||||
index that needed the same `{provider}/{model}` retry.
|
||||
"""
|
||||
from litellm import Router
|
||||
from litellm.proxy.auth.auth_checks import _can_object_call_model
|
||||
|
||||
router = Router(
|
||||
model_list=[
|
||||
{
|
||||
"model_name": "openai/*_team-a_abc",
|
||||
"litellm_params": {"model": "openai/*", "api_key": "fake"},
|
||||
"model_info": {
|
||||
"id": "team-byok-wildcard",
|
||||
"team_id": "team-a",
|
||||
"team_public_model_name": "openai/*",
|
||||
"access_groups": ["team-models"],
|
||||
},
|
||||
}
|
||||
]
|
||||
)
|
||||
|
||||
for model in ("gpt-4o", "openai/gpt-4o"):
|
||||
assert (
|
||||
_can_object_call_model(
|
||||
model=model,
|
||||
llm_router=router,
|
||||
models=["team-models"],
|
||||
object_type="team",
|
||||
team_id="team-a",
|
||||
)
|
||||
is True
|
||||
)
|
||||
|
|
|
|||
|
|
@ -25,6 +25,7 @@ from litellm.proxy.client.cli.commands.auth import (
|
|||
save_token,
|
||||
whoami,
|
||||
)
|
||||
from litellm.proxy.client.cli.commands.claude_settings import SettingsFileOwner
|
||||
|
||||
|
||||
def _mock_cli_sso_start_response(
|
||||
|
|
@ -267,7 +268,7 @@ class TestTokenUtilities:
|
|||
def test_load_token_io_error(self):
|
||||
"""Test loading token with IO error"""
|
||||
with (
|
||||
patch("builtins.open", side_effect=IOError("Permission denied")),
|
||||
patch("builtins.open", side_effect=OSError("Permission denied")),
|
||||
patch("litellm.proxy.client.cli.commands.auth.get_token_file_path") as mock_path,
|
||||
patch("os.path.exists", return_value=True),
|
||||
):
|
||||
|
|
@ -1029,3 +1030,83 @@ class TestSaveTokenPrivateWrite:
|
|||
|
||||
assert json.loads(token_file.read_text()) == {"key": "sk-original", "timestamp": 1234567890}
|
||||
assert list(token_file.parent.glob(".tmp-*")) == []
|
||||
|
||||
|
||||
class TestLoginConfigClaude:
|
||||
"""`lite login --config-claude` wiring into ~/.claude/settings.json"""
|
||||
|
||||
def setup_method(self):
|
||||
self.runner = CliRunner()
|
||||
|
||||
def _run_login(self, tmp_path, args, base_url="https://test.example.com"):
|
||||
settings_path = tmp_path / "claude" / "settings.json"
|
||||
backup_path = tmp_path / "claude_settings_backup.json"
|
||||
poll_response = Mock()
|
||||
poll_response.status_code = 200
|
||||
poll_response.json.return_value = {
|
||||
"status": "ready",
|
||||
"key": "eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.test.jwt",
|
||||
"user_id": "test-user-123",
|
||||
"team_id": "team-1",
|
||||
"teams": ["team-1"],
|
||||
}
|
||||
with (
|
||||
patch("webbrowser.open"),
|
||||
patch("requests.post", return_value=_mock_cli_sso_start_response()),
|
||||
patch("requests.get", return_value=poll_response),
|
||||
patch("litellm.proxy.client.cli.commands.auth.save_token"),
|
||||
patch("litellm.proxy.client.cli.interface.show_commands"),
|
||||
patch("litellm.proxy.client.cli.commands.auth.CLAUDE_SETTINGS_PATH", settings_path),
|
||||
patch(
|
||||
"litellm.proxy.client.cli.commands.auth.SETTINGS_FILE_OWNERS",
|
||||
(SettingsFileOwner(backup_path, "lite up", "lite down"),),
|
||||
),
|
||||
patch(
|
||||
"litellm.proxy.client.cli.commands.claude_settings.shutil.which",
|
||||
return_value="/usr/local/bin/lite",
|
||||
),
|
||||
):
|
||||
result = self.runner.invoke(login, args, obj={"base_url": base_url})
|
||||
return result, settings_path, backup_path
|
||||
|
||||
def test_default_login_does_not_touch_claude_settings(self, tmp_path):
|
||||
result, settings_path, _backup_path = self._run_login(tmp_path, [])
|
||||
|
||||
assert result.exit_code == 0
|
||||
assert "Login successful!" in result.output
|
||||
assert not settings_path.exists()
|
||||
assert "Configured Claude Code" not in result.output
|
||||
|
||||
def test_flag_writes_the_settings_file_and_reports_success(self, tmp_path):
|
||||
result, settings_path, _backup_path = self._run_login(tmp_path, ["--config-claude"])
|
||||
|
||||
assert result.exit_code == 0
|
||||
written = json.loads(settings_path.read_text())
|
||||
assert written["env"]["ANTHROPIC_BASE_URL"] == "https://test.example.com"
|
||||
assert written["apiKeyHelper"] == "/usr/local/bin/lite --base-url https://test.example.com auth print-token"
|
||||
assert "Configured Claude Code" in result.output
|
||||
|
||||
def test_flag_preserves_unrelated_settings_on_an_existing_file(self, tmp_path):
|
||||
settings_path = tmp_path / "claude" / "settings.json"
|
||||
settings_path.parent.mkdir(parents=True)
|
||||
settings_path.write_text(json.dumps({"theme": "dark", "env": {"KEEP": "me"}}))
|
||||
|
||||
result, _settings_path, _backup_path = self._run_login(tmp_path, ["--config-claude"])
|
||||
|
||||
assert result.exit_code == 0
|
||||
written = json.loads(settings_path.read_text())
|
||||
assert written["theme"] == "dark"
|
||||
assert written["env"]["KEEP"] == "me"
|
||||
|
||||
def test_settings_failure_is_reported_without_claiming_login_failed(self, tmp_path):
|
||||
settings_path = tmp_path / "claude" / "settings.json"
|
||||
settings_path.parent.mkdir(parents=True)
|
||||
settings_path.write_text("not json at all {{{")
|
||||
|
||||
result, _settings_path, _backup_path = self._run_login(tmp_path, ["--config-claude"])
|
||||
|
||||
assert result.exit_code != 0
|
||||
assert "Login successful!" in result.output
|
||||
assert "could not configure Claude Code" in result.output
|
||||
assert "invalid JSON" in result.output
|
||||
assert "Authentication failed" not in result.output
|
||||
|
|
|
|||
269
tests/test_litellm/proxy/client/cli/test_claude_settings.py
Normal file
269
tests/test_litellm/proxy/client/cli/test_claude_settings.py
Normal file
|
|
@ -0,0 +1,269 @@
|
|||
import json
|
||||
import shlex
|
||||
import stat
|
||||
import time
|
||||
from unittest.mock import patch
|
||||
|
||||
import pytest
|
||||
from click.testing import CliRunner
|
||||
|
||||
from litellm.proxy.client.cli import cli
|
||||
from litellm.proxy.client.cli.commands.claude_settings import (
|
||||
AUTOROUTE_BACKUP_PATH,
|
||||
BACKUP_PATH,
|
||||
SETTINGS_FILE_OWNERS,
|
||||
ClaudeSettingsError,
|
||||
SettingsFileOwner,
|
||||
resolve_api_key_helper,
|
||||
write_claude_settings,
|
||||
)
|
||||
|
||||
|
||||
def _owners(*backup_paths):
|
||||
"""Stand-in owners for the real `lite up` / `lite autoroute up` registry."""
|
||||
return tuple(SettingsFileOwner(path, "lite up", "lite down") for path in backup_paths)
|
||||
|
||||
CLAUDE_SETTINGS_MODULE = "litellm.proxy.client.cli.commands.claude_settings"
|
||||
AUTH_MODULE = "litellm.proxy.client.cli.commands.auth"
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def paths(tmp_path):
|
||||
return tmp_path / "claude" / "settings.json", tmp_path / "backup.json"
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def lite_on_path():
|
||||
with patch(f"{CLAUDE_SETTINGS_MODULE}.shutil.which", return_value="/usr/local/bin/lite"):
|
||||
yield
|
||||
|
||||
|
||||
class TestWriteClaudeSettings:
|
||||
def test_creates_the_file_and_its_parent_when_missing(self, paths, lite_on_path):
|
||||
settings_path, backup_path = paths
|
||||
assert not settings_path.parent.exists()
|
||||
|
||||
write_claude_settings("https://proxy.example.com/", settings_path, _owners(backup_path))
|
||||
|
||||
written = json.loads(settings_path.read_text())
|
||||
assert written["env"]["ANTHROPIC_BASE_URL"] == "https://proxy.example.com"
|
||||
assert written["apiKeyHelper"] == "/usr/local/bin/lite --base-url https://proxy.example.com auth print-token"
|
||||
|
||||
def test_updates_an_existing_file_preserving_unrelated_settings(self, paths, lite_on_path):
|
||||
settings_path, backup_path = paths
|
||||
settings_path.parent.mkdir(parents=True)
|
||||
settings_path.write_text(
|
||||
json.dumps(
|
||||
{
|
||||
"theme": "dark",
|
||||
"permissions": {"allow": ["Bash"]},
|
||||
"env": {"SOME_OTHER_VAR": "keep-me", "ANTHROPIC_BASE_URL": "https://old.example.com"},
|
||||
"apiKeyHelper": "old-helper",
|
||||
}
|
||||
)
|
||||
)
|
||||
|
||||
write_claude_settings("https://proxy.example.com", settings_path, _owners(backup_path))
|
||||
|
||||
written = json.loads(settings_path.read_text())
|
||||
assert written["theme"] == "dark"
|
||||
assert written["permissions"] == {"allow": ["Bash"]}
|
||||
assert written["env"]["SOME_OTHER_VAR"] == "keep-me"
|
||||
assert written["env"]["ANTHROPIC_BASE_URL"] == "https://proxy.example.com"
|
||||
assert written["apiKeyHelper"] != "old-helper"
|
||||
|
||||
def test_rerunning_against_a_new_proxy_refreshes_both_base_url_and_helper(self, paths, lite_on_path):
|
||||
settings_path, backup_path = paths
|
||||
|
||||
write_claude_settings("https://first.example.com", settings_path, _owners(backup_path))
|
||||
write_claude_settings("https://second.example.com", settings_path, _owners(backup_path))
|
||||
|
||||
written = json.loads(settings_path.read_text())
|
||||
assert written["env"]["ANTHROPIC_BASE_URL"] == "https://second.example.com"
|
||||
assert "second.example.com" in written["apiKeyHelper"]
|
||||
assert "first.example.com" not in written["apiKeyHelper"]
|
||||
|
||||
def test_drops_a_stray_static_api_key_so_the_helper_token_wins(self, paths, lite_on_path):
|
||||
settings_path, backup_path = paths
|
||||
settings_path.parent.mkdir(parents=True)
|
||||
settings_path.write_text(json.dumps({"env": {"ANTHROPIC_API_KEY": "sk-leaked"}}))
|
||||
|
||||
write_claude_settings("https://proxy.example.com", settings_path, _owners(backup_path))
|
||||
|
||||
assert "ANTHROPIC_API_KEY" not in json.loads(settings_path.read_text())["env"]
|
||||
|
||||
def test_written_file_is_owner_only(self, paths, lite_on_path):
|
||||
settings_path, backup_path = paths
|
||||
write_claude_settings("https://proxy.example.com", settings_path, _owners(backup_path))
|
||||
assert stat.S_IMODE(settings_path.stat().st_mode) == 0o600
|
||||
|
||||
def test_refuses_while_lite_up_holds_a_backup(self, paths, lite_on_path):
|
||||
settings_path, backup_path = paths
|
||||
backup_path.write_text("{}")
|
||||
|
||||
with pytest.raises(ClaudeSettingsError, match="lite down"):
|
||||
write_claude_settings("https://proxy.example.com", settings_path, _owners(backup_path))
|
||||
|
||||
assert not settings_path.exists()
|
||||
|
||||
def test_refuses_on_corrupt_existing_settings_without_touching_the_file(self, paths, lite_on_path):
|
||||
settings_path, backup_path = paths
|
||||
settings_path.parent.mkdir(parents=True)
|
||||
settings_path.write_text("not json at all {{{")
|
||||
|
||||
with pytest.raises(ClaudeSettingsError, match="invalid JSON"):
|
||||
write_claude_settings("https://proxy.example.com", settings_path, _owners(backup_path))
|
||||
|
||||
assert settings_path.read_text() == "not json at all {{{"
|
||||
|
||||
def test_reports_an_actionable_error_when_lite_is_not_on_path(self, paths):
|
||||
settings_path, backup_path = paths
|
||||
with patch(f"{CLAUDE_SETTINGS_MODULE}.shutil.which", return_value=None):
|
||||
with pytest.raises(ClaudeSettingsError, match="Could not find `lite`"):
|
||||
write_claude_settings("https://proxy.example.com", settings_path, _owners(backup_path))
|
||||
|
||||
assert not settings_path.exists()
|
||||
|
||||
def test_reports_an_actionable_error_on_a_non_utf8_file(self, paths, lite_on_path):
|
||||
"""Bytes that are not valid UTF-8 must not escape as UnicodeDecodeError.
|
||||
|
||||
UnicodeDecodeError is a ValueError, not an OSError, so a decode-side catch
|
||||
is easy to miss; login's broad `except Exception` would then relabel it as
|
||||
an authentication failure and exit 0.
|
||||
"""
|
||||
settings_path, backup_path = paths
|
||||
settings_path.parent.mkdir(parents=True)
|
||||
settings_path.write_bytes(b'{"theme": "\xff\xfe"}')
|
||||
|
||||
with pytest.raises(ClaudeSettingsError, match="invalid JSON"):
|
||||
write_claude_settings("https://proxy.example.com", settings_path, _owners(backup_path))
|
||||
|
||||
def test_reports_an_actionable_error_when_the_file_cannot_be_read(self, paths, lite_on_path):
|
||||
"""An unreadable settings file must not surface as "Authentication failed".
|
||||
|
||||
login wraps the whole flow in a broad `except Exception`, so any OSError
|
||||
escaping this function gets relabelled as an auth failure and sends the
|
||||
user looking at their SSO config instead of at file permissions.
|
||||
"""
|
||||
settings_path, backup_path = paths
|
||||
settings_path.parent.mkdir(parents=True)
|
||||
settings_path.mkdir()
|
||||
|
||||
with pytest.raises(ClaudeSettingsError, match="Could not read"):
|
||||
write_claude_settings("https://proxy.example.com", settings_path, _owners(backup_path))
|
||||
|
||||
def test_reports_an_actionable_error_when_the_file_cannot_be_written(self, paths, lite_on_path):
|
||||
settings_path, backup_path = paths
|
||||
with patch(
|
||||
f"{CLAUDE_SETTINGS_MODULE}.write_private_json",
|
||||
side_effect=OSError("Read-only file system"),
|
||||
):
|
||||
with pytest.raises(ClaudeSettingsError, match="Read-only file system"):
|
||||
write_claude_settings("https://proxy.example.com", settings_path, _owners(backup_path))
|
||||
|
||||
|
||||
class TestApiKeyHelperIsActuallyInvocable:
|
||||
"""The helper string is executed verbatim by Claude Code, so it has to parse.
|
||||
|
||||
Asserting only on its text is what let a malformed command (`--base-url`, a
|
||||
top-level group option, placed after the `print-token` subcommand) ship: click
|
||||
rejects it with "No such option" and every Claude Code request loses its token.
|
||||
"""
|
||||
|
||||
def _helper_args(self, base_url):
|
||||
with patch(f"{CLAUDE_SETTINGS_MODULE}.shutil.which", return_value="/usr/local/bin/lite"):
|
||||
return shlex.split(resolve_api_key_helper(base_url))[1:]
|
||||
|
||||
def test_the_generated_command_parses(self):
|
||||
result = CliRunner().invoke(cli, self._helper_args("http://localhost:4000"))
|
||||
|
||||
assert "No such option" not in result.output
|
||||
assert result.exit_code != 2
|
||||
|
||||
def test_the_generated_command_reaches_print_token(self):
|
||||
with patch(f"{AUTH_MODULE}.load_token", return_value=None):
|
||||
result = CliRunner().invoke(cli, self._helper_args("http://localhost:4000"))
|
||||
|
||||
assert "Not authenticated" in result.output
|
||||
|
||||
def test_the_generated_command_carries_the_base_url_through(self):
|
||||
stale = {
|
||||
"base_url": "http://other-proxy.example.com",
|
||||
"key": "sk-stale",
|
||||
"timestamp": time.time(),
|
||||
}
|
||||
with patch(f"{AUTH_MODULE}.load_token", return_value=stale):
|
||||
result = CliRunner().invoke(cli, self._helper_args("http://localhost:4000"))
|
||||
|
||||
assert "Not authenticated for this server" in result.output
|
||||
|
||||
|
||||
class TestConflictingOwnersOfTheSettingsFile:
|
||||
"""Both `lite up` and `lite autoroute up` restore a backup when they stop.
|
||||
|
||||
Guarding only one of them leaves the other free to silently revert this
|
||||
write, which is the exact hazard the guard exists to prevent.
|
||||
"""
|
||||
|
||||
def test_any_owner_holding_a_backup_blocks_the_write(self, tmp_path, lite_on_path):
|
||||
settings_path = tmp_path / "claude" / "settings.json"
|
||||
|
||||
for index, owner in enumerate(SETTINGS_FILE_OWNERS):
|
||||
backup = tmp_path / f"backup-{index}.json"
|
||||
backup.write_text("{}")
|
||||
stand_in = SettingsFileOwner(backup, owner.start_command, owner.stop_command)
|
||||
with pytest.raises(ClaudeSettingsError, match="currently managing"):
|
||||
write_claude_settings("https://proxy.example.com", settings_path, (stand_in,))
|
||||
backup.unlink()
|
||||
assert not settings_path.exists()
|
||||
|
||||
def test_the_error_names_the_owner_that_actually_holds_the_file(self, tmp_path, lite_on_path):
|
||||
settings_path = tmp_path / "claude" / "settings.json"
|
||||
backup = tmp_path / "auto.json"
|
||||
backup.write_text("{}")
|
||||
autoroute = SettingsFileOwner(backup, "lite autoroute up", "lite autoroute down")
|
||||
|
||||
with pytest.raises(ClaudeSettingsError, match="`lite autoroute up` is currently managing"):
|
||||
write_claude_settings("https://proxy.example.com", settings_path, (autoroute,))
|
||||
with pytest.raises(ClaudeSettingsError, match="Run `lite autoroute down` first"):
|
||||
write_claude_settings("https://proxy.example.com", settings_path, (autoroute,))
|
||||
|
||||
def test_the_registry_matches_the_paths_the_commands_actually_use(self):
|
||||
"""A second definition of the autoroute dir must not drift from this one."""
|
||||
from litellm.proxy.client.cli.commands.autoroute.process import AUTOROUTE_DIR
|
||||
|
||||
assert AUTOROUTE_BACKUP_PATH == AUTOROUTE_DIR / "claude_settings_backup.json"
|
||||
assert {o.backup_path for o in SETTINGS_FILE_OWNERS} == {BACKUP_PATH, AUTOROUTE_BACKUP_PATH}
|
||||
assert {o.stop_command for o in SETTINGS_FILE_OWNERS} == {"lite down", "lite autoroute down"}
|
||||
|
||||
|
||||
class TestDoesNotDestroyUserOwnedStructure:
|
||||
def test_writes_through_a_symlinked_settings_file(self, tmp_path, lite_on_path):
|
||||
"""os.replace() swaps the symlink for a regular file, detaching a dotfiles repo.
|
||||
|
||||
There is no backup here to undo that, so the link must survive and its
|
||||
target must be the thing that gets updated.
|
||||
"""
|
||||
real = tmp_path / "dotfiles" / "settings.json"
|
||||
real.parent.mkdir()
|
||||
real.write_text(json.dumps({"theme": "dark"}))
|
||||
link = tmp_path / "claude" / "settings.json"
|
||||
link.parent.mkdir()
|
||||
link.symlink_to(real)
|
||||
|
||||
write_claude_settings("https://proxy.example.com", link, ())
|
||||
|
||||
assert link.is_symlink()
|
||||
assert json.loads(real.read_text())["env"]["ANTHROPIC_BASE_URL"] == "https://proxy.example.com"
|
||||
assert json.loads(real.read_text())["theme"] == "dark"
|
||||
|
||||
def test_refuses_rather_than_discarding_a_non_object_env(self, paths, lite_on_path):
|
||||
"""merge coerces a non-dict env to {}; that is silent data loss on a persistent write."""
|
||||
settings_path, backup_path = paths
|
||||
settings_path.parent.mkdir(parents=True)
|
||||
settings_path.write_text(json.dumps({"theme": "dark", "env": "not-an-object"}))
|
||||
|
||||
with pytest.raises(ClaudeSettingsError, match="non-object"):
|
||||
write_claude_settings("https://proxy.example.com", settings_path, _owners(backup_path))
|
||||
|
||||
assert json.loads(settings_path.read_text())["env"] == "not-an-object"
|
||||
|
|
@ -10,6 +10,7 @@ from click.testing import CliRunner
|
|||
|
||||
from litellm.proxy.client.cli.commands import up as up_module
|
||||
from litellm.proxy.client.cli.commands.agents import AgentRunError
|
||||
from litellm.proxy.client.cli.commands.claude_settings import ClaudeSettingsError
|
||||
from litellm.proxy.client.cli.commands.up import (
|
||||
BackupRecord,
|
||||
UpError,
|
||||
|
|
@ -25,6 +26,7 @@ from litellm.proxy.client.cli.commands.up import (
|
|||
)
|
||||
|
||||
UP_MODULE = "litellm.proxy.client.cli.commands.up"
|
||||
AUTH_MODULE = "litellm.proxy.client.cli.commands.auth"
|
||||
|
||||
|
||||
def _patch_paths(monkeypatch, tmp_path):
|
||||
|
|
@ -92,13 +94,13 @@ class TestLoadJsonOrEmpty:
|
|||
def test_raises_clean_error_on_invalid_json(self, tmp_path):
|
||||
path = tmp_path / "settings.json"
|
||||
path.write_text("not json at all {{{")
|
||||
with pytest.raises(UpError, match="invalid JSON"):
|
||||
with pytest.raises(ClaudeSettingsError, match="invalid JSON"):
|
||||
load_json_or_empty(path)
|
||||
|
||||
def test_raises_clean_error_on_non_object_root(self, tmp_path):
|
||||
path = tmp_path / "settings.json"
|
||||
path.write_text(json.dumps([1, 2, 3]))
|
||||
with pytest.raises(UpError, match="invalid JSON"):
|
||||
with pytest.raises(ClaudeSettingsError, match="invalid JSON"):
|
||||
load_json_or_empty(path)
|
||||
|
||||
|
||||
|
|
@ -201,16 +203,16 @@ class TestResolveApiKeyHelper:
|
|||
def test_returns_helper_command_bound_to_the_selected_proxy(self, monkeypatch):
|
||||
monkeypatch.setattr(shutil, "which", lambda name: "/usr/local/bin/lite")
|
||||
helper = resolve_api_key_helper("http://localhost:4000")
|
||||
assert helper == "/usr/local/bin/lite auth print-token --base-url http://localhost:4000"
|
||||
assert helper == "/usr/local/bin/lite --base-url http://localhost:4000 auth print-token"
|
||||
|
||||
def test_quotes_a_base_url_containing_shell_metacharacters(self, monkeypatch):
|
||||
monkeypatch.setattr(shutil, "which", lambda name: "/usr/local/bin/lite")
|
||||
helper = resolve_api_key_helper("http://example.com/path; rm -rf /")
|
||||
assert helper == "/usr/local/bin/lite auth print-token --base-url 'http://example.com/path; rm -rf /'"
|
||||
assert helper == "/usr/local/bin/lite --base-url 'http://example.com/path; rm -rf /' auth print-token"
|
||||
|
||||
def test_raises_when_lite_not_on_path(self, monkeypatch):
|
||||
monkeypatch.setattr(shutil, "which", lambda name: None)
|
||||
with pytest.raises(UpError, match="Could not find `lite`"):
|
||||
with pytest.raises(ClaudeSettingsError, match="Could not find `lite`"):
|
||||
resolve_api_key_helper("http://localhost:4000")
|
||||
|
||||
|
||||
|
|
@ -396,3 +398,50 @@ class TestDownCommand:
|
|||
assert result.exit_code != 0
|
||||
assert result.exception is None or isinstance(result.exception, SystemExit)
|
||||
assert "invalid or unexpected JSON" in result.output
|
||||
|
||||
|
||||
class TestUpCanInvokeTheRealLoginCommand:
|
||||
"""`lite up` calls ctx.invoke(login) on the real command object.
|
||||
|
||||
Every other test in this file monkeypatches `up_module.login` with a fake, so
|
||||
none of them would notice a login parameter that ctx.invoke cannot supply.
|
||||
"""
|
||||
|
||||
def test_ctx_invoke_supplies_every_login_parameter(self):
|
||||
from litellm.proxy.client.cli.commands.auth import login as real_login
|
||||
|
||||
reached = []
|
||||
|
||||
@click.command()
|
||||
@click.pass_context
|
||||
def driver(ctx):
|
||||
ctx.obj = {"base_url": "http://127.0.0.1:9"}
|
||||
ctx.invoke(real_login)
|
||||
|
||||
with patch(
|
||||
f"{AUTH_MODULE}._start_cli_sso_flow",
|
||||
side_effect=lambda base_url: reached.append(base_url) or RuntimeError("stop"),
|
||||
):
|
||||
result = CliRunner().invoke(driver, [], standalone_mode=False)
|
||||
|
||||
assert not isinstance(result.exception, TypeError), result.exception
|
||||
assert reached == ["http://127.0.0.1:9"]
|
||||
|
||||
def test_ctx_invoke_leaves_claude_settings_alone(self, tmp_path):
|
||||
from litellm.proxy.client.cli.commands.auth import login as real_login
|
||||
|
||||
settings_path = tmp_path / "settings.json"
|
||||
|
||||
@click.command()
|
||||
@click.pass_context
|
||||
def driver(ctx):
|
||||
ctx.obj = {"base_url": "http://127.0.0.1:9"}
|
||||
ctx.invoke(real_login)
|
||||
|
||||
with (
|
||||
patch(f"{AUTH_MODULE}.CLAUDE_SETTINGS_PATH", settings_path),
|
||||
patch(f"{AUTH_MODULE}._start_cli_sso_flow", side_effect=RuntimeError("stop")),
|
||||
):
|
||||
CliRunner().invoke(driver, [], standalone_mode=False)
|
||||
|
||||
assert not settings_path.exists()
|
||||
|
|
|
|||
|
|
@ -1739,18 +1739,18 @@ def _make_batch_input_bytes(n_rows: int, padding: int = 200) -> bytes:
|
|||
return ("\n".join(rows)).encode("utf-8")
|
||||
|
||||
|
||||
def test_iter_batch_input_entries_matches_dict_list():
|
||||
def test_iter_batch_output_entries_matches_dict_list():
|
||||
from litellm.batches.batch_utils import (
|
||||
_get_file_content_as_dictionary,
|
||||
_iter_batch_input_entries,
|
||||
_iter_batch_output_entries,
|
||||
)
|
||||
|
||||
raw = _make_batch_input_bytes(50)
|
||||
streamed = list(_iter_batch_input_entries(raw))
|
||||
streamed = list(_iter_batch_output_entries(raw))
|
||||
assert streamed == _get_file_content_as_dictionary(raw)
|
||||
assert streamed[0]["custom_id"] == "request-0"
|
||||
# tolerant of blank lines and a missing trailing newline
|
||||
assert list(_iter_batch_input_entries(raw + b"\n\n")) == streamed
|
||||
assert list(_iter_batch_output_entries(raw + b"\n\n")) == streamed
|
||||
|
||||
|
||||
def test_streaming_count_peak_below_dict_list():
|
||||
|
|
@ -1759,7 +1759,7 @@ def test_streaming_count_peak_below_dict_list():
|
|||
|
||||
from litellm.batches.batch_utils import (
|
||||
_get_file_content_as_dictionary,
|
||||
_iter_batch_input_entries,
|
||||
_iter_batch_output_entries,
|
||||
)
|
||||
|
||||
raw = _make_batch_input_bytes(8000)
|
||||
|
|
@ -1777,7 +1777,7 @@ def test_streaming_count_peak_below_dict_list():
|
|||
def _stream():
|
||||
count = 0
|
||||
models: set = set()
|
||||
for entry in _iter_batch_input_entries(raw):
|
||||
for entry in _iter_batch_output_entries(raw):
|
||||
count += 1
|
||||
model = (entry.get("body") or {}).get("model")
|
||||
if model:
|
||||
|
|
|
|||
|
|
@ -1,6 +1,6 @@
|
|||
{
|
||||
"LIT001": {
|
||||
"limit": 22808
|
||||
"limit": 22805
|
||||
},
|
||||
"LIT002": {
|
||||
"limit": 26878
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue