mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-13 08:52:02 +00:00
Merge pull request #50487 from BerriAI/backport_49984_lts_1108
Some checks failed
LiteLLM Rust / rust-lint (push) Waiting to run
LiteLLM Rust / rust-test (push) Waiting to run
LiteLLM Rust / rust-wheel (push) Waiting to run
Terraform Provider / gofmt, vet, build, test (push) Has been cancelled
Terraform Provider / Provider endpoints vs proxy OpenAPI schema (push) Has been cancelled
Some checks failed
LiteLLM Rust / rust-lint (push) Waiting to run
LiteLLM Rust / rust-test (push) Waiting to run
LiteLLM Rust / rust-wheel (push) Waiting to run
Terraform Provider / gofmt, vet, build, test (push) Has been cancelled
Terraform Provider / Provider endpoints vs proxy OpenAPI schema (push) Has been cancelled
test(e2e): backport the e2e speedup and main's CI cleanups to lts/1.108.x
This commit is contained in:
yuneng-jiang
GitHub
commit
7e95523094
47 changed files with 555 additions and 3842 deletions
69
.github/e2e-stack/assert_tests_ran.py
vendored
69
.github/e2e-stack/assert_tests_ran.py
vendored
|
|
@ -1,69 +0,0 @@
|
|||
import os
|
||||
import re
|
||||
import sys
|
||||
import xml.etree.ElementTree as ET
|
||||
from pathlib import Path
|
||||
from typing import Final
|
||||
|
||||
sys.path.insert(0, str(Path(__file__).resolve().parents[2] / "tests/e2e"))
|
||||
from coverage_registry.management_cases import MANAGEMENT_CASES
|
||||
|
||||
|
||||
def main() -> int:
|
||||
selected: Final = tuple(sys.argv[2:])
|
||||
try:
|
||||
report: Final = ET.parse(Path(sys.argv[1])).getroot()
|
||||
except (ET.ParseError, OSError):
|
||||
_ = sys.stdout.write("::error::could not read the test execution report\n")
|
||||
return 1
|
||||
cases: Final = tuple(report.iter("testcase"))
|
||||
expected_count: Final = os.environ.get("E2E_REQUIRED_TEST_COUNT")
|
||||
passed: Final = frozenset(
|
||||
case.get("file") for case in cases if all(case.find(tag) is None for tag in ("skipped", "failure", "error"))
|
||||
)
|
||||
missing: Final = tuple(path for path in selected if path not in passed)
|
||||
required_nodes: Final = frozenset(case.node for case in MANAGEMENT_CASES if case.node.split("::", 1)[0] in selected)
|
||||
passed_nodes: Final = frozenset(
|
||||
prop.get("value")
|
||||
for case in cases
|
||||
if all(case.find(tag) is None for tag in ("skipped", "failure", "error"))
|
||||
for prop in case.findall("./properties/property")
|
||||
if prop.get("name") == "management_node"
|
||||
)
|
||||
missing_nodes: Final = required_nodes - passed_nodes
|
||||
for node in sorted(missing_nodes):
|
||||
_ = sys.stdout.write(f"::error::required management case did not pass: {node}\n")
|
||||
for path in selected:
|
||||
collected: Final = sum(case.get("file") == path for case in cases)
|
||||
skipped: Final = sum(case.get("file") == path and case.find("skipped") is not None for case in cases)
|
||||
_ = sys.stdout.write(f"{path}: {collected} collected, {skipped} skipped\n")
|
||||
for case in cases:
|
||||
if case.get("file") != path or all(case.find(tag) is None for tag in ("failure", "error", "skipped")):
|
||||
continue
|
||||
outcome = "skipped" if case.find("skipped") is not None else "failed"
|
||||
_ = sys.stdout.write(f" {outcome}: {case.get('classname', '')}::{case.get('name', '')}\n")
|
||||
for prop in case.findall("./properties/property"):
|
||||
name = prop.get("name", "")
|
||||
value = prop.get("value", "")
|
||||
if name in ("oauth_failure_phase", "oauth_exception_type", "oauth_frame") and re.fullmatch(
|
||||
r"[A-Za-z0-9_.:<>-]{1,240}", value
|
||||
):
|
||||
_ = sys.stdout.write(f" {name}: {value}\n")
|
||||
if expected_count is not None and (
|
||||
len(cases) != int(expected_count) or any(case.find("skipped") is not None for case in cases)
|
||||
):
|
||||
_ = sys.stdout.write("::error::required test count was not met or a required case was skipped\n")
|
||||
return 1
|
||||
if (
|
||||
selected
|
||||
and not missing
|
||||
and not missing_nodes
|
||||
and not any(case.find(tag) is not None for case in cases for tag in ("failure", "error"))
|
||||
):
|
||||
return 0
|
||||
_ = sys.stdout.write("::error::every selected file must execute a passing test, with no failures or errors\n")
|
||||
return 1
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(main())
|
||||
|
|
@ -1,319 +0,0 @@
|
|||
import asyncio
|
||||
import aiohttp
|
||||
import json
|
||||
import math
|
||||
from typing import Any
|
||||
|
||||
# Asynchronously fetch data from a given URL
|
||||
async def fetch_data(url):
|
||||
try:
|
||||
# Create an asynchronous session
|
||||
async with aiohttp.ClientSession() as session:
|
||||
# Send a GET request to the URL
|
||||
async with session.get(url) as resp:
|
||||
# Raise an error if the response status is not OK
|
||||
resp.raise_for_status()
|
||||
# Parse the response JSON
|
||||
resp_json = await resp.json()
|
||||
print("Fetch the data from URL.")
|
||||
# Return the 'data' field from the JSON response
|
||||
return resp_json['data']
|
||||
except Exception as e:
|
||||
# Print an error message if fetching data fails
|
||||
print("Error fetching data from URL:", e)
|
||||
return None
|
||||
|
||||
|
||||
FRIENDLI_API_URL = "https://api.friendli.ai/serverless/v1/models"
|
||||
FRIENDLI_PROVIDER = "friendliai"
|
||||
|
||||
INHERITABLE_BASE_KEYS = (
|
||||
"supports_pdf_input",
|
||||
"supports_assistant_prefill",
|
||||
"supports_adaptive_thinking",
|
||||
"supports_output_config",
|
||||
)
|
||||
|
||||
REASONING_EFFORT_LEVEL_ORDER = ("none", "minimal", "low", "medium", "high", "xhigh", "max")
|
||||
|
||||
|
||||
def _find_base_model_entry(base_model: str, local_data: dict) -> str | None:
|
||||
if not base_model:
|
||||
return None
|
||||
bm_tail = base_model.split("/")[-1].lower()
|
||||
if base_model in local_data:
|
||||
return base_model
|
||||
for key in local_data:
|
||||
if key.startswith("sample_spec") or key == "fallback_generalizations":
|
||||
continue
|
||||
if key.split("/")[-1].lower() == bm_tail:
|
||||
return key
|
||||
return None
|
||||
|
||||
|
||||
def _reasoning_effort_levels(reasoning_options: list) -> list:
|
||||
offered = {
|
||||
val
|
||||
for opt in reasoning_options or []
|
||||
if opt.get("type") == "effort"
|
||||
for val in opt.get("values", [])
|
||||
}
|
||||
return [level for level in REASONING_EFFORT_LEVEL_ORDER if level in offered]
|
||||
|
||||
|
||||
def _valid_token_price(value: object) -> bool:
|
||||
try:
|
||||
price = float(value) # pyright: ignore[reportArgumentType] # non-numeric values are rejected via the except
|
||||
except (TypeError, ValueError):
|
||||
return False
|
||||
return math.isfinite(price) and price >= 0
|
||||
|
||||
|
||||
def _has_valid_token_prices(pricing: dict | None) -> bool:
|
||||
prices = pricing or {}
|
||||
return _valid_token_price(prices.get("input")) and _valid_token_price(prices.get("output"))
|
||||
|
||||
|
||||
def _pricing(pricing: dict) -> dict:
|
||||
out: dict[str, Any] = {}
|
||||
if not pricing:
|
||||
return out
|
||||
if "input" in pricing:
|
||||
out["input_cost_per_token"] = float(pricing["input"])
|
||||
if "output" in pricing:
|
||||
out["output_cost_per_token"] = float(pricing["output"])
|
||||
if "input_cache_read" in pricing and pricing["input_cache_read"] is not None:
|
||||
out["cache_read_input_token_cost"] = float(pricing["input_cache_read"])
|
||||
return out
|
||||
|
||||
|
||||
def _modality_flags(input_mods: list) -> dict:
|
||||
mods = input_mods or []
|
||||
has_image = "image" in mods
|
||||
return {
|
||||
"supports_vision": has_image,
|
||||
"supports_image_input": has_image,
|
||||
"supports_video_input": "video" in mods,
|
||||
}
|
||||
|
||||
|
||||
def transform_friendli_data(data: list, local_data: dict) -> dict:
|
||||
transformed: dict[str, dict] = {}
|
||||
if not data:
|
||||
return transformed
|
||||
for model in data:
|
||||
# An unpriced row must never wholesale-replace an already priced local entry:
|
||||
# missing prices cost-calculate as zero, silently zeroing tracked spend
|
||||
if not _has_valid_token_prices(model.get("pricing")):
|
||||
continue
|
||||
model_id = model["id"]
|
||||
base_model = model.get("base_model") or ""
|
||||
entry: dict[str, Any] = {
|
||||
"litellm_provider": FRIENDLI_PROVIDER,
|
||||
}
|
||||
|
||||
base_key = _find_base_model_entry(base_model, local_data)
|
||||
if base_key:
|
||||
base_entry = local_data[base_key]
|
||||
for k in INHERITABLE_BASE_KEYS:
|
||||
if k in base_entry:
|
||||
entry[k] = base_entry[k]
|
||||
|
||||
ctx = model.get("context_length")
|
||||
if ctx is not None:
|
||||
entry["max_input_tokens"] = int(ctx)
|
||||
max_out = model.get("max_completion_tokens")
|
||||
if max_out is not None:
|
||||
entry["max_output_tokens"] = int(max_out)
|
||||
entry["max_tokens"] = int(max_out)
|
||||
|
||||
pricing = _pricing(model.get("pricing", {}))
|
||||
entry.update(pricing)
|
||||
entry["supports_prompt_caching"] = "cache_read_input_token_cost" in pricing
|
||||
|
||||
reasoning = model.get("reasoning") is True
|
||||
entry["supports_reasoning"] = reasoning
|
||||
if reasoning:
|
||||
entry["reasoning_effort_levels"] = _reasoning_effort_levels(
|
||||
model.get("reasoning_options", [])
|
||||
)
|
||||
|
||||
func = model.get("functionality", {})
|
||||
entry["supports_function_calling"] = func.get("tool_call") is True
|
||||
entry["supports_parallel_function_calling"] = func.get("parallel_tool_call") is True
|
||||
is_struct = func.get("structured_output") is True
|
||||
entry["supports_response_schema"] = is_struct
|
||||
entry["supports_native_structured_output"] = is_struct
|
||||
entry["supports_system_messages"] = func.get("system_messages") is True
|
||||
entry["supports_tool_choice"] = func.get("tool_choice") is True
|
||||
|
||||
entry.update(_modality_flags(model.get("input_modalities", [])))
|
||||
|
||||
entry["mode"] = model.get("mode", "chat")
|
||||
|
||||
desc = model.get("description")
|
||||
if desc:
|
||||
entry["comment"] = desc
|
||||
|
||||
dep = model.get("deprecation_date")
|
||||
if dep:
|
||||
entry["deprecation_date"] = dep.split("T")[0]
|
||||
|
||||
entry["source"] = FRIENDLI_API_URL
|
||||
|
||||
transformed[f"{FRIENDLI_PROVIDER}/{model_id}"] = entry
|
||||
return transformed
|
||||
|
||||
# Synchronize local data with remote data
|
||||
def sync_local_data_with_remote(local_data, remote_data, replace_keys=frozenset()):
|
||||
# Update existing keys in local_data with values from remote_data
|
||||
# (replace_keys entries are swapped wholesale so a field the remote catalog
|
||||
# dropped, e.g. cache pricing, cannot survive as a stale value)
|
||||
for key in (set(local_data) & set(remote_data)):
|
||||
if key in replace_keys:
|
||||
local_data[key] = remote_data[key]
|
||||
else:
|
||||
local_data[key].update(remote_data[key])
|
||||
|
||||
# Add new keys from remote_data to local_data
|
||||
for key in (set(remote_data) - set(local_data)):
|
||||
local_data[key] = remote_data[key]
|
||||
|
||||
# Write data to the json file
|
||||
def write_to_file(file_path, data):
|
||||
try:
|
||||
# Open the file in write mode
|
||||
with open(file_path, "w") as file:
|
||||
# Dump the data as JSON into the file
|
||||
json.dump(data, file, indent=4)
|
||||
print("Values updated successfully.")
|
||||
except Exception as e:
|
||||
# Print an error message if writing to file fails
|
||||
print("Error updating JSON file:", e)
|
||||
|
||||
# Update the existing models and add the missing models for OpenRouter
|
||||
def transform_openrouter_data(data):
|
||||
transformed = {}
|
||||
if not data:
|
||||
return transformed
|
||||
for row in data:
|
||||
# Add the fields 'max_tokens' and 'input_cost_per_token'
|
||||
obj = {
|
||||
"max_tokens": row["context_length"],
|
||||
"input_cost_per_token": float(row["pricing"]["prompt"]),
|
||||
}
|
||||
|
||||
# Add 'max_output_tokens' as a field if it is not None
|
||||
if "top_provider" in row and "max_completion_tokens" in row["top_provider"] and row["top_provider"]["max_completion_tokens"] is not None:
|
||||
obj['max_output_tokens'] = int(row["top_provider"]["max_completion_tokens"])
|
||||
|
||||
# Add the field 'output_cost_per_token'
|
||||
obj.update({
|
||||
"output_cost_per_token": float(row["pricing"]["completion"]),
|
||||
})
|
||||
|
||||
# Add field 'input_cost_per_image' if it exists and is non-zero
|
||||
if "pricing" in row and "image" in row["pricing"] and float(row["pricing"]["image"]) != 0.0:
|
||||
obj['input_cost_per_image'] = float(row["pricing"]["image"])
|
||||
|
||||
# Add the fields 'litellm_provider' and 'mode'
|
||||
obj.update({
|
||||
"litellm_provider": "openrouter",
|
||||
"mode": "chat"
|
||||
})
|
||||
|
||||
# Add the 'supports_vision' field if the modality is 'multimodal'
|
||||
if row.get('architecture', {}).get('modality') == 'multimodal':
|
||||
obj['supports_vision'] = True
|
||||
|
||||
# Use a composite key to store the transformed object
|
||||
transformed[f'openrouter/{row["id"]}'] = obj
|
||||
|
||||
return transformed
|
||||
|
||||
# Update the existing models and add the missing models for Vercel AI Gateway
|
||||
def transform_vercel_ai_gateway_data(data):
|
||||
transformed = {}
|
||||
if not data:
|
||||
return transformed
|
||||
for row in data:
|
||||
# Rows without token pricing or token limits (video/embedding models) previously KeyError'd the whole sync
|
||||
if any(row.get(k) is None for k in ("context_window", "max_tokens")) or any(
|
||||
row.get("pricing", {}).get(k) is None for k in ("input", "output")
|
||||
):
|
||||
continue
|
||||
obj = {
|
||||
"max_tokens": row["context_window"],
|
||||
"input_cost_per_token": float(row["pricing"]["input"]),
|
||||
"output_cost_per_token": float(row["pricing"]["output"]),
|
||||
'max_output_tokens': row['max_tokens'],
|
||||
'max_input_tokens': row["context_window"],
|
||||
}
|
||||
|
||||
# Handle cache pricing if available
|
||||
if "pricing" in row:
|
||||
if "input_cache_read" in row["pricing"] and row["pricing"]["input_cache_read"] is not None:
|
||||
obj['cache_read_input_token_cost'] = float(f"{float(row['pricing']['input_cache_read']):e}")
|
||||
|
||||
if "input_cache_write" in row["pricing"] and row["pricing"]["input_cache_write"] is not None:
|
||||
obj['cache_creation_input_token_cost'] = float(f"{float(row['pricing']['input_cache_write']):e}")
|
||||
|
||||
mode = "embedding" if "embedding" in row["id"].lower() else "chat"
|
||||
|
||||
obj.update({"litellm_provider": "vercel_ai_gateway", "mode": mode})
|
||||
|
||||
transformed[f'vercel_ai_gateway/{row["id"]}'] = obj
|
||||
|
||||
return transformed
|
||||
|
||||
|
||||
# Load local data from a specified file
|
||||
def load_local_data(file_path):
|
||||
try:
|
||||
# Open the file in read mode
|
||||
with open(file_path, "r") as file:
|
||||
# Load and return the JSON data
|
||||
return json.load(file)
|
||||
except FileNotFoundError:
|
||||
# Print an error message if the file is not found
|
||||
print("File not found:", file_path)
|
||||
return None
|
||||
except json.JSONDecodeError as e:
|
||||
# Print an error message if JSON decoding fails
|
||||
print("Error decoding JSON:", e)
|
||||
return None
|
||||
|
||||
def main():
|
||||
local_file_path = "model_prices_and_context_window.json" # Path to the local data file
|
||||
openrouter_url = "https://openrouter.ai/api/v1/models" # URL to fetch OpenRouter data
|
||||
vercel_ai_gateway_url = "https://ai-gateway.vercel.sh/v1/models" # URL to fetch Vercel AI Gateway data
|
||||
|
||||
# Load local data from file
|
||||
local_data = load_local_data(local_file_path)
|
||||
|
||||
# Fetch OpenRouter data
|
||||
openrouter_data = asyncio.run(fetch_data(openrouter_url))
|
||||
# Transform the fetched OpenRouter data
|
||||
openrouter_data = transform_openrouter_data(openrouter_data)
|
||||
|
||||
# Fetch Vercel AI Gateway data
|
||||
vercel_data = asyncio.run(fetch_data(vercel_ai_gateway_url))
|
||||
# Transform the fetched Vercel AI Gateway data
|
||||
vercel_data = transform_vercel_ai_gateway_data(vercel_data)
|
||||
|
||||
friendli_data = asyncio.run(fetch_data(FRIENDLI_API_URL))
|
||||
friendli_data = transform_friendli_data(friendli_data, local_data)
|
||||
|
||||
# Combine both datasets
|
||||
all_remote_data = {**openrouter_data, **vercel_data, **friendli_data}
|
||||
|
||||
# If both local and openrouter data are available, synchronize and save
|
||||
if local_data and all_remote_data:
|
||||
sync_local_data_with_remote(local_data, all_remote_data, replace_keys=frozenset(friendli_data))
|
||||
write_to_file(local_file_path, local_data)
|
||||
else:
|
||||
print("Failed to fetch model data from either local file or URL.")
|
||||
|
||||
# Entry point of the script
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
|
|
@ -1,39 +0,0 @@
|
|||
name: Updates model_prices_and_context_window.json and Create Pull Request
|
||||
|
||||
on:
|
||||
schedule:
|
||||
- cron: "0 0 * * 0" # Run every Sundays at midnight
|
||||
#- cron: "0 0 * * *" # Run daily at midnight
|
||||
|
||||
permissions:
|
||||
contents: write
|
||||
pull-requests: write
|
||||
|
||||
jobs:
|
||||
auto_update_price_and_context_window:
|
||||
if: github.repository == 'BerriAI/litellm'
|
||||
runs-on: ubuntu-latest
|
||||
steps:
|
||||
- uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0
|
||||
with:
|
||||
persist-credentials: false
|
||||
- name: Set up uv
|
||||
uses: ./.github/actions/setup-uv-with-retries
|
||||
with:
|
||||
version: "0.10.9"
|
||||
- name: Update JSON Data
|
||||
run: |
|
||||
uv run --frozen --with 'aiohttp==3.13.3' python ".github/scripts/auto_update_price_and_context_window_file.py"
|
||||
- name: Regenerate JSON Schema
|
||||
run: |
|
||||
uv run --frozen python ci_cd/generate_model_prices_schema.py
|
||||
- name: Create Pull Request
|
||||
run: |
|
||||
git add model_prices_and_context_window.json model_prices_and_context_window.schema.json
|
||||
git commit -m "Update model_prices_and_context_window.json file: $(date +'%Y-%m-%d')"
|
||||
gh pr create --title "Update model_prices_and_context_window.json file" \
|
||||
--body "Automated update for model_prices_and_context_window.json" \
|
||||
--head auto-update-price-and-context-window-$(date +'%Y-%m-%d') \
|
||||
--base main
|
||||
env:
|
||||
GH_TOKEN: ${{ secrets.GH_TOKEN }}
|
||||
2
.github/workflows/create-rc-branch.yml
vendored
2
.github/workflows/create-rc-branch.yml
vendored
|
|
@ -2,7 +2,7 @@ name: Create RC Branch
|
|||
|
||||
on:
|
||||
schedule:
|
||||
- cron: "0 3 * * 5"
|
||||
- cron: "0 3 * * 1"
|
||||
timezone: "America/Los_Angeles"
|
||||
workflow_dispatch:
|
||||
|
||||
|
|
|
|||
30
.github/workflows/deploy-dev-on-main.yml
vendored
30
.github/workflows/deploy-dev-on-main.yml
vendored
|
|
@ -1,30 +0,0 @@
|
|||
name: Deploy main to gateway-dev
|
||||
|
||||
on:
|
||||
push:
|
||||
branches: [main]
|
||||
|
||||
permissions: {}
|
||||
|
||||
concurrency:
|
||||
group: deploy-dev-dispatch
|
||||
cancel-in-progress: true
|
||||
|
||||
jobs:
|
||||
dispatch:
|
||||
if: github.repository == 'BerriAI/litellm'
|
||||
runs-on: ubuntu-latest
|
||||
timeout-minutes: 5
|
||||
steps:
|
||||
- uses: actions/create-github-app-token@fee1f7d63c2ff003460e3d139729b119787bc349 # v2.2.2
|
||||
id: token
|
||||
with:
|
||||
app-id: ${{ secrets.LITELLM_OPS_DISPATCH_APP_ID }}
|
||||
private-key: ${{ secrets.LITELLM_OPS_DISPATCH_APP_PRIVATE_KEY }}
|
||||
owner: BerriAI
|
||||
repositories: litellm-ops
|
||||
permission-actions: write
|
||||
- name: Dispatch the litellm-ops dev deploy
|
||||
env:
|
||||
GH_TOKEN: ${{ steps.token.outputs.token }}
|
||||
run: gh api --method POST repos/BerriAI/litellm-ops/actions/workflows/deploy-dev.yml/dispatches -f ref=main
|
||||
70
.github/workflows/sync-together-ai-models.yml
vendored
70
.github/workflows/sync-together-ai-models.yml
vendored
|
|
@ -1,70 +0,0 @@
|
|||
name: Sync Together AI model registry
|
||||
|
||||
on:
|
||||
schedule:
|
||||
- cron: "30 6 * * *"
|
||||
workflow_dispatch:
|
||||
|
||||
permissions:
|
||||
contents: write
|
||||
pull-requests: write
|
||||
|
||||
jobs:
|
||||
sync_together_ai_models:
|
||||
if: github.repository == 'BerriAI/litellm'
|
||||
runs-on: ubuntu-latest
|
||||
env:
|
||||
BASE_BRANCH: ${{ github.event.repository.default_branch }}
|
||||
steps:
|
||||
- uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0
|
||||
with:
|
||||
ref: ${{ env.BASE_BRANCH }}
|
||||
persist-credentials: false
|
||||
- name: Set up uv
|
||||
uses: ./.github/actions/setup-uv-with-retries
|
||||
with:
|
||||
version: "0.10.9"
|
||||
- name: Look for an already-open sync PR
|
||||
id: existing
|
||||
run: |
|
||||
open_pr="$(gh pr list --repo "$GITHUB_REPOSITORY" --state open --limit 1000 --json headRefName \
|
||||
--jq '[.[].headRefName | select(startswith("litellm_together_registry_sync_"))] | first // empty')"
|
||||
echo "open_pr=$open_pr" >> "$GITHUB_OUTPUT"
|
||||
if [ -n "$open_pr" ]; then
|
||||
echo "An open sync PR already exists on branch $open_pr; skipping this run."
|
||||
fi
|
||||
env:
|
||||
GH_TOKEN: ${{ secrets.GH_TOKEN || github.token }}
|
||||
- name: Run the sync
|
||||
if: steps.existing.outputs.open_pr == ''
|
||||
run: |
|
||||
uv run --frozen python scripts/sync_together_ai_models.py --write --pr-body-file "$RUNNER_TEMP/pr_body.md"
|
||||
env:
|
||||
TOGETHER_API_KEY: ${{ secrets.TOGETHER_API_KEY }}
|
||||
- name: Regenerate the JSON schema
|
||||
if: steps.existing.outputs.open_pr == ''
|
||||
run: |
|
||||
uv run --frozen python ci_cd/generate_model_prices_schema.py
|
||||
- name: Create a pull request when the registry changed
|
||||
if: steps.existing.outputs.open_pr == ''
|
||||
run: |
|
||||
if git diff --quiet; then
|
||||
echo "Registry already in sync; no PR needed."
|
||||
exit 0
|
||||
fi
|
||||
branch="litellm_together_registry_sync_$(date +'%Y-%m-%d')"
|
||||
git config user.name "github-actions[bot]"
|
||||
git config user.email "41898282+github-actions[bot]@users.noreply.github.com"
|
||||
git checkout -b "$branch"
|
||||
git add model_prices_and_context_window.json \
|
||||
litellm/model_prices_and_context_window_backup.json \
|
||||
model_prices_and_context_window.schema.json
|
||||
git commit -m "feat(models): sync together_ai model registry $(date +'%Y-%m-%d')"
|
||||
gh auth setup-git
|
||||
git push origin "$branch"
|
||||
gh pr create --title "feat(models): sync together_ai model registry" \
|
||||
--body-file "$RUNNER_TEMP/pr_body.md" \
|
||||
--head "$branch" \
|
||||
--base "$BASE_BRANCH"
|
||||
env:
|
||||
GH_TOKEN: ${{ secrets.GH_TOKEN || github.token }}
|
||||
8
.github/workflows/test-linting.yml
vendored
8
.github/workflows/test-linting.yml
vendored
|
|
@ -376,10 +376,10 @@ jobs:
|
|||
exit 0
|
||||
fi
|
||||
files="$(git diff --name-only "$base" HEAD)"
|
||||
if grep -Eq '^(litellm/(proxy|types)/|ui/litellm-dashboard/(src/lib/http/schema\.d\.ts|scripts/gen-api-types\.mjs|package(-lock)?\.json)$|\.github/workflows/test-linting\.yml$)' <<< "$files"; then
|
||||
if grep -Eq '^(litellm/|enterprise/|(pyproject\.toml|uv\.lock)$|ui/litellm-dashboard/(src/lib/http/schema\.d\.ts|scripts/gen-api-types\.mjs|package(-lock)?\.json)$|\.github/workflows/test-linting\.yml$)' <<< "$files"; then
|
||||
echo "relevant=true" >> "$GITHUB_OUTPUT"
|
||||
else
|
||||
echo "No proxy, types or generator changes in this pull request, nothing to verify."
|
||||
echo "No litellm, enterprise, dependency or generator changes in this pull request, nothing to verify."
|
||||
echo "relevant=false" >> "$GITHUB_OUTPUT"
|
||||
fi
|
||||
|
||||
|
|
@ -539,8 +539,8 @@ jobs:
|
|||
- name: test_unit_passed_gate
|
||||
run: uv run --no-sync pytest -q --noconftest -p no:cacheprovider -c /dev/null tests/code_coverage_tests/test_unit_passed_gate.py
|
||||
|
||||
- name: test_e2e_changed_gate
|
||||
run: uv run --no-sync pytest -q --noconftest -p no:cacheprovider -c /dev/null tests/code_coverage_tests/test_e2e_changed_gate.py tests/code_coverage_tests/test_e2e_idp_stack.py
|
||||
- name: test_e2e_idp_stack
|
||||
run: uv run --no-sync pytest -q --noconftest -p no:cacheprovider -c /dev/null tests/code_coverage_tests/test_e2e_idp_stack.py
|
||||
|
||||
- name: test_e2e_metadata
|
||||
env:
|
||||
|
|
|
|||
180
.github/workflows/test-mcp-oauth-e2e.yml
vendored
180
.github/workflows/test-mcp-oauth-e2e.yml
vendored
|
|
@ -1,180 +0,0 @@
|
|||
name: MCP OAuth happy path
|
||||
|
||||
on:
|
||||
pull_request:
|
||||
paths:
|
||||
- '.github/workflows/test-mcp-oauth-e2e.yml'
|
||||
- '.github/e2e-stack/**'
|
||||
- 'tests/e2e/*.py'
|
||||
- 'tests/e2e/pytest.ini'
|
||||
- 'tests/e2e/idp_realm.json'
|
||||
- 'tests/e2e/mcp/**'
|
||||
- 'litellm/experimental_mcp_client/**'
|
||||
- 'litellm/proxy/_experimental/mcp_server/**'
|
||||
- 'litellm/proxy/auth/**'
|
||||
- 'litellm/proxy/management_endpoints/mcp_management_endpoints.py'
|
||||
- 'litellm/proxy/_types.py'
|
||||
- 'litellm/types/mcp_server/mcp_server_manager.py'
|
||||
- 'litellm/proxy/management_endpoints/*sso*.py'
|
||||
- 'litellm/proxy/management_endpoints/sso/**'
|
||||
- 'litellm/proxy/common_utils/encrypt_decrypt_utils.py'
|
||||
- 'litellm/proxy/proxy_server.py'
|
||||
- 'litellm/proxy/schema.prisma'
|
||||
- 'ui/litellm-dashboard/src/app/connect/**'
|
||||
- 'ui/litellm-dashboard/src/app/mcp/oauth/**'
|
||||
- 'pyproject.toml'
|
||||
- 'uv.lock'
|
||||
workflow_dispatch:
|
||||
|
||||
permissions: {}
|
||||
|
||||
concurrency:
|
||||
group: mcp-oauth-${{ github.ref }}
|
||||
cancel-in-progress: true
|
||||
|
||||
jobs:
|
||||
oauth:
|
||||
if: github.event_name == 'workflow_dispatch' || github.event.pull_request.head.repo.full_name == github.repository
|
||||
runs-on: ubuntu-latest
|
||||
environment: e2e-changed
|
||||
timeout-minutes: 45
|
||||
permissions:
|
||||
contents: read
|
||||
id-token: write
|
||||
services:
|
||||
postgres:
|
||||
image: postgres:16.6
|
||||
env:
|
||||
POSTGRES_USER: litellm
|
||||
POSTGRES_PASSWORD: dbpassword9090
|
||||
POSTGRES_DB: litellm
|
||||
ports:
|
||||
- 5432:5432
|
||||
options: >-
|
||||
--health-cmd "pg_isready -U litellm"
|
||||
--health-interval 5s
|
||||
--health-timeout 5s
|
||||
--health-retries 10
|
||||
env:
|
||||
DATABASE_HOST: 127.0.0.1
|
||||
DATABASE_PORT: '5432'
|
||||
DATABASE_USER: litellm
|
||||
DATABASE_PASSWORD: dbpassword9090
|
||||
DATABASE_NAME: litellm
|
||||
DATABASE_URL: postgresql://litellm:dbpassword9090@127.0.0.1:5432/litellm
|
||||
E2E_KEYCLOAK_URL: http://127.0.0.1:8081
|
||||
E2E_KEYCLOAK_ADMIN_USER: admin
|
||||
E2E_KEYCLOAK_ADMIN_PASSWORD: e2e-ephemeral-idp-not-a-secret
|
||||
E2E_FIXTURE_MODE: live
|
||||
E2E_PROVIDER_CACHE: '0'
|
||||
E2E_MCP_OAUTH_LIVE: '1'
|
||||
E2E_REQUIRED_TEST_COUNT: '4'
|
||||
steps:
|
||||
- name: Checkout the tested source
|
||||
uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0
|
||||
with:
|
||||
ref: ${{ github.event.pull_request.head.sha || github.sha }}
|
||||
persist-credentials: false
|
||||
|
||||
- name: Require and materialize the upstream login
|
||||
env:
|
||||
STORAGE_STATE: ${{ secrets.E2E_LINEAR_STORAGE_STATE_B64 }}
|
||||
run: |
|
||||
umask 077
|
||||
python3 - <<'PY'
|
||||
import base64
|
||||
import json
|
||||
import os
|
||||
import secrets
|
||||
from pathlib import Path
|
||||
encoded = os.environ.get("STORAGE_STATE", "")
|
||||
if not encoded:
|
||||
raise SystemExit("E2E_LINEAR_STORAGE_STATE_B64 is required; capture and provision a test-account login")
|
||||
state = json.loads(base64.b64decode(encoded, validate=True))
|
||||
if not isinstance(state, dict) or not state.get("cookies"):
|
||||
raise SystemExit("The captured login must contain browser cookies")
|
||||
directory = Path(os.environ["RUNNER_TEMP"]) / "mcp-oauth-private"
|
||||
directory.mkdir(mode=0o700)
|
||||
path = directory / "linear-state.json"
|
||||
path.write_text(json.dumps(state))
|
||||
with open(os.environ["GITHUB_ENV"], "a") as output:
|
||||
output.write(f"E2E_LINEAR_STORAGE_STATE={path}\n")
|
||||
for name in ("LITELLM_MASTER_KEY", "LITELLM_SALT_KEY"):
|
||||
value = "sk-e2e-" + secrets.token_hex(24)
|
||||
print(f"::add-mask::{value}")
|
||||
output.write(f"{name}={value}\n")
|
||||
PY
|
||||
|
||||
- uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5.6.0
|
||||
with:
|
||||
python-version: '3.13'
|
||||
- uses: ./.github/actions/setup-uv-with-retries
|
||||
with:
|
||||
version: '0.10.9'
|
||||
- uses: ./.github/actions/cache-cargo-build
|
||||
- name: Install the frozen E2E environment
|
||||
run: |
|
||||
.github/scripts/uv_sync_with_retries.sh --frozen --extra proxy --extra proxy-runtime --extra extra_proxy --group ci --group proxy-dev --group e2e-dev
|
||||
uv run --no-sync python scripts/prisma_generate_if_needed.py
|
||||
uv run --no-sync playwright install --with-deps chromium
|
||||
|
||||
- name: Configure license access
|
||||
id: aws
|
||||
uses: aws-actions/configure-aws-credentials@e7f100cf4c008499ea8adda475de1042d6975c7b # v6.2.0
|
||||
with:
|
||||
role-to-assume: ${{ vars.E2E_AWS_ROLE_TO_ASSUME }}
|
||||
aws-region: us-east-1
|
||||
role-session-name: mcp-oauth-${{ github.run_id }}
|
||||
role-duration-seconds: 900
|
||||
output-env-credentials: false
|
||||
output-credentials: true
|
||||
- name: Load the E2E license
|
||||
env:
|
||||
AWS_ACCESS_KEY_ID: ${{ steps.aws.outputs.aws-access-key-id }}
|
||||
AWS_SECRET_ACCESS_KEY: ${{ steps.aws.outputs.aws-secret-access-key }}
|
||||
AWS_SESSION_TOKEN: ${{ steps.aws.outputs.aws-session-token }}
|
||||
AWS_DEFAULT_REGION: us-east-1
|
||||
run: |
|
||||
license="$(aws secretsmanager get-secret-value --secret-id litellm-e2e-changed-license --query SecretString --output text)"
|
||||
test -n "${license}"
|
||||
echo "::add-mask::${license}"
|
||||
echo "LITELLM_LICENSE=${license}" >> "${GITHUB_ENV}"
|
||||
|
||||
- uses: actions/setup-node@49933ea5288caeca8642d1e84afbd3f7d6820020 # v4.4.0
|
||||
with:
|
||||
node-version-file: ui/litellm-dashboard/.nvmrc
|
||||
- name: Build the gateway consent UI at the tested commit
|
||||
run: |
|
||||
cd ui/litellm-dashboard
|
||||
../../scripts/with_dashboard_node.sh npm ci
|
||||
../../scripts/with_dashboard_node.sh npm run build
|
||||
mkdir -p ../../litellm/proxy/_experimental/out
|
||||
cp -r out/. ../../litellm/proxy/_experimental/out/
|
||||
find ../../litellm/proxy/_experimental/out -name '*.html' ! -name index.html | while read -r page; do
|
||||
mkdir -p "${page%.html}"
|
||||
mv "${page}" "${page%.html}/index.html"
|
||||
done
|
||||
|
||||
- name: Prepare the isolated database and IdP
|
||||
run: |
|
||||
umask 077
|
||||
bash .github/e2e-stack/start-idp.sh
|
||||
uv run --no-sync python migrations/run.py > "${RUNNER_TEMP}/mcp-oauth-private/migrations.log" 2>&1
|
||||
|
||||
- name: Run every required OAuth variant without retries
|
||||
run: |
|
||||
umask 077
|
||||
uv run --no-sync pytest -c tests/e2e/pytest.ini tests/e2e/mcp/test_mcp_oauth_happy_path_e2e.py \
|
||||
--rootdir=. --reruns 0 --tb=short -o junit_family=xunit1 \
|
||||
--junitxml="${RUNNER_TEMP}/mcp-oauth-private/results.xml" \
|
||||
> "${RUNNER_TEMP}/mcp-oauth-private/pytest.log" 2>&1
|
||||
- name: Report JUnit results and reject skipped or missing cases
|
||||
if: always()
|
||||
run: |
|
||||
uv run --no-sync python .github/e2e-stack/assert_tests_ran.py \
|
||||
"${RUNNER_TEMP}/mcp-oauth-private/results.xml" tests/e2e/mcp/test_mcp_oauth_happy_path_e2e.py
|
||||
- name: Remove private login and logs
|
||||
if: always()
|
||||
run: |
|
||||
docker rm -f e2e-keycloak >/dev/null 2>&1 || true
|
||||
rm -rf "${RUNNER_TEMP}/mcp-oauth-private"
|
||||
|
|
@ -12,6 +12,7 @@ on:
|
|||
paths:
|
||||
- "terraform/provider/**"
|
||||
- "litellm/proxy/**"
|
||||
- "enterprise/**"
|
||||
- ".github/workflows/test-terraform-provider.yml"
|
||||
|
||||
permissions:
|
||||
|
|
|
|||
85
.github/workflows/weekly_load_anomaly.yml
vendored
85
.github/workflows/weekly_load_anomaly.yml
vendored
|
|
@ -1,85 +0,0 @@
|
|||
name: "Weekly Load Anomaly Check"
|
||||
|
||||
on:
|
||||
schedule:
|
||||
- cron: "0 12 * * 6"
|
||||
workflow_dispatch:
|
||||
|
||||
permissions:
|
||||
contents: read
|
||||
|
||||
jobs:
|
||||
weekly-load-anomaly:
|
||||
if: github.event_name != 'schedule' || github.repository == 'BerriAI/litellm'
|
||||
runs-on: ubuntu-latest
|
||||
timeout-minutes: 45
|
||||
services:
|
||||
postgres:
|
||||
image: postgres:16.6
|
||||
env:
|
||||
POSTGRES_USER: llmproxy
|
||||
POSTGRES_PASSWORD: dbpassword9090
|
||||
POSTGRES_DB: litellm
|
||||
ports:
|
||||
- 5432:5432
|
||||
options: >-
|
||||
--health-cmd "pg_isready -U llmproxy"
|
||||
--health-interval 5s
|
||||
--health-timeout 5s
|
||||
--health-retries 10
|
||||
env:
|
||||
DATABASE_URL: postgresql://llmproxy:dbpassword9090@localhost:5432/litellm
|
||||
LITELLM_MASTER_KEY: sk-weekly-anomaly-check
|
||||
ANTHROPIC_API_KEY: ${{ secrets.ANTHROPIC_API_KEY }}
|
||||
AWS_BEARER_TOKEN_BEDROCK: ${{ secrets.AWS_BEARER_TOKEN_BEDROCK }}
|
||||
steps:
|
||||
- uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0
|
||||
with:
|
||||
persist-credentials: false
|
||||
|
||||
- name: Set up Python
|
||||
uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5.6.0
|
||||
with:
|
||||
python-version: "3.12"
|
||||
|
||||
- name: Set up uv
|
||||
uses: ./.github/actions/setup-uv-with-retries
|
||||
with:
|
||||
version: "0.10.9"
|
||||
|
||||
- name: Cache the Rust build
|
||||
uses: ./.github/actions/cache-cargo-build
|
||||
|
||||
- name: Install dependencies
|
||||
run: |
|
||||
.github/scripts/uv_sync_with_retries.sh --frozen --group ci --group proxy-dev --extra proxy
|
||||
|
||||
- name: Cache Prisma binaries
|
||||
uses: ./.github/actions/cache-prisma-binaries
|
||||
|
||||
- name: Generate Prisma client
|
||||
run: |
|
||||
uv run --no-sync prisma generate --schema litellm/proxy/schema.prisma
|
||||
|
||||
- name: Start the proxy
|
||||
run: |
|
||||
nohup uv run --no-sync litellm --config tests/e2e/perf/weekly_anomaly_config.yml --port 4000 > proxy.log 2>&1 &
|
||||
for _ in $(seq 1 90); do
|
||||
if curl -fs http://localhost:4000/health/liveliness > /dev/null; then
|
||||
exit 0
|
||||
fi
|
||||
sleep 2
|
||||
done
|
||||
echo "proxy never became live"
|
||||
tail -n 100 proxy.log
|
||||
exit 1
|
||||
|
||||
- name: Run the weekly session anomaly test
|
||||
env:
|
||||
E2E_WEEKLY_ANOMALY: "1"
|
||||
run: |
|
||||
uv run --no-sync pytest tests/e2e/perf/test_weekly_session_anomaly_e2e.py -v --tb=short -rA
|
||||
|
||||
- name: Show proxy log on failure
|
||||
if: failure()
|
||||
run: tail -n 300 proxy.log
|
||||
|
|
@ -18,7 +18,8 @@
|
|||
# -> ruff over ruff-tests.toml + `make lint-test-quality` (test-linting.yml's
|
||||
# test-tree ruff and test-quality gate steps)
|
||||
# - dashboard -> prettier + eslint + lint budgets (test-linting.yml's ui job)
|
||||
# - proxy/types -> regenerate the lazy OpenAPI snapshot and dashboard API types, fail on drift (test-linting.yml's ui-api-types job)
|
||||
# - spec_pattern -> regenerate the lazy OpenAPI snapshot and dashboard API types, fail on drift
|
||||
# (test-linting.yml's ui-api-types job)
|
||||
#
|
||||
# Each block is skipped when no matching files are in scope, so unrelated commits
|
||||
# stay fast. This is intentionally not auto-installed as a git hook (see
|
||||
|
|
@ -98,7 +99,7 @@ existing_files() {
|
|||
litellm_py_pattern='^litellm/.*\.py$'
|
||||
e2e_py_pattern='^tests/e2e(_harness)?/.*\.py$'
|
||||
test_tree_pattern='^(tests/.*\.py|ruff-tests\.toml|scripts/(check_test_quality|test_quality_gate)\.py)$'
|
||||
spec_pattern='^(litellm/(proxy|types)/.*|ui/litellm-dashboard/(scripts/gen-api-types\.mjs|package\.json|package-lock\.json|src/lib/http/schema\.d\.ts))$'
|
||||
spec_pattern='^(litellm/.*|enterprise/.*|pyproject\.toml|uv\.lock|ui/litellm-dashboard/(scripts/gen-api-types\.mjs|package\.json|package-lock\.json|src/lib/http/schema\.d\.ts))$'
|
||||
ui_prettier_pattern='^ui/litellm-dashboard/.*\.(js|jsx|ts|tsx|mjs|cjs|json|css|scss|md|mdx|yml|yaml|html)$'
|
||||
ui_eslint_pattern='^ui/litellm-dashboard/.*\.(js|jsx|ts|tsx|mjs|cjs)$'
|
||||
|
||||
|
|
@ -110,9 +111,6 @@ e2e_py_files=$(scope_match "$e2e_py_pattern")
|
|||
test_tree_files=$(scope_match "$test_tree_pattern")
|
||||
# ruff format (and CI's format step) skip enterprise; the rest of make lint covers it.
|
||||
fmt_files=$(printf '%s\n' "$litellm_py_files" | grep -v '^litellm/enterprise/' | existing_files)
|
||||
# test-linting.yml triggers on any file under litellm/proxy or litellm/types
|
||||
# (Prisma schema and configs included, not just Python) plus the generator and its
|
||||
# lockfiles, so match that whole trigger set rather than a Python subset.
|
||||
spec_files=$(scope_match "$spec_pattern")
|
||||
# CI's ui job runs prettier over a wider extension set than eslint; keep that
|
||||
# split so this flags exactly what the job would.
|
||||
|
|
@ -337,7 +335,7 @@ summary_item "tests/e2e checks (basedpyright + raw HTTP client ban)" "$e2e_py_fi
|
|||
summary_item "test-tree lint (ruff-tests.toml + test-quality gate)" "$test_tree_files" \
|
||||
"no tests/ Python files or test-tree lint inputs in scope"
|
||||
summary_item "dashboard lint (prettier + eslint + lint budgets)" "$ui_prettier_changed$ui_eslint_changed" "no dashboard files in scope"
|
||||
summary_item "dashboard API-type sync (npm run gen:api)" "$spec_files" "no litellm/proxy, litellm/types, or generator files in scope"
|
||||
summary_item "dashboard API-type sync (npm run gen:api)" "$spec_files" "no litellm/, enterprise/, dependency, or generator files in scope"
|
||||
|
||||
if [ -z "$litellm_py_files$e2e_py_files$test_tree_files$ui_prettier_changed$ui_eslint_changed$spec_files" ]; then
|
||||
echo "check: NOTE - no gating lint check matches the files in scope, so nothing ran:" >&2
|
||||
|
|
|
|||
|
|
@ -1,547 +0,0 @@
|
|||
"""Sync the together_ai entries of model_prices_and_context_window.json with Together's live serverless catalog.
|
||||
|
||||
Pulls ``GET https://api.together.ai/v1/models?serverless`` plus the deprecations doc, maps API fields onto
|
||||
registry fields, merges the reviewed capability rules below for everything the API cannot express, and diffs
|
||||
the result against the registry. Dry run (the default) prints the diff summary and the generated PR body;
|
||||
``--write`` applies the changes to the root cost map and its ``litellm/`` backup copy.
|
||||
|
||||
Policy highlights:
|
||||
- Prices arrive per 1M tokens with float artifacts and are normalized to clean per-token values.
|
||||
- A registry entry absent from the serverless catalog is marked with ``deprecation_date`` from the docs
|
||||
deprecation table, never deleted; absences with no docs date are surfaced for a human call.
|
||||
- Availability comes from the API: a model the docs list as removed but the API still serves stays live,
|
||||
with the conflict surfaced as a warning.
|
||||
- Manually curated values the API cannot express (``metadata.successor``, ``max_output_tokens`` on existing
|
||||
entries, capability flags no rule covers) are never overwritten; conflicts are surfaced instead.
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import json
|
||||
import os
|
||||
import re
|
||||
import sys
|
||||
from collections.abc import Mapping, Sequence
|
||||
from dataclasses import dataclass, field
|
||||
from pathlib import Path
|
||||
from types import MappingProxyType
|
||||
from typing import Final
|
||||
|
||||
import httpx
|
||||
from pydantic import BaseModel, TypeAdapter, ValidationError
|
||||
|
||||
MODELS_URL: Final = "https://api.together.ai/v1/models?serverless"
|
||||
DEPRECATIONS_URL: Final = "https://docs.together.ai/docs/deprecations.md"
|
||||
PROVIDER: Final = "together_ai"
|
||||
PREFIX: Final = "together_ai/"
|
||||
SOURCE_URL: Final = "https://docs.together.ai/docs/serverless-models"
|
||||
COST_MAP_RELPATHS: Final = (
|
||||
"model_prices_and_context_window.json",
|
||||
"litellm/model_prices_and_context_window_backup.json",
|
||||
)
|
||||
|
||||
TYPE_TO_MODE: Final = MappingProxyType({"chat": "chat", "embedding": "embedding", "moderation": "chat"})
|
||||
|
||||
|
||||
class SyncError(RuntimeError):
|
||||
pass
|
||||
|
||||
|
||||
class CatalogPricing(BaseModel):
|
||||
input: float
|
||||
output: float
|
||||
cached_input: float | None = None
|
||||
|
||||
|
||||
class CatalogModel(BaseModel):
|
||||
id: str
|
||||
type: str
|
||||
context_length: int | None = None
|
||||
pricing: CatalogPricing
|
||||
|
||||
|
||||
CATALOG_ADAPTER: Final = TypeAdapter(list[CatalogModel])
|
||||
|
||||
RegistryEntry = dict[str, object]
|
||||
CostMap = dict[str, object]
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class CapabilityRule:
|
||||
model_id: str
|
||||
fields: Mapping[str, bool | int]
|
||||
provenance: str
|
||||
|
||||
|
||||
def _rule(model_id: str, provenance: str, **fields: bool | int) -> CapabilityRule:
|
||||
return CapabilityRule(model_id=model_id, fields=MappingProxyType(dict(fields)), provenance=provenance)
|
||||
|
||||
|
||||
_TOOLS: Final = MappingProxyType(
|
||||
{
|
||||
"supports_function_calling": True,
|
||||
"supports_parallel_function_calling": True,
|
||||
"supports_response_schema": True,
|
||||
"supports_tool_choice": True,
|
||||
}
|
||||
)
|
||||
|
||||
CAPABILITY_RULES: Final = (
|
||||
_rule(
|
||||
"MiniMaxAI/MiniMax-M3",
|
||||
"reviewed for the LIT-5968 backfill against https://www.together.ai/models/minimax-m3",
|
||||
**_TOOLS,
|
||||
supports_reasoning=True,
|
||||
supports_vision=True,
|
||||
),
|
||||
_rule("Prism-ML/Ternary-Bonsai-27B", "reviewed for the LIT-5968 backfill; no tool or vision support documented"),
|
||||
_rule(
|
||||
"Qwen/Qwen3.5-9B",
|
||||
"reviewed for the LIT-5968 backfill against https://www.together.ai/models/qwen3-5-9b",
|
||||
**_TOOLS,
|
||||
supports_reasoning=True,
|
||||
supports_vision=True,
|
||||
),
|
||||
_rule(
|
||||
"Qwen/Qwen3.6-Plus",
|
||||
"reviewed for the LIT-5968 backfill; hybrid reasoning model without a documented tools contract",
|
||||
supports_reasoning=True,
|
||||
),
|
||||
_rule("Qwen/Qwen3.7-Max", "reviewed for the LIT-5968 backfill; no tool or vision support documented"),
|
||||
_rule("Qwen/Qwen3.7-Plus", "reviewed for the LIT-5968 backfill; no tool or vision support documented"),
|
||||
_rule("Qwen/Qwen3.8-2.4T-A95B", "reviewed for the LIT-5968 backfill; no tool or vision support documented"),
|
||||
_rule("arize-ai/qwen-2-1.5b-instruct", "reviewed for the LIT-5968 backfill; no tool or vision support documented"),
|
||||
_rule(
|
||||
"deepseek-ai/DeepSeek-V4-Flash-0731",
|
||||
"reviewed for the LIT-5968 backfill against https://www.together.ai/models/deepseek-v4-flash",
|
||||
**_TOOLS,
|
||||
),
|
||||
_rule(
|
||||
"deepseek-ai/DeepSeek-V4-Pro",
|
||||
"reviewed for the LIT-5968 backfill against https://www.together.ai/models/deepseek-v4-pro",
|
||||
**_TOOLS,
|
||||
supports_reasoning=True,
|
||||
),
|
||||
_rule(
|
||||
"deepseek-ai/DeepSeek-V4-Pro-0813",
|
||||
"reviewed for the LIT-5968 backfill against https://www.together.ai/models/deepseek-v4-pro",
|
||||
**_TOOLS,
|
||||
),
|
||||
_rule("google/gemma-3n-E4B-it", "reviewed for the LIT-5968 backfill; no tool or vision support documented"),
|
||||
_rule(
|
||||
"google/gemma-4-31B-it",
|
||||
"reviewed for the LIT-5968 backfill against https://www.together.ai/models/gemma-4-31b-it",
|
||||
**_TOOLS,
|
||||
supports_vision=True,
|
||||
),
|
||||
_rule(
|
||||
"intfloat/multilingual-e5-large-instruct",
|
||||
"embedding dims per https://huggingface.co/intfloat/multilingual-e5-large-instruct",
|
||||
output_vector_size=1024,
|
||||
),
|
||||
_rule(
|
||||
"meta-llama/Llama-3.3-70B-Instruct-Turbo",
|
||||
"reviewed for the LIT-5968 backfill against https://docs.together.ai/docs/function-calling",
|
||||
**_TOOLS,
|
||||
),
|
||||
_rule(
|
||||
"meta-llama/Llama-Guard-4-12B",
|
||||
"moderation classifier with a chat-shaped API; no tools per the LIT-5968 backfill review",
|
||||
),
|
||||
_rule("meta-models/Muse-Glimmer-30B", "reviewed for the LIT-5968 backfill; no tool or vision support documented"),
|
||||
_rule(
|
||||
"moonshotai/Kimi-K2.7-Code",
|
||||
"reviewed for the LIT-5968 backfill against https://www.together.ai/models/kimi-k2-7-code",
|
||||
**_TOOLS,
|
||||
supports_vision=True,
|
||||
),
|
||||
_rule(
|
||||
"moonshotai/Kimi-K3",
|
||||
"reviewed for the LIT-5968 backfill against https://www.together.ai/models/kimi-k3",
|
||||
**_TOOLS,
|
||||
supports_reasoning=True,
|
||||
supports_vision=True,
|
||||
),
|
||||
_rule(
|
||||
"nvidia/nemotron-3-ultra-550b-a55b",
|
||||
"reviewed for the LIT-5968 backfill against https://www.together.ai/models/nemotron-3-ultra",
|
||||
**_TOOLS,
|
||||
supports_reasoning=True,
|
||||
),
|
||||
_rule(
|
||||
"openai/gpt-oss-120b",
|
||||
"reviewed for the LIT-5968 backfill against https://www.together.ai/models/gpt-oss-120b",
|
||||
**_TOOLS,
|
||||
supports_reasoning=True,
|
||||
),
|
||||
_rule(
|
||||
"openai/gpt-oss-20b",
|
||||
"reviewed for the LIT-5968 backfill against https://www.together.ai/models/gpt-oss-20b",
|
||||
**_TOOLS,
|
||||
),
|
||||
_rule("pearl-ai/gemma-4-31b-it", "reviewed for the LIT-5968 backfill; no tool or vision support documented"),
|
||||
_rule(
|
||||
"thinkingmachines/Inkling",
|
||||
"reviewed for the LIT-5968 backfill against https://www.together.ai/models/inkling",
|
||||
**_TOOLS,
|
||||
),
|
||||
_rule("thinkingmachines/Inkling-Small", "reviewed for the LIT-5968 backfill; no tool or vision support documented"),
|
||||
_rule(
|
||||
"zai-org/GLM-5.2",
|
||||
"reviewed for the LIT-5968 backfill against https://www.together.ai/models/glm-5-2;"
|
||||
" 128K output ceiling per https://docs.z.ai/guides/llm/glm-5.2",
|
||||
**_TOOLS,
|
||||
supports_reasoning=True,
|
||||
max_output_tokens=128000,
|
||||
max_tokens=128000,
|
||||
),
|
||||
_rule(
|
||||
"zai-org/GLM-5.3-Flash",
|
||||
"reviewed for LIT-6489 against https://www.together.ai/models/glm-5-3-flash;"
|
||||
" 128K output ceiling per https://docs.z.ai/guides/llm/glm-5.3",
|
||||
**_TOOLS,
|
||||
supports_reasoning=True,
|
||||
supports_vision=True,
|
||||
max_output_tokens=128000,
|
||||
max_tokens=128000,
|
||||
),
|
||||
)
|
||||
|
||||
RULES_BY_ID: Final = MappingProxyType({rule.model_id: rule for rule in CAPABILITY_RULES})
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class DeprecationDoc:
|
||||
removal_dates: Mapping[str, str]
|
||||
redirects: Mapping[str, str]
|
||||
|
||||
|
||||
_REDIRECT_ROW: Final = re.compile(r"^\|\s*`([^`]+)`\s*\|\s*`([^`]+)`\s*\|")
|
||||
_REMOVAL_ROW: Final = re.compile(r"^\|\s*(\d{4}-\d{2}-\d{2})\s*\|\s*`([^`]+)`\s*\|")
|
||||
|
||||
|
||||
def _section(markdown: str, heading: str) -> str:
|
||||
level: Final = heading.split(" ", 1)[0]
|
||||
start: Final = markdown.find(f"\n{heading}\n")
|
||||
if start < 0:
|
||||
return ""
|
||||
body: Final = markdown[start + 1 + len(heading) :]
|
||||
next_heading: Final = re.search(rf"^{re.escape(level)} ", body, flags=re.MULTILINE)
|
||||
return body[: next_heading.start()] if next_heading else body
|
||||
|
||||
|
||||
def parse_deprecations(markdown: str) -> DeprecationDoc:
|
||||
redirect_rows: Final = tuple(
|
||||
m.groups()
|
||||
for m in (_REDIRECT_ROW.match(line) for line in _section(markdown, "## Active model redirects").splitlines())
|
||||
if m
|
||||
)
|
||||
inference: Final = _section(_section(markdown, "## Deprecation history"), "### Inference")
|
||||
removal_rows: Final = tuple(m.groups() for m in (_REMOVAL_ROW.match(line) for line in inference.splitlines()) if m)
|
||||
if not redirect_rows or not removal_rows:
|
||||
raise SyncError(
|
||||
"deprecations doc parsed to zero redirect or removal rows; the table format at "
|
||||
f"{DEPRECATIONS_URL} changed and the parser needs updating"
|
||||
)
|
||||
removal_dates: Final = {model: date for date, model in reversed(removal_rows)}
|
||||
return DeprecationDoc(
|
||||
removal_dates=MappingProxyType(dict(reversed(removal_dates.items()))),
|
||||
redirects=MappingProxyType({original: target for original, target in redirect_rows}),
|
||||
)
|
||||
|
||||
|
||||
def per_token(price_per_million: float) -> float:
|
||||
return float(f"{price_per_million / 1e6:.6g}")
|
||||
|
||||
|
||||
def _resolve_name(name: str, universe: frozenset[str]) -> str | None:
|
||||
if name in universe:
|
||||
return name
|
||||
suffix_matches: Final = tuple(candidate for candidate in universe if candidate.endswith(f"/{name}"))
|
||||
return suffix_matches[0] if len(suffix_matches) == 1 else None
|
||||
|
||||
|
||||
def resolve_successor(model_id: str, doc: DeprecationDoc, live_ids: frozenset[str]) -> str | None:
|
||||
canonical: Final = live_ids | frozenset(doc.removal_dates)
|
||||
redirects: Final = {
|
||||
(_resolve_name(raw_source, canonical) or raw_source): (_resolve_name(raw_target, canonical) or raw_target)
|
||||
for raw_source, raw_target in doc.redirects.items()
|
||||
}
|
||||
seen: Final = set()
|
||||
current = model_id # rebind-ok: walks the redirect chain
|
||||
while current in redirects and current not in seen:
|
||||
seen.add(current)
|
||||
current = redirects[current] # rebind-ok: walks the redirect chain
|
||||
return current if current != model_id and current in live_ids else None
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class SyncOutcome:
|
||||
cost_map: CostMap
|
||||
added: tuple[str, ...] = ()
|
||||
updated: tuple[str, ...] = ()
|
||||
deprecated: tuple[str, ...] = ()
|
||||
reappeared: tuple[str, ...] = ()
|
||||
warnings: tuple[str, ...] = ()
|
||||
skipped_types: Mapping[str, int] = field(default_factory=dict)
|
||||
|
||||
@property
|
||||
def has_changes(self) -> bool:
|
||||
return bool(self.added or self.updated or self.deprecated or self.reappeared)
|
||||
|
||||
|
||||
def _api_fields(model: CatalogModel) -> RegistryEntry:
|
||||
cached: Final = model.pricing.cached_input
|
||||
return {
|
||||
"input_cost_per_token": per_token(model.pricing.input),
|
||||
"output_cost_per_token": per_token(model.pricing.output),
|
||||
**({"cache_read_input_token_cost": per_token(cached), "supports_prompt_caching": True} if cached else {}),
|
||||
**({"max_input_tokens": model.context_length} if model.context_length is not None else {}),
|
||||
}
|
||||
|
||||
|
||||
def _new_entry(model: CatalogModel, mode: str) -> RegistryEntry:
|
||||
rule: Final = RULES_BY_ID.get(model.id)
|
||||
legacy_ceiling: Final = {} if model.context_length is None else {"max_tokens": model.context_length}
|
||||
merged: Final = {
|
||||
**_api_fields(model),
|
||||
**legacy_ceiling,
|
||||
"litellm_provider": PROVIDER,
|
||||
"mode": mode,
|
||||
"source": SOURCE_URL,
|
||||
**(dict(rule.fields) if rule else {}),
|
||||
}
|
||||
return dict(sorted(merged.items()))
|
||||
|
||||
|
||||
def _updated_entry(entry: RegistryEntry, model: CatalogModel) -> tuple[RegistryEntry, tuple[str, ...]]:
|
||||
rule: Final = RULES_BY_ID.get(model.id)
|
||||
desired: Final = {**_api_fields(model), **(dict(rule.fields) if rule else {})}
|
||||
dropped: Final = () if model.pricing.cached_input else ("cache_read_input_token_cost", "supports_prompt_caching")
|
||||
changes: Final = tuple(
|
||||
f"{name}: {entry.get(name)!r} -> {value!r}" for name, value in desired.items() if entry.get(name) != value
|
||||
) + tuple(
|
||||
f"{name}: {entry[name]!r} removed (no longer in the catalog pricing)" for name in dropped if name in entry
|
||||
)
|
||||
merged: Final = {name: value for name, value in {**entry, **desired}.items() if name not in dropped}
|
||||
return dict(sorted(merged.items())), changes
|
||||
|
||||
|
||||
def _with_new_keys_in_block(original: CostMap, result: CostMap, new_keys: Sequence[str]) -> CostMap:
|
||||
provider_keys: Final = tuple(key for key in original if key.startswith(PREFIX))
|
||||
if not new_keys or not provider_keys:
|
||||
return result
|
||||
block_end: Final = provider_keys[-1]
|
||||
return {
|
||||
key: value
|
||||
for existing in original
|
||||
for key, value in (
|
||||
(existing, result[existing]),
|
||||
*((new, result[new]) for new in sorted(new_keys) if existing == block_end),
|
||||
)
|
||||
}
|
||||
|
||||
|
||||
def compute_sync(cost_map: CostMap, catalog: Sequence[CatalogModel], doc: DeprecationDoc) -> SyncOutcome:
|
||||
live_ids: Final = frozenset(model.id for model in catalog)
|
||||
token_models: Final = {model.id: model for model in catalog if model.type in TYPE_TO_MODE}
|
||||
skipped: Final = {
|
||||
model.type: sum(1 for m in catalog if m.type == model.type)
|
||||
for model in catalog
|
||||
if model.type not in TYPE_TO_MODE
|
||||
}
|
||||
registry_ids: Final = {key.removeprefix(PREFIX): key for key in cost_map if key.startswith(PREFIX)}
|
||||
|
||||
added: Final[list[str]] = []
|
||||
updated: Final[list[str]] = []
|
||||
deprecated: Final[list[str]] = []
|
||||
reappeared: Final[list[str]] = []
|
||||
warnings: Final[list[str]] = []
|
||||
result: Final[CostMap] = dict(cost_map)
|
||||
|
||||
for model_id, model in sorted(token_models.items()):
|
||||
mode: Final = TYPE_TO_MODE[model.type]
|
||||
key: Final = f"{PREFIX}{model_id}"
|
||||
if model_id in doc.removal_dates:
|
||||
warnings.append(
|
||||
f"`{key}` is listed as removed on {doc.removal_dates[model_id]} in the docs but the serverless "
|
||||
"catalog still serves it; availability kept from the API"
|
||||
)
|
||||
entry = result.get(key)
|
||||
if not isinstance(entry, dict):
|
||||
result[key] = _new_entry(model, mode)
|
||||
added.append(key)
|
||||
if model.type == "chat" and model_id not in RULES_BY_ID:
|
||||
warnings.append(
|
||||
f"`{key}` added without a capability rule; review its tools/vision/reasoning support and add one"
|
||||
)
|
||||
continue
|
||||
if entry.get("mode") != mode:
|
||||
warnings.append(
|
||||
f"`{key}` has curated mode {entry.get('mode')!r} but the catalog maps to {mode!r}; left unchanged"
|
||||
)
|
||||
new_entry, changes = _updated_entry(entry, model)
|
||||
if "deprecation_date" in new_entry:
|
||||
new_entry.pop("deprecation_date")
|
||||
reappeared.append(key)
|
||||
if changes:
|
||||
updated.append(f"{key}: " + "; ".join(changes))
|
||||
if changes or key in reappeared:
|
||||
result[key] = new_entry
|
||||
|
||||
for model_id, key in sorted(registry_ids.items()):
|
||||
if model_id in token_models:
|
||||
continue
|
||||
entry = result.get(key)
|
||||
if not isinstance(entry, dict):
|
||||
continue
|
||||
removal_date: Final = doc.removal_dates.get(model_id)
|
||||
successor: Final = resolve_successor(model_id, doc, live_ids)
|
||||
metadata = entry.get("metadata")
|
||||
curated_successor: Final = metadata.get("successor") if isinstance(metadata, dict) else None
|
||||
new_entry = dict(entry)
|
||||
if removal_date is not None and entry.get("deprecation_date") != removal_date:
|
||||
if "deprecation_date" in entry:
|
||||
warnings.append(
|
||||
f"`{key}` has curated deprecation_date {entry.get('deprecation_date')!r} but the docs list "
|
||||
f"{removal_date!r}; left unchanged"
|
||||
)
|
||||
else:
|
||||
new_entry["deprecation_date"] = removal_date
|
||||
if removal_date is None and "deprecation_date" not in entry:
|
||||
warnings.append(
|
||||
f"`{key}` is absent from the serverless catalog with no removal date in the docs; "
|
||||
"needs a human deprecation call"
|
||||
)
|
||||
if successor is not None:
|
||||
desired_successor: Final = f"{PREFIX}{successor}"
|
||||
if curated_successor is None:
|
||||
new_entry["metadata"] = dict(
|
||||
sorted({**(metadata if isinstance(metadata, dict) else {}), "successor": desired_successor}.items())
|
||||
)
|
||||
elif curated_successor != desired_successor:
|
||||
warnings.append(
|
||||
f"`{key}` has curated successor {curated_successor!r} but the docs redirects resolve to "
|
||||
f"{desired_successor!r}; left unchanged"
|
||||
)
|
||||
if new_entry != entry:
|
||||
result[key] = dict(sorted(new_entry.items()))
|
||||
deprecated.append(f"{key}: " + ", ".join(sorted(set(new_entry) - set(entry)) or ["updated"]))
|
||||
|
||||
return SyncOutcome(
|
||||
cost_map=_with_new_keys_in_block(cost_map, result, tuple(added)),
|
||||
added=tuple(added),
|
||||
updated=tuple(updated),
|
||||
deprecated=tuple(deprecated),
|
||||
reappeared=tuple(reappeared),
|
||||
warnings=tuple(warnings),
|
||||
skipped_types=MappingProxyType(skipped),
|
||||
)
|
||||
|
||||
|
||||
def _section_block(title: str, lines: Sequence[str], backtick: bool) -> str:
|
||||
bullets: Final = "\n".join(f"- `{line}`" if backtick else f"- {line}" for line in lines) or "- none"
|
||||
return f"### {title} ({len(lines)})\n{bullets}\n"
|
||||
|
||||
|
||||
def render_pr_body(outcome: SyncOutcome) -> str:
|
||||
skipped: Final = ", ".join(f"{kind} ({count})" for kind, count in sorted(outcome.skipped_types.items())) or "none"
|
||||
return (
|
||||
"Automated daily sync of the together_ai entries in model_prices_and_context_window.json against "
|
||||
f"`GET {MODELS_URL}` and {DEPRECATIONS_URL} by scripts/sync_together_ai_models.py.\n"
|
||||
"\n"
|
||||
f"{_section_block('Added', outcome.added, backtick=True)}"
|
||||
"\n"
|
||||
f"{_section_block('Updated', outcome.updated, backtick=True)}"
|
||||
"\n"
|
||||
f"{_section_block('Marked deprecated', outcome.deprecated, backtick=True)}"
|
||||
"\n"
|
||||
f"{_section_block('Returned to the catalog', outcome.reappeared, backtick=True)}"
|
||||
"\n"
|
||||
f"{_section_block('Warnings needing a human call', outcome.warnings, backtick=False)}"
|
||||
"\n"
|
||||
f"Catalog model types outside the sync's token-pricing scope, skipped: {skipped}\n"
|
||||
)
|
||||
|
||||
|
||||
def render_summary(outcome: SyncOutcome) -> str:
|
||||
return (
|
||||
f"added={len(outcome.added)} updated={len(outcome.updated)} deprecated={len(outcome.deprecated)} "
|
||||
f"reappeared={len(outcome.reappeared)} warnings={len(outcome.warnings)}"
|
||||
)
|
||||
|
||||
|
||||
def load_catalog(raw: bytes) -> list[CatalogModel]:
|
||||
parsed: Final = json.loads(raw)
|
||||
entries: Final = parsed.get("data") if isinstance(parsed, dict) else parsed
|
||||
try:
|
||||
catalog: Final = CATALOG_ADAPTER.validate_python(entries)
|
||||
except ValidationError as error:
|
||||
raise SyncError(f"the catalog response no longer matches the expected shape: {error}") from error
|
||||
if not any(model.type in TYPE_TO_MODE for model in catalog):
|
||||
raise SyncError(
|
||||
"the catalog response contains no token-priced models; refusing to mark the whole registry deprecated"
|
||||
)
|
||||
return catalog
|
||||
|
||||
|
||||
def _fetch(url: str, headers: Mapping[str, str]) -> bytes:
|
||||
response: Final = httpx.get(url, headers=dict(headers), timeout=30, follow_redirects=True)
|
||||
if response.status_code != 200:
|
||||
raise SyncError(f"GET {url} returned {response.status_code}")
|
||||
return response.content
|
||||
|
||||
|
||||
def _serialize(cost_map: CostMap) -> str:
|
||||
return json.dumps(cost_map, indent=4, ensure_ascii=False) + "\n"
|
||||
|
||||
|
||||
def main(argv: Sequence[str]) -> int:
|
||||
parser: Final = argparse.ArgumentParser(description=__doc__)
|
||||
parser.add_argument("--write", action="store_true", help="apply the sync to the cost map files (default: dry run)")
|
||||
parser.add_argument("--models-json", type=Path, help="recorded catalog response to use instead of the live API")
|
||||
parser.add_argument(
|
||||
"--deprecations-md", type=Path, help="recorded deprecations doc to use instead of the live docs"
|
||||
)
|
||||
parser.add_argument("--pr-body-file", type=Path, help="write the generated PR body to this path")
|
||||
parser.add_argument("--repo-root", type=Path, default=Path(__file__).resolve().parent.parent)
|
||||
args: Final = parser.parse_args(argv)
|
||||
|
||||
if args.models_json is not None:
|
||||
catalog_raw: Final = args.models_json.read_bytes()
|
||||
else:
|
||||
api_key: Final = os.environ.get("TOGETHER_API_KEY")
|
||||
if not api_key:
|
||||
raise SyncError("TOGETHER_API_KEY is not set and --models-json was not given")
|
||||
catalog_raw = _fetch(MODELS_URL, {"Authorization": f"Bearer {api_key}"}) # rebind-ok: branch-dependent source
|
||||
catalog: Final = load_catalog(catalog_raw)
|
||||
markdown: Final = (
|
||||
args.deprecations_md.read_text() if args.deprecations_md is not None else _fetch(DEPRECATIONS_URL, {}).decode()
|
||||
)
|
||||
doc: Final = parse_deprecations(markdown)
|
||||
|
||||
cost_map_path: Final = args.repo_root / COST_MAP_RELPATHS[0]
|
||||
cost_map: Final = json.loads(cost_map_path.read_text())
|
||||
outcome: Final = compute_sync(cost_map, catalog, doc)
|
||||
body: Final = render_pr_body(outcome)
|
||||
|
||||
if args.pr_body_file is not None:
|
||||
args.pr_body_file.write_text(body)
|
||||
if args.write and outcome.has_changes:
|
||||
for relpath in COST_MAP_RELPATHS:
|
||||
(args.repo_root / relpath).write_text(_serialize(outcome.cost_map))
|
||||
print(render_summary(outcome))
|
||||
print()
|
||||
print(body)
|
||||
if not args.write:
|
||||
print("dry run: no files were touched")
|
||||
elif not outcome.has_changes:
|
||||
print("registry already in sync: no files were touched")
|
||||
return 0
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
try:
|
||||
raise SystemExit(main(sys.argv[1:]))
|
||||
except SyncError as error:
|
||||
print(f"SYNC FAILED: {error}", file=sys.stderr)
|
||||
raise SystemExit(1) from error
|
||||
|
|
@ -25,7 +25,6 @@ ALLOWED_RAW_CLIENT_FILES = {
|
|||
"e2e/e2e_http.py": ("requests",),
|
||||
"e2e/conftest.py": ("requests",),
|
||||
"e2e/claude_code/pr_gate_version_resolver.py": ("urllib.request",),
|
||||
"e2e/mcp/oauth_chat_client.py": ("httpx",),
|
||||
}
|
||||
|
||||
EXCEPTION_ONLY_NAMES = frozenset({"RequestException", "ConnectionError", "Timeout", "HTTPError"})
|
||||
|
|
|
|||
|
|
@ -1,156 +0,0 @@
|
|||
import os
|
||||
import subprocess
|
||||
import sys
|
||||
import xml.etree.ElementTree as ET
|
||||
from pathlib import Path
|
||||
from typing import Final
|
||||
|
||||
import pytest
|
||||
|
||||
GATE: Final = Path(__file__).resolve().parents[2] / ".github/e2e-stack/assert_tests_ran.py"
|
||||
SELECTED: Final = ("tests/e2e/access_control/test_a.py", "tests/e2e/access_control/test_b.py")
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("second_outcome", "expected_status"),
|
||||
(("passed", 0), ("skipped", 1), ("failure", 1), ("error", 1), ("deselected", 1)),
|
||||
)
|
||||
def test_each_changed_file_must_run(tmp_path: Path, second_outcome: str, expected_status: int) -> None:
|
||||
suite: Final = ET.Element("testsuite")
|
||||
_ = ET.SubElement(suite, "testcase", file=SELECTED[0])
|
||||
if second_outcome != "deselected":
|
||||
second: Final = ET.SubElement(suite, "testcase", file=SELECTED[1])
|
||||
if second_outcome != "passed":
|
||||
_ = ET.SubElement(second, second_outcome)
|
||||
report: Final = tmp_path / "report.xml"
|
||||
ET.ElementTree(suite).write(report)
|
||||
|
||||
result: Final = subprocess.run([sys.executable, str(GATE), str(report), *SELECTED], capture_output=True, text=True)
|
||||
|
||||
assert result.returncode == expected_status, result.stdout
|
||||
|
||||
|
||||
@pytest.mark.parametrize("outcome", ("failure", "error"))
|
||||
def test_passing_case_does_not_hide_a_failure_in_the_same_file(tmp_path: Path, outcome: str) -> None:
|
||||
suite: Final = ET.Element("testsuite")
|
||||
_ = ET.SubElement(suite, "testcase", file=SELECTED[0])
|
||||
failed: Final = ET.SubElement(suite, "testcase", file=SELECTED[0])
|
||||
_ = ET.SubElement(failed, outcome)
|
||||
report: Final = tmp_path / "report.xml"
|
||||
ET.ElementTree(suite).write(report)
|
||||
|
||||
result: Final = subprocess.run(
|
||||
[sys.executable, str(GATE), str(report), SELECTED[0]], capture_output=True, text=True
|
||||
)
|
||||
|
||||
assert result.returncode == 1
|
||||
|
||||
|
||||
def test_failed_cases_are_named_per_selected_file(tmp_path: Path) -> None:
|
||||
suite: Final = ET.Element("testsuite")
|
||||
_ = ET.SubElement(suite, "testcase", file=SELECTED[0], classname="tests.e2e.access_control.test_a", name="test_ok")
|
||||
failed: Final = ET.SubElement(
|
||||
suite, "testcase", file=SELECTED[0], classname="tests.e2e.access_control.test_a", name="test_boom"
|
||||
)
|
||||
_ = ET.SubElement(failed, "failure", message="secret-bearing message")
|
||||
errored: Final = ET.SubElement(
|
||||
suite, "testcase", file=SELECTED[1], classname="tests.e2e.access_control.test_b", name="test_setup"
|
||||
)
|
||||
_ = ET.SubElement(errored, "error")
|
||||
report: Final = tmp_path / "report.xml"
|
||||
ET.ElementTree(suite).write(report)
|
||||
|
||||
result: Final = subprocess.run([sys.executable, str(GATE), str(report), *SELECTED], capture_output=True, text=True)
|
||||
|
||||
assert result.returncode == 1
|
||||
assert " failed: tests.e2e.access_control.test_a::test_boom\n" in result.stdout
|
||||
assert " failed: tests.e2e.access_control.test_b::test_setup\n" in result.stdout
|
||||
assert "test_ok" not in result.stdout
|
||||
assert "secret-bearing message" not in result.stdout
|
||||
|
||||
|
||||
@pytest.mark.parametrize("contents", ("<testsuite/>", "<testsuite", '<testsuite><testcase name="a"/></testsuite>'))
|
||||
def test_missing_execution_evidence_fails(tmp_path: Path, contents: str) -> None:
|
||||
report: Final = tmp_path / "report.xml"
|
||||
_ = report.write_text(contents)
|
||||
|
||||
result: Final = subprocess.run([sys.executable, str(GATE), str(report), *SELECTED], capture_output=True, text=True)
|
||||
|
||||
assert result.returncode == 1
|
||||
|
||||
|
||||
@pytest.mark.parametrize("omitted_role", ("proxy_admin", "team_member", "internal_user_viewer"))
|
||||
def test_one_passing_management_case_cannot_hide_a_missing_actor(tmp_path: Path, omitted_role: str) -> None:
|
||||
suite: Final = ET.Element("testsuite")
|
||||
path: Final = "tests/e2e/management/test_jwt_management_e2e.py"
|
||||
case: Final = ET.SubElement(suite, "testcase", file=path)
|
||||
properties: Final = ET.SubElement(case, "properties")
|
||||
_ = ET.SubElement(
|
||||
properties,
|
||||
"property",
|
||||
name="management_node",
|
||||
value=f"{path}::TestJwtManagement::test_actor_subject_and_database_role[proxy_admin_viewer]",
|
||||
)
|
||||
report: Final = tmp_path / "report.xml"
|
||||
ET.ElementTree(suite).write(report)
|
||||
result: Final = subprocess.run([sys.executable, str(GATE), str(report), path], capture_output=True, text=True)
|
||||
assert result.returncode == 1
|
||||
assert f"test_actor_subject_and_database_role[{omitted_role}]" in result.stdout
|
||||
|
||||
|
||||
@pytest.mark.parametrize("phase", ("setup", "call", "teardown"))
|
||||
@pytest.mark.parametrize("required_count", ("1", "4"))
|
||||
def test_oauth_failure_diagnostics_do_not_publish_private_payloads(
|
||||
tmp_path: Path, phase: str, required_count: str
|
||||
) -> None:
|
||||
suite: Final = ET.Element("testsuite")
|
||||
case: Final = ET.SubElement(suite, "testcase", file=SELECTED[0])
|
||||
private: Final = "private-token-in-exception-message"
|
||||
failure: Final = ET.SubElement(case, "failure", message=private)
|
||||
failure.text = private
|
||||
properties: Final = ET.SubElement(case, "properties")
|
||||
for name, value in (
|
||||
("oauth_failure_phase", phase),
|
||||
("oauth_exception_type", "AssertionError"),
|
||||
("oauth_frame", "oauth_gateway.py:120:start"),
|
||||
("oauth_frame", f"injected\\n{private}"),
|
||||
("unrelated_property", private),
|
||||
):
|
||||
_ = ET.SubElement(properties, "property", name=name, value=value)
|
||||
report: Final = tmp_path / "report.xml"
|
||||
ET.ElementTree(suite).write(report)
|
||||
result: Final = subprocess.run(
|
||||
[sys.executable, "-I", str(GATE), str(report), SELECTED[0]],
|
||||
capture_output=True,
|
||||
text=True,
|
||||
env={**os.environ, "E2E_REQUIRED_TEST_COUNT": required_count},
|
||||
)
|
||||
assert result.returncode == 1
|
||||
assert f"oauth_failure_phase: {phase}" in result.stdout
|
||||
assert "oauth_exception_type: AssertionError" in result.stdout
|
||||
assert "oauth_frame: oauth_gateway.py:120:start" in result.stdout
|
||||
assert private not in result.stdout + result.stderr
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("count", "skip", "expected"), ((0, False, 1), (3, False, 1), (4, False, 0), (5, False, 1), (4, True, 1))
|
||||
)
|
||||
def test_required_count_reports_cases_before_rejecting(tmp_path: Path, count: int, skip: bool, expected: int) -> None:
|
||||
suite = ET.Element("testsuite")
|
||||
for index in range(count):
|
||||
case = ET.SubElement(suite, "testcase", file=SELECTED[0], classname="OAuth", name=f"variant{index}")
|
||||
if skip and index == 0:
|
||||
ET.SubElement(case, "skipped", message="private-skip-reason")
|
||||
report = tmp_path / "report.xml"
|
||||
ET.ElementTree(suite).write(report)
|
||||
result = subprocess.run(
|
||||
[sys.executable, "-I", str(GATE), str(report), SELECTED[0]],
|
||||
env={**os.environ, "E2E_REQUIRED_TEST_COUNT": "4"},
|
||||
capture_output=True,
|
||||
text=True,
|
||||
)
|
||||
assert result.returncode == expected
|
||||
assert f"{count} collected, {int(skip)} skipped" in result.stdout
|
||||
if skip:
|
||||
assert "skipped: OAuth::variant0" in result.stdout
|
||||
assert "private-skip-reason" not in result.stdout + result.stderr
|
||||
|
|
@ -172,12 +172,7 @@ def open_consent() -> None:
|
|||
raise RuntimeError("consent page timed out")
|
||||
|
||||
|
||||
@pytest.mark.mcp_oauth_live
|
||||
def test_oauth_dies_on_consent() -> None:
|
||||
open_consent()
|
||||
|
||||
|
||||
def test_plain_dies_on_consent() -> None:
|
||||
def test_dies_on_consent() -> None:
|
||||
open_consent()
|
||||
"""
|
||||
|
||||
|
|
@ -353,16 +348,13 @@ class TestStepsReachTheReport:
|
|||
|
||||
def test_a_failed_phase_s_own_report_carries_the_steps(tmp_path: Path) -> None:
|
||||
"""Plugins that read the failed setup or call report, not the teardown one
|
||||
junitxml writes from, see where the test died too, oauth-live or not."""
|
||||
junitxml writes from, see where the test died too."""
|
||||
write_suite(tmp_path, {"test_consent.py": FAILED_PHASE_SUITE, "report_spy.py": REPORT_SPY_PLUGIN})
|
||||
child: Final = run_child_pytest(tmp_path, "-p", "report_spy", env={"E2E_MCP_OAUTH_LIVE": "1"})
|
||||
child: Final = run_child_pytest(tmp_path, "-p", "report_spy")
|
||||
seen_path: Final = tmp_path / "failed-reports.jsonl"
|
||||
assert seen_path.exists(), f"no failed report reached the spy:\n{child.stdout}\n{child.stderr}"
|
||||
seen: Final = dict(map(FailedReport.validate_json, seen_path.read_text().splitlines()))
|
||||
assert seen == {
|
||||
"test_oauth_dies_on_consent": ("open the consent page",),
|
||||
"test_plain_dies_on_consent": ("open the consent page",),
|
||||
}, child.stdout
|
||||
assert seen == {"test_dies_on_consent": ("open the consent page",)}, child.stdout
|
||||
|
||||
|
||||
class TestDeclaredPropertiesReachTheReport:
|
||||
|
|
|
|||
File diff suppressed because one or more lines are too long
|
|
@ -250,55 +250,3 @@ The semantic header set is `content-type`, `accept`, `anthropic-version`, `anthr
|
|||
Excluded transport and telemetry headers are `host`, `content-length`, `connection`, `accept-encoding`, `user-agent`, `traceparent`, `tracestate`, `x-request-id`, `x-client-request-id` and `x-stainless-*`. Inbound transfer-encoding is unsupported; send JSON with content-length framing. The destination represents host identity and the relay carries original body bytes. Replay does not verify credentials, SDK timeout/retry behavior, transport performance, model availability or stateful remote IDs. Live relay uses original request bytes and header values, never the stored identity
|
||||
|
||||
Strict replay harness regression tests live in `tests/code_coverage_tests/test_provider_replay_harness.py`. The `python` job in `.github/workflows/test-linting.yml` runs them alongside the provider-edge and fixture tests at the root of `tests/e2e_harness/` with `--noconftest -o addopts= -o "pythonpath=tests/e2e tests/e2e_harness" -p no:rerunfailures`; they need only synthetic HTTP providers and temporary fixture storage
|
||||
|
||||
|
||||
## MCP OAuth happy path
|
||||
|
||||
`test_mcp_oauth_happy_path_e2e.py` runs one shared scenario with four variants:
|
||||
aggregate gateway SSO and explicitly configured per-server JWT, each directly
|
||||
against Linear and through the live provider edge. The edge forwards to real
|
||||
Linear without replay and compares the forwarded bearer to the encrypted
|
||||
canonical user/server credential. This observes the forwarding boundary, not
|
||||
Linear's internal logs. Direct variants independently exercise discovery
|
||||
|
||||
Use the existing database preparation, Prisma generation and Keycloak setup.
|
||||
Build and stage the dashboard from the tested checkout as in the UI runner.
|
||||
Provide `DATABASE_URL`, `LITELLM_MASTER_KEY`, `LITELLM_SALT_KEY`, `LITELLM_LICENSE`,
|
||||
and the `E2E_KEYCLOAK_*` settings. Capture a test-account Linear login using
|
||||
`mcp/linear_session_capture.py` and set `E2E_LINEAR_STORAGE_STATE` to that private
|
||||
file. The test workspace must contain a team. Do not publish browser state or
|
||||
raw test/proxy output
|
||||
|
||||
```bash
|
||||
E2E_MCP_OAUTH_LIVE=1 E2E_FIXTURE_MODE=live E2E_PROVIDER_CACHE=0 \
|
||||
uv run --no-sync pytest tests/e2e/mcp/test_mcp_oauth_happy_path_e2e.py \
|
||||
--rootdir=. --reruns 0
|
||||
```
|
||||
|
||||
The test starts and restarts its own source-built proxy on a free loopback port,
|
||||
retaining its database and SSO client but no Redis or process-local cache. It
|
||||
does not restart an existing proxy or clear shared databases. Gateway login,
|
||||
consent, immediate list/call and post-restart reconnect must all succeed. The
|
||||
aggregate client never injects a gateway header; the explicitly labeled JWT
|
||||
variant configures `x-litellm-api-key` for the first consent and reconnects with
|
||||
only its gateway JWT after restart
|
||||
|
||||
`.github/workflows/test-mcp-oauth-e2e.yml` automatically requests a run for
|
||||
same-repository pull requests changing MCP, gateway authentication/SSO, consent
|
||||
UI, dependencies or the relevant E2E harness/workflow paths. It retains manual
|
||||
`workflow_dispatch` for targeted verification. The four cases run in the
|
||||
protected `e2e-changed` environment after its normal deployment approval;
|
||||
reviewers should approve and inspect this separate OAuth check when it appears.
|
||||
Fork pull requests do not run this credentialed job; use a reviewed
|
||||
same-repository branch for their verification. The workflow's path-filtered
|
||||
check is not configured here as a globally required branch-protection check.
|
||||
Provision `E2E_LINEAR_STORAGE_STATE_B64` as a secret there and retain the existing E2E license/AWS role configuration. A missing or
|
||||
expired session fails the job; collection, deselection and skips are not passes.
|
||||
No LLM call is needed
|
||||
|
||||
Coverage remains limited to authorization-code OAuth over HTTP. M2M, OBO,
|
||||
PKCE passthrough, static/BYOK, ID-JAG, forwarding, SigV4 and stdio are outside this
|
||||
scenario; consult the registry and LIT-3559 for their existing coverage and gaps.
|
||||
LIT-4506 owns broader isolation/failure regressions. LIT-7737 retains ownership
|
||||
of dependency/Python compatibility and its matrix; this test reuses its delivered
|
||||
environment and does not change dependency constraints or compatibility gates
|
||||
|
|
|
|||
|
|
@ -597,14 +597,21 @@ def _build_control_plane_client(proxy_config: ProxyConfig):
|
|||
)
|
||||
|
||||
|
||||
def _register_deployment(proxy, deployment: CompatDeployment) -> str:
|
||||
"""Register one deployment and return its proxy-assigned model_id
|
||||
once it is servable on the data plane. The aliases are shared by every
|
||||
cell and, under xdist, by every worker, so no call to them belongs to
|
||||
one test and none is cached: the matrix exists to reach real providers."""
|
||||
return proxy.create_model(
|
||||
deployment.model_name,
|
||||
deployment.litellm_params,
|
||||
def _register_deployments(proxy, deployments: tuple[CompatDeployment, ...]):
|
||||
"""The aliases are shared by every cell and, under xdist, by every worker,
|
||||
so no call to them belongs to one test and none is cached: the matrix
|
||||
exists to reach real providers."""
|
||||
from models import ModelInfoBody, ModelNewBody
|
||||
|
||||
return proxy.register_models(
|
||||
tuple(
|
||||
ModelNewBody(
|
||||
model_name=deployment.model_name,
|
||||
litellm_params=deployment.litellm_params,
|
||||
model_info=ModelInfoBody(),
|
||||
)
|
||||
for deployment in deployments
|
||||
),
|
||||
provider_live=True,
|
||||
)
|
||||
|
||||
|
|
@ -632,18 +639,14 @@ def _compat_models_registered() -> Any:
|
|||
yield
|
||||
return
|
||||
|
||||
from proxy_client import ModelNotRegistered, ModelRegistered
|
||||
from requests import RequestException
|
||||
|
||||
proxy = _build_control_plane_client(proxy_config)
|
||||
registered_ids: list[str] = []
|
||||
failures: list[tuple[str, str]] = []
|
||||
registrations = _register_deployments(proxy, load_all_deployments())
|
||||
registered_ids = tuple(r.model_id for r in registrations if isinstance(r, ModelRegistered))
|
||||
failures = tuple((r.model_name, r.reason) for r in registrations if isinstance(r, ModelNotRegistered))
|
||||
try:
|
||||
for deployment in load_all_deployments():
|
||||
try:
|
||||
model_id = _register_deployment(proxy, deployment)
|
||||
registered_ids.append(model_id)
|
||||
except (AssertionError, RequestException) as exc:
|
||||
failures.append((deployment.model_name, str(exc)))
|
||||
if failures:
|
||||
summary = "\n".join(
|
||||
f" - {name}: {reason}" for name, reason in failures
|
||||
|
|
|
|||
|
|
@ -17,7 +17,6 @@ import functools
|
|||
import os
|
||||
from collections.abc import Generator, Iterator
|
||||
from datetime import datetime, timezone
|
||||
from pathlib import Path
|
||||
from types import MappingProxyType
|
||||
from typing import Final
|
||||
|
||||
|
|
@ -29,7 +28,6 @@ from e2e_config import (
|
|||
FIXTURE_DIR,
|
||||
FIXTURE_MODE_RAW,
|
||||
MANAGED_FILES_OPT_IN_ENV,
|
||||
MCP_OAUTH_LIVE_OPT_IN_ENV,
|
||||
OTEL_TLS_OPT_IN_ENV,
|
||||
OTEL_V2_OPT_IN_ENV,
|
||||
OWNED_GATEWAY_OPT_IN_ENV,
|
||||
|
|
@ -37,6 +35,7 @@ from e2e_config import (
|
|||
PROMPT_CACHING_OPT_IN_ENV,
|
||||
PROVIDER_EDGE_HOST_OPT_IN_ENV,
|
||||
PROXY_BASE_URL,
|
||||
PROXY_WIDE_CONFIG_OPT_IN_ENV,
|
||||
SAGEMAKER_NOVA2_LITE_OPT_IN_ENV,
|
||||
SAGEMAKER_NOVA_OPT_IN_ENV,
|
||||
SECRET_MANAGER_OPT_IN_ENV,
|
||||
|
|
@ -56,7 +55,6 @@ from models import TeamNewBody, UserNewBody, UserNewResponse
|
|||
from provider_cache_routing import LIVE_PROVIDER_REQUIRED
|
||||
from provider_edge import replay_leftover_error
|
||||
from proxy_client import ProxyClient, build_proxy_client
|
||||
from stack_lock import stack_lock
|
||||
|
||||
_E2E_TEST_RAN = pytest.StashKey[bool]()
|
||||
_CALL_PASSED = pytest.StashKey[bool]()
|
||||
|
|
@ -72,7 +70,6 @@ OPT_IN_MARKERS: Final = MappingProxyType(
|
|||
"redis_chaos": PERF_TESTS_OPT_IN_ENV,
|
||||
"perf": PERF_TESTS_OPT_IN_ENV,
|
||||
"cli_determinism": CLI_DETERMINISM_OPT_IN_ENV,
|
||||
"mcp_oauth_live": MCP_OAUTH_LIVE_OPT_IN_ENV,
|
||||
"provider_edge_host": PROVIDER_EDGE_HOST_OPT_IN_ENV,
|
||||
"owned_gateway": OWNED_GATEWAY_OPT_IN_ENV,
|
||||
"otel_v2": OTEL_V2_OPT_IN_ENV,
|
||||
|
|
@ -80,6 +77,7 @@ OPT_IN_MARKERS: Final = MappingProxyType(
|
|||
"secret_manager": SECRET_MANAGER_OPT_IN_ENV,
|
||||
"sagemaker_nova": SAGEMAKER_NOVA_OPT_IN_ENV,
|
||||
"sagemaker_nova2_lite": SAGEMAKER_NOVA2_LITE_OPT_IN_ENV,
|
||||
"proxy_wide_config": PROXY_WIDE_CONFIG_OPT_IN_ENV,
|
||||
}
|
||||
)
|
||||
|
||||
|
|
@ -165,16 +163,6 @@ def pytest_configure(config: pytest.Config) -> None:
|
|||
"redis_chaos: load test that pauses the proxy's Redis outright mid-run; needs a proxy booted from "
|
||||
"gateway/redis_chaos_ci_config.yml on the same host, and is deselected unless E2E_PERF_TESTS is set",
|
||||
)
|
||||
config.addinivalue_line(
|
||||
"markers",
|
||||
"quiet_stack: measures the proxy itself, so it runs while no other test on this host is hitting the stack; "
|
||||
"every other test waits for it to finish",
|
||||
)
|
||||
config.addinivalue_line(
|
||||
"markers",
|
||||
"mcp_oauth_live: real Linear OAuth consent via a captured browser session; deselected unless "
|
||||
"E2E_MCP_OAUTH_LIVE is set",
|
||||
)
|
||||
config.addinivalue_line(
|
||||
"markers",
|
||||
"provider_edge_host: routes provider traffic through the pytest host's edge in every fixture mode, so the "
|
||||
|
|
@ -208,6 +196,11 @@ def pytest_configure(config: pytest.Config) -> None:
|
|||
"sagemaker_nova2_lite: needs a deployed SageMaker Nova 2 Lite endpoint the proxy's AWS credentials can invoke; "
|
||||
"deselected unless SAGEMAKER_NOVA2_LITE_ENDPOINT names it",
|
||||
)
|
||||
config.addinivalue_line(
|
||||
"markers",
|
||||
"proxy_wide_config: turns on a proxy-wide setting through /config/update, so it runs on a stack no other suite "
|
||||
"is driving; deselected unless E2E_PROXY_WIDE_CONFIG_STACK is set",
|
||||
)
|
||||
|
||||
|
||||
def pytest_sessionstart(session: pytest.Session) -> None:
|
||||
|
|
@ -277,7 +270,7 @@ def pytest_collection_finish(session: pytest.Session) -> None:
|
|||
ahead of xdist's own collection-finish report, and the controller schedules no
|
||||
test until every worker has reported, so this is the idle footprint of a stack
|
||||
that just passed its readiness gate. The fixture hands the capture to the
|
||||
idle-budget test in router/test_reliability_memory_e2e.py."""
|
||||
idle-budget test in perf/test_reliability_memory_e2e.py."""
|
||||
if session.config.getoption("collectonly") or not any(_uses_idle_rss(item) for item in session.items):
|
||||
return
|
||||
session.config.stash[_IDLE_RSS] = read_rss_everywhere(build_proxy_client(), timeout=IDLE_RSS_READ_TIMEOUT_SECONDS)
|
||||
|
|
@ -307,12 +300,6 @@ def _proxy_fail_reason() -> str | None:
|
|||
return None
|
||||
|
||||
|
||||
@pytest.hookimpl(wrapper=True)
|
||||
def pytest_runtest_protocol(item: pytest.Item, nextitem: pytest.Item | None) -> Generator[None, object, object]:
|
||||
with stack_lock(exclusive=item.get_closest_marker("quiet_stack") is not None):
|
||||
return (yield)
|
||||
|
||||
|
||||
@pytest.hookimpl(tryfirst=True)
|
||||
def pytest_runtest_setup(item: pytest.Item) -> None:
|
||||
"""Hard-fail `e2e`-marked tests unless a proxy answers its liveness probe.
|
||||
|
|
@ -331,8 +318,6 @@ def pytest_runtest_setup(item: pytest.Item) -> None:
|
|||
item.user_properties.extend(item.config.stash[_IDLE_RSS].junit_properties)
|
||||
if not _reaches_proxy(item):
|
||||
return
|
||||
if isinstance(item, pytest.Function) and "oauth_gateway" in item.fixturenames:
|
||||
return
|
||||
reason = _proxy_fail_reason()
|
||||
if reason is not None:
|
||||
pytest.fail(reason)
|
||||
|
|
@ -376,12 +361,6 @@ def pytest_runtest_makereport(
|
|||
report = yield
|
||||
if report.when in ("setup", "call"):
|
||||
attach_step_properties(item)
|
||||
if item.get_closest_marker("mcp_oauth_live") is not None and call.excinfo is not None:
|
||||
# Publish code locations only, never exception messages, source text or locals.
|
||||
item.user_properties.append(("oauth_failure_phase", report.when))
|
||||
item.user_properties.append(("oauth_exception_type", call.excinfo.type.__name__))
|
||||
for entry in call.excinfo.traceback:
|
||||
item.user_properties.append(("oauth_frame", f"{Path(entry.path).name}:{entry.lineno + 1}:{entry.name}"))
|
||||
if report.when == "call":
|
||||
item.stash[_CALL_PASSED] = report.passed
|
||||
report.user_properties = list(item.user_properties)
|
||||
|
|
|
|||
|
|
@ -79,14 +79,6 @@
|
|||
assertions: [succeeds]
|
||||
source: "db.py user_oauth_credential lookup"
|
||||
rationale: OAuth2 token passthrough; per-user credential storage
|
||||
- id: mcp.call_tool.oauth.persists_across_processes
|
||||
module: mcp
|
||||
tier: P1
|
||||
operation: call_tool
|
||||
auth_family: oauth
|
||||
assertions: [persists_across_processes]
|
||||
source: "outbound_credentials/per_user_oauth_store.py V2PerUserTokenStore"
|
||||
rationale: Stored per-user token survives a verified restart of an owned gateway with no Redis cache
|
||||
- id: mcp.list_tools.none.succeeds
|
||||
module: mcp
|
||||
tier: P1
|
||||
|
|
|
|||
|
|
@ -44,8 +44,6 @@
|
|||
- {id: quota_management.spend_tracking.messages_bridge.logs_cost, module: quota_management, tier: P1, behavior: spend_tracking, variant: messages_bridge, assertions: [logs_cost], exercised_on: [messages], source: "llms/anthropic/pass_through/responses_adapters/handler.py", rationale: "A streaming /v1/messages request served by an openai-provider model is bridged through the anthropic-messages -> Responses adapter and must aggregate the consumed SSE stream into one spend row with nonzero cost and token counts, attributed to custom_llm_provider openai under call_type anthropic_messages"}
|
||||
- {id: quota_management.spend_tracking.embeddings.logs_cost, module: quota_management, tier: P1, behavior: spend_tracking, variant: embeddings, assertions: [logs_cost], exercised_on: [embeddings], source: "proxy/spend_tracking/spend_tracking_utils.py", rationale: "Embedding calls write nonzero spend rows"}
|
||||
- {id: quota_management.spend_tracking.cache_hit.zero_cost, module: quota_management, tier: P1, behavior: spend_tracking, variant: cache_hit, assertions: [zero_cost], exercised_on: [chat_completions], source: "proxy/spend_tracking/spend_tracking_utils.py", rationale: "A response-cache hit logs at zero cost with the cache-hit marker"}
|
||||
- {id: quota_management.spend_tracking.discount_config.logs_cost, module: quota_management, tier: P1, behavior: spend_tracking, variant: discount_config, assertions: [logs_cost], exercised_on: [chat_completions], source: "proxy/management_endpoints/cost_tracking_settings.py + cost_calculator.py", rationale: "A persisted provider discount changes the real request's spend row"}
|
||||
- {id: quota_management.spend_tracking.margin_config.logs_cost, module: quota_management, tier: P1, behavior: spend_tracking, variant: margin_config, assertions: [logs_cost], exercised_on: [chat_completions], source: "proxy/management_endpoints/cost_tracking_settings.py + cost_calculator.py", rationale: "A persisted provider margin is added to the real request's spend row"}
|
||||
- {id: quota_management.spend_tracking.guardrail_cost.logs_cost, module: quota_management, tier: P1, behavior: spend_tracking, variant: guardrail_cost, assertions: [logs_cost], exercised_on: [chat_completions], source: "litellm_core_utils/llm_cost_calc/guardrail_cost.py + proxy/spend_tracking/spend_tracking_utils.py", rationale: "A priced Bedrock guardrail charge reaches the persisted spend row"}
|
||||
- {id: quota_management.spend_tracking.key_rollup.matches_sum_of_logs, module: quota_management, tier: P1, behavior: spend_tracking, variant: key_rollup, assertions: [matches_sum_of_logs], exercised_on: [chat_completions], source: "proxy/db/db_spend_update_writer.py", rationale: "A key's rolled-up spend equals the sum of its log rows"}
|
||||
- {id: quota_management.spend_tracking.concurrent_burst.loses_no_spend, module: quota_management, tier: P1, behavior: spend_tracking, variant: concurrent_burst, assertions: [loses_no_spend], exercised_on: [chat_completions], source: "proxy/db/db_spend_update_writer.py", rationale: "Concurrent calls all land as spend; no row lost to write contention"}
|
||||
|
|
|
|||
|
|
@ -111,10 +111,6 @@ CHEAP_ANTHROPIC_MODEL = os.environ.get("E2E_CHEAP_ANTHROPIC_MODEL", "claude-haik
|
|||
CHEAP_OPENAI_MODEL = os.environ.get("E2E_CHEAP_OPENAI_MODEL", "gpt-5.5")
|
||||
S3_PARTITION_GRANULARITY = os.environ.get("E2E_S3_PARTITION_GRANULARITY", "day")
|
||||
|
||||
LINEAR_MCP_URL = os.environ.get("E2E_LINEAR_MCP_URL", "https://mcp.linear.app/mcp")
|
||||
LINEAR_STORAGE_STATE = os.environ.get("E2E_LINEAR_STORAGE_STATE", "")
|
||||
LINEAR_READONLY_TOOL: Final = "list_teams" # as listed by tools/list on mcp.linear.app when PR #33787 landed
|
||||
|
||||
# Jaeger query API of the compose stack's OTEL trace destination (the `jaeger`
|
||||
# service in docker-compose.yml maps it to host 16686). Trace-completeness tests
|
||||
# read exported spans back through it.
|
||||
|
|
@ -207,7 +203,6 @@ MANAGED_FILES_OPT_IN_ENV = "E2E_MANAGED_FILES_STACK"
|
|||
PROMPT_CACHING_OPT_IN_ENV = "E2E_PROMPT_CACHING_STACK"
|
||||
PERF_TESTS_OPT_IN_ENV: Final = "E2E_PERF_TESTS"
|
||||
CLI_DETERMINISM_OPT_IN_ENV = "E2E_CLI_DETERMINISM"
|
||||
MCP_OAUTH_LIVE_OPT_IN_ENV: Final = "E2E_MCP_OAUTH_LIVE"
|
||||
PROVIDER_EDGE_HOST_OPT_IN_ENV: Final = "E2E_PROVIDER_EDGE_HOST_REACHABLE"
|
||||
OWNED_GATEWAY_OPT_IN_ENV: Final = "E2E_OWNED_GATEWAY"
|
||||
OTEL_V2_OPT_IN_ENV: Final = "E2E_OTEL_V2"
|
||||
|
|
@ -215,6 +210,7 @@ OTEL_TLS_OPT_IN_ENV: Final = "E2E_OTEL_EXPORTER_ENDPOINT"
|
|||
SECRET_MANAGER_OPT_IN_ENV: Final = "E2E_SECRET_MANAGER"
|
||||
SAGEMAKER_NOVA_OPT_IN_ENV: Final = "SAGEMAKER_NOVA_ENDPOINT"
|
||||
SAGEMAKER_NOVA2_LITE_OPT_IN_ENV: Final = "SAGEMAKER_NOVA2_LITE_ENDPOINT"
|
||||
PROXY_WIDE_CONFIG_OPT_IN_ENV: Final = "E2E_PROXY_WIDE_CONFIG_STACK"
|
||||
ANOMALY_SESSIONS = int(os.environ.get("E2E_ANOMALY_SESSIONS", "6"))
|
||||
ANOMALY_TURNS_PER_SESSION = int(os.environ.get("E2E_ANOMALY_TURNS_PER_SESSION", "6"))
|
||||
ANOMALY_TURN_ATTEMPTS = int(os.environ.get("E2E_ANOMALY_TURN_ATTEMPTS", "3"))
|
||||
|
|
@ -239,7 +235,6 @@ MEMORY_RSS_SETTLE_SAMPLES = int(os.environ.get("E2E_MEMORY_RSS_SETTLE_SAMPLES",
|
|||
MEMORY_RSS_SAMPLE_INTERVAL_SECONDS = float(os.environ.get("E2E_MEMORY_RSS_SAMPLE_INTERVAL_SECONDS", "1"))
|
||||
MEMORY_RSS_CEILING_MB = float(os.environ.get("E2E_MEMORY_RSS_CEILING_MB", "200"))
|
||||
MEMORY_IDLE_RSS_BUDGET_MB = float(os.environ.get("E2E_MEMORY_IDLE_RSS_BUDGET_MB", "768"))
|
||||
MEMORY_STORED_REQUEST_BUDGET_KB = float(os.environ.get("E2E_MEMORY_STORED_REQUEST_BUDGET_KB", "64"))
|
||||
|
||||
|
||||
def ws_base_url() -> str:
|
||||
|
|
|
|||
|
|
@ -1,58 +0,0 @@
|
|||
"""One-time helper to capture a logged-in Linear browser session for the
|
||||
real-Linear MCP e2e test.
|
||||
|
||||
The real-Linear test drives the genuine gateway-managed authorization_code
|
||||
dance against ``mcp.linear.app``. The only step that cannot be scripted is
|
||||
Linear's login (magic link / SSO), so a human authenticates once here and the
|
||||
resulting session (cookies + local storage) is persisted to disk. The e2e test
|
||||
then loads that session in a headless Playwright context and clicks Approve on
|
||||
Linear's consent screen every run, with no human and no login automation.
|
||||
|
||||
Run it with the e2e venv, log into Linear in the window that opens, then return
|
||||
to the terminal and press Enter:
|
||||
|
||||
LITELLM=~/litellm-mcpe2e
|
||||
"$LITELLM"/.venv/bin/python "$LITELLM"/tests/e2e/mcp/linear_session_capture.py
|
||||
|
||||
The session is written to ``E2E_LINEAR_STORAGE_STATE`` (default
|
||||
``~/.litellm-e2e/linear_storage_state.json``), outside the repo. It is a
|
||||
secret: never commit it. Re-run this whenever Linear expires the session.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
from pathlib import Path
|
||||
|
||||
from playwright.sync_api import sync_playwright
|
||||
|
||||
DEFAULT_STATE_PATH = Path.home() / ".litellm-e2e" / "linear_storage_state.json"
|
||||
|
||||
|
||||
def capture(state_path: Path) -> None:
|
||||
"""Open a headed browser at Linear, wait for the human to log in, then save
|
||||
the authenticated session to ``state_path``."""
|
||||
state_path.parent.mkdir(parents=True, exist_ok=True)
|
||||
with sync_playwright() as playwright:
|
||||
browser = playwright.chromium.launch(headless=False)
|
||||
context = browser.new_context()
|
||||
page = context.new_page()
|
||||
page.goto("https://linear.app/login", wait_until="domcontentloaded")
|
||||
print("\n" + "=" * 72)
|
||||
print("Log into Linear in the browser window that just opened.")
|
||||
print("If Linear emails you a magic link, paste the link into THIS window's")
|
||||
print("address bar (opening it in your default browser won't capture the")
|
||||
print("session). Google SSO works too as long as you complete it here.")
|
||||
print("When your Linear workspace has loaded, come back and press Enter.")
|
||||
print("=" * 72)
|
||||
input("Press Enter once you are logged in... ")
|
||||
page.goto("https://mcp.linear.app/", wait_until="domcontentloaded")
|
||||
context.storage_state(path=str(state_path))
|
||||
browser.close()
|
||||
print(f"\nSaved Linear session to {state_path}")
|
||||
print("Point the e2e test at it with:")
|
||||
print(f' export E2E_LINEAR_STORAGE_STATE="{state_path}"')
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
capture(Path(os.environ.get("E2E_LINEAR_STORAGE_STATE", str(DEFAULT_STATE_PATH))))
|
||||
|
|
@ -1,442 +0,0 @@
|
|||
"""Client for the mcp chat-completion OAuth e2e suite.
|
||||
|
||||
Registers a gateway-managed OAuth (authorization_code) MCP server, seeds the
|
||||
per-user upstream token by driving the interactive authorize dance with the
|
||||
official mcp SDK's OAuthClientProvider (the browser leg is a headless Chromium
|
||||
primed with a human's saved Linear session), then exercises the server through
|
||||
/chat/completions, where the gateway lists and executes its tools with the
|
||||
stored per-user token.
|
||||
|
||||
Management routes (/v1/mcp/server CRUD, /chat/completions) go through the
|
||||
shared ProxyClient transport. The MCP protocol used to seed the token goes through
|
||||
the mcp SDK, the same library production MCP hosts run.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import re
|
||||
import time
|
||||
from dataclasses import dataclass
|
||||
from typing import TYPE_CHECKING, Final
|
||||
from urllib.parse import parse_qsl
|
||||
|
||||
import httpx
|
||||
import httpx2
|
||||
import pytest
|
||||
from e2e_config import PROXY_BASE_URL, REQUEST_TIMEOUT
|
||||
from e2e_http import AuthHeaders, NoBody, unwrap
|
||||
from e2e_metadata import step
|
||||
from idp import Identity
|
||||
from mcp import ClientSession
|
||||
from mcp.client.auth import OAuthClientProvider
|
||||
from mcp.client.streamable_http import streamable_http_client
|
||||
from mcp.shared.auth import AuthorizationCodeResult, OAuthClientInformationFull, OAuthClientMetadata, OAuthToken
|
||||
from mcp.types import TextContent
|
||||
from models import (
|
||||
ChatBody,
|
||||
ChatResponse,
|
||||
McpOauthUserCredentialStatus,
|
||||
McpServerCreateBody,
|
||||
McpServerInfo,
|
||||
McpServerUserCredentialListResponse,
|
||||
McpServerUserCredentialRow,
|
||||
)
|
||||
from proxy_client import ProxyClient
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from playwright.async_api import Route
|
||||
|
||||
# Where the "browser" lands at the end of the authorize dance. Nothing listens
|
||||
# here: the route interceptor short-circuits the final redirect and reads the
|
||||
# code/state off its query string, exactly like a desktop MCP host intercepting
|
||||
# its loopback redirect.
|
||||
OAUTH_CLIENT_REDIRECT_URI = "http://127.0.0.1:53682/e2e/callback"
|
||||
BROWSER_CONSENT_TIMEOUT = 60.0
|
||||
|
||||
|
||||
def _mcp_url(alias: str, base_url: str = PROXY_BASE_URL) -> str:
|
||||
return f"{base_url.rstrip('/')}/{alias}/mcp"
|
||||
|
||||
|
||||
class InMemoryTokenStorage:
|
||||
"""The mcp SDK's TokenStorage protocol, in memory for one dance: the
|
||||
DCR-registered client and the gateway tokens minted for it."""
|
||||
|
||||
def __init__(self) -> None:
|
||||
self._tokens: OAuthToken | None = None
|
||||
self._client_info: OAuthClientInformationFull | None = None
|
||||
|
||||
async def get_tokens(self) -> OAuthToken | None:
|
||||
return self._tokens
|
||||
|
||||
async def set_tokens(self, tokens: OAuthToken) -> None:
|
||||
self._tokens = tokens
|
||||
|
||||
async def get_client_info(self) -> OAuthClientInformationFull | None:
|
||||
return self._client_info
|
||||
|
||||
async def set_client_info(self, client_info: OAuthClientInformationFull) -> None:
|
||||
self._client_info = client_info
|
||||
|
||||
|
||||
async def _browser_follow_authorize(
|
||||
start_url: str,
|
||||
storage_state_path: str,
|
||||
identity: Identity | None = None,
|
||||
server_alias: str | None = None,
|
||||
allow_upstream_consent: bool = True,
|
||||
) -> tuple[str, str | None]:
|
||||
"""Play the browser's role for a real upstream whose authorize endpoint
|
||||
serves an interactive consent page (Linear). A headless Chromium primed
|
||||
with a human's saved Linear session opens the gateway authorize URL and
|
||||
clicks through Linear's consent screens (the mcp.linear.app Approve form,
|
||||
then the linear.app workspace-selection page), riding the rest of the chain
|
||||
(Linear -> gateway callback -> host redirect_uri). The final hop is
|
||||
intercepted and short-circuited, since nothing listens there, and its
|
||||
code/state are read off the query string."""
|
||||
from playwright.async_api import async_playwright
|
||||
|
||||
captured: dict[str, str] = {} # mutable-ok: hand-off from the request listener
|
||||
trail: list[str] = [] # mutable-ok: navigation diagnostics for a failed dance
|
||||
|
||||
def _note_request(request: object) -> None:
|
||||
url = getattr(request, "url", "")
|
||||
host = httpx.URL(url).host
|
||||
if not allow_upstream_consent and (host == "linear.app" or host.endswith(".linear.app")):
|
||||
captured["upstream_consent"] = "seen"
|
||||
if url.startswith(OAUTH_CLIENT_REDIRECT_URI) and "url" not in captured:
|
||||
captured["url"] = url
|
||||
|
||||
async def _swallow_redirect(route: Route) -> None:
|
||||
await route.fulfill(status=200, content_type="text/plain", body="ok")
|
||||
|
||||
async with async_playwright() as playwright:
|
||||
browser = await playwright.chromium.launch(headless=True)
|
||||
context = await browser.new_context(storage_state=storage_state_path)
|
||||
await context.route(re.compile(re.escape(OAUTH_CLIENT_REDIRECT_URI) + r".*"), _swallow_redirect)
|
||||
page = await context.new_page()
|
||||
context.on("request", _note_request)
|
||||
page.on("framenavigated", lambda frame: trail.append(frame.url.split("?", 1)[0]))
|
||||
await page.goto(start_url, wait_until="domcontentloaded")
|
||||
deadline = time.monotonic() + BROWSER_CONSENT_TIMEOUT
|
||||
while "url" not in captured and time.monotonic() < deadline:
|
||||
try:
|
||||
await page.wait_for_load_state("networkidle", timeout=8000)
|
||||
except Exception: # noqa: BLE001 - a busy consent page never idles; fall through and try to advance it
|
||||
pass
|
||||
if "upstream_consent" in captured or "url" in captured:
|
||||
break
|
||||
if await page.locator("#username").count() and identity is not None:
|
||||
await page.locator("#username").fill(identity.username)
|
||||
await page.locator("#password").fill(identity.password)
|
||||
await page.locator("#kc-login").click()
|
||||
continue
|
||||
if "/ui/connect" in page.url and server_alias is not None:
|
||||
card = page.locator("div.cursor-pointer").filter(has=page.get_by_text(server_alias, exact=True))
|
||||
if await card.count() != 1:
|
||||
await asyncio.sleep(0.5)
|
||||
continue
|
||||
connect = card.get_by_text("Connect", exact=True)
|
||||
if await connect.count():
|
||||
await connect.click()
|
||||
continue
|
||||
if not await card.locator("svg.text-success").count():
|
||||
await asyncio.sleep(0.5)
|
||||
continue
|
||||
finish = page.get_by_role("button", name="Finish connecting", exact=True)
|
||||
if await finish.count() and await finish.is_enabled():
|
||||
await finish.click()
|
||||
continue
|
||||
control = page.locator(
|
||||
'button[name="action"][value="approve"], button:has-text("Authorize"), '
|
||||
'button:has-text("Allow"), button:has-text("@"), a:has-text("@")'
|
||||
).first
|
||||
try:
|
||||
await control.click(timeout=5000)
|
||||
except Exception: # noqa: BLE001 - nothing to advance yet; loop and re-check
|
||||
await asyncio.sleep(0.5)
|
||||
final_url = page.url
|
||||
await browser.close()
|
||||
|
||||
# A redirect chain can finish inside goto/networkidle before the loop checks the page.
|
||||
assert "upstream_consent" not in captured, "cold reconnect required upstream consent"
|
||||
landing = captured.get("url")
|
||||
assert landing is not None, (
|
||||
f"consent flow never reached {OAUTH_CLIENT_REDIRECT_URI}; "
|
||||
f"final={final_url.split('?', 1)[0]!r}; trail={trail[-6:]}"
|
||||
)
|
||||
params = dict(parse_qsl(httpx.URL(landing).query.decode()))
|
||||
assert "code" in params, "client redirect_uri carried no authorization code"
|
||||
return params["code"], params.get("state")
|
||||
|
||||
|
||||
def _oauth_provider(
|
||||
url: str,
|
||||
storage: InMemoryTokenStorage,
|
||||
storage_state_path: str | None,
|
||||
identity: Identity | None = None,
|
||||
server_alias: str | None = None,
|
||||
allow_upstream_consent: bool = True,
|
||||
) -> OAuthClientProvider:
|
||||
"""The SDK's real OAuth machinery (RFC 9728/8414 discovery, RFC 7591 DCR,
|
||||
PKCE, token exchange) with the browser leg driven by Playwright against the
|
||||
upstream's consent screen."""
|
||||
code_holder: dict[str, str | None] = {} # mutable-ok: hand-off between the two SDK callbacks
|
||||
|
||||
async def _reject_redirect(_: str) -> None:
|
||||
raise AssertionError("gateway demanded a fresh upstream consent; stored per-user token was not reused")
|
||||
|
||||
async def _follow_redirect(authorize_url: str) -> None:
|
||||
assert storage_state_path is not None
|
||||
code, state = await _browser_follow_authorize(
|
||||
authorize_url, storage_state_path, identity, server_alias, allow_upstream_consent
|
||||
)
|
||||
code_holder["code"] = code
|
||||
code_holder["state"] = state
|
||||
|
||||
redirect_handler: Final = _reject_redirect if storage_state_path is None else _follow_redirect
|
||||
|
||||
async def callback_handler() -> AuthorizationCodeResult:
|
||||
code = code_holder.get("code")
|
||||
assert code is not None, "callback_handler ran before the authorize redirect completed"
|
||||
return AuthorizationCodeResult(code=code, state=code_holder.get("state"))
|
||||
|
||||
return OAuthClientProvider(
|
||||
server_url=url,
|
||||
client_metadata=OAuthClientMetadata.model_validate(
|
||||
{
|
||||
"redirect_uris": [OAUTH_CLIENT_REDIRECT_URI],
|
||||
"token_endpoint_auth_method": "none",
|
||||
"grant_types": ["authorization_code", "refresh_token"],
|
||||
"response_types": ["code"],
|
||||
"client_name": "e2e-mcp-host",
|
||||
}
|
||||
),
|
||||
storage=storage,
|
||||
redirect_handler=redirect_handler,
|
||||
callback_handler=callback_handler,
|
||||
)
|
||||
|
||||
|
||||
class _HeaderInjectingTransport(httpx2.AsyncBaseTransport):
|
||||
"""Adds the caller's LiteLLM key header to every outgoing SDK request
|
||||
(discovery, DCR, token exchange), so the gateway resolves which user to
|
||||
store the upstream token for from the key on the token exchange, exactly
|
||||
like a production MCP host configured with a LiteLLM key header."""
|
||||
|
||||
def __init__(self, inner: httpx2.AsyncBaseTransport, headers: dict[str, str], gateway_url: str) -> None:
|
||||
self._inner = inner
|
||||
self._headers = headers
|
||||
self._gateway_url = httpx2.URL(gateway_url)
|
||||
|
||||
@staticmethod
|
||||
def _port(url: httpx2.URL) -> int | None:
|
||||
if url.port is not None:
|
||||
return url.port
|
||||
return {"http": 80, "https": 443}.get(url.scheme)
|
||||
|
||||
async def handle_async_request(self, request: httpx2.Request) -> httpx2.Response:
|
||||
same_origin: Final = (
|
||||
request.url.scheme == self._gateway_url.scheme
|
||||
and request.url.host == self._gateway_url.host
|
||||
and self._port(request.url) == self._port(self._gateway_url)
|
||||
)
|
||||
if same_origin:
|
||||
for name, value in self._headers.items():
|
||||
if name not in request.headers:
|
||||
request.headers[name] = value
|
||||
else:
|
||||
for name, value in self._headers.items():
|
||||
if request.headers.get(name) == value:
|
||||
del request.headers[name]
|
||||
return await self._inner.handle_async_request(request)
|
||||
|
||||
async def aclose(self) -> None:
|
||||
await self._inner.aclose()
|
||||
|
||||
|
||||
def _oauth_http_client(
|
||||
headers: dict[str, str], auth: OAuthClientProvider, gateway_url: str = PROXY_BASE_URL
|
||||
) -> httpx2.AsyncClient:
|
||||
return httpx2.AsyncClient(
|
||||
auth=auth,
|
||||
timeout=httpx2.Timeout(REQUEST_TIMEOUT),
|
||||
follow_redirects=True,
|
||||
transport=_HeaderInjectingTransport(httpx2.AsyncHTTPTransport(), headers, gateway_url),
|
||||
)
|
||||
|
||||
|
||||
async def _seed_via_dance(
|
||||
url: str, headers: dict[str, str], storage: InMemoryTokenStorage, storage_state_path: str
|
||||
) -> tuple[str, ...]:
|
||||
async with _oauth_http_client(headers, _oauth_provider(url, storage, storage_state_path)) as http_client:
|
||||
async with streamable_http_client(url, http_client=http_client) as (read, write):
|
||||
async with ClientSession(read, write) as session:
|
||||
await session.initialize()
|
||||
listed = await session.list_tools()
|
||||
return tuple(sorted(tool.name for tool in listed.tools))
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class OauthToolRun:
|
||||
tools: tuple[str, ...]
|
||||
is_error: bool
|
||||
text: str
|
||||
|
||||
|
||||
async def _list_and_call(
|
||||
url: str,
|
||||
headers: dict[str, str],
|
||||
storage: InMemoryTokenStorage,
|
||||
storage_state_path: str | None,
|
||||
tool: str,
|
||||
arguments: dict[str, str],
|
||||
gateway_url: str = PROXY_BASE_URL,
|
||||
identity: Identity | None = None,
|
||||
server_alias: str | None = None,
|
||||
allow_upstream_consent: bool = True,
|
||||
) -> OauthToolRun:
|
||||
async with _oauth_http_client(
|
||||
headers,
|
||||
_oauth_provider(url, storage, storage_state_path, identity, server_alias, allow_upstream_consent),
|
||||
gateway_url,
|
||||
) as http_client:
|
||||
async with streamable_http_client(url, http_client=http_client) as (read, write):
|
||||
async with ClientSession(read, write) as session:
|
||||
await session.initialize()
|
||||
listed: Final = await session.list_tools()
|
||||
result: Final = await session.call_tool(tool, arguments)
|
||||
text: Final = "".join(content.text for content in result.content if isinstance(content, TextContent))
|
||||
return OauthToolRun(
|
||||
tools=tuple(sorted(tool_item.name for tool_item in listed.tools)),
|
||||
is_error=result.is_error,
|
||||
text=text,
|
||||
)
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class ChatMcpClient:
|
||||
proxy: ProxyClient
|
||||
|
||||
@step("Register the MCP server with the alias {body.alias}")
|
||||
def create_server(self, body: McpServerCreateBody) -> McpServerInfo:
|
||||
return unwrap(
|
||||
self.proxy.transport.post(
|
||||
"/v1/mcp/server",
|
||||
headers=self.proxy.transport.master,
|
||||
json=body,
|
||||
response_type=McpServerInfo,
|
||||
)
|
||||
)
|
||||
|
||||
@step("Read the MCP server back from /v1/mcp/server")
|
||||
def server_info(self, server_id: str) -> McpServerInfo:
|
||||
return unwrap(
|
||||
self.proxy.transport.get(
|
||||
f"/v1/mcp/server/{server_id}",
|
||||
headers=self.proxy.transport.master,
|
||||
params=NoBody(),
|
||||
response_type=McpServerInfo,
|
||||
)
|
||||
)
|
||||
|
||||
@step("Delete the MCP server")
|
||||
def delete_server(self, server_id: str) -> None:
|
||||
_ = self.proxy.transport.delete(
|
||||
f"/v1/mcp/server/{server_id}",
|
||||
headers=self.proxy.transport.master,
|
||||
json=NoBody(),
|
||||
response_type=NoBody,
|
||||
)
|
||||
|
||||
@step("Sign the key's user in to the MCP server {alias} through the OAuth consent flow")
|
||||
def seed_user_token(self, alias: str, key: str, storage_state_path: str) -> tuple[str, ...]:
|
||||
"""Drive the interactive authorize dance for `key`'s user so the gateway
|
||||
stores their upstream token, retried to the shared deadline since the
|
||||
just-created server and key propagate asynchronously. The LiteLLM key
|
||||
rides x-litellm-api-key so the gateway binds the token to that user.
|
||||
Returns the upstream tool names the dance listed, proof the token works."""
|
||||
headers = {"x-litellm-api-key": f"Bearer {key}"}
|
||||
storage = InMemoryTokenStorage()
|
||||
deadline = time.monotonic() + self.proxy.poll_timeout
|
||||
last_error: Exception | None = None
|
||||
while time.monotonic() < deadline:
|
||||
try:
|
||||
return asyncio.run(_seed_via_dance(_mcp_url(alias), headers, storage, storage_state_path))
|
||||
except Exception as exc: # noqa: BLE001 - retried to the deadline; the last error surfaces below
|
||||
last_error = exc
|
||||
time.sleep(self.proxy.poll_interval)
|
||||
pytest.fail(
|
||||
f"authorize dance for {alias!r} never completed within {self.proxy.poll_timeout}s; "
|
||||
f"last error: {last_error!r}"
|
||||
)
|
||||
|
||||
@step("List the tools on the MCP server {alias} and call {tool} over the MCP protocol")
|
||||
def list_and_call(
|
||||
self,
|
||||
alias: str,
|
||||
headers: dict[str, str],
|
||||
storage: InMemoryTokenStorage,
|
||||
storage_state_path: str | None,
|
||||
tool: str,
|
||||
arguments: dict[str, str],
|
||||
base_url: str = PROXY_BASE_URL,
|
||||
identity: Identity | None = None,
|
||||
allow_upstream_consent: bool = True,
|
||||
) -> OauthToolRun:
|
||||
return asyncio.run(
|
||||
_list_and_call(
|
||||
f"{base_url.rstrip('/')}/mcp" if identity is not None else _mcp_url(alias, base_url),
|
||||
headers,
|
||||
storage,
|
||||
storage_state_path,
|
||||
tool,
|
||||
arguments,
|
||||
base_url,
|
||||
identity,
|
||||
alias,
|
||||
allow_upstream_consent,
|
||||
)
|
||||
)
|
||||
|
||||
@step("List the users with a stored OAuth token for the MCP server")
|
||||
def server_user_credentials(self, server_id: str) -> tuple[McpServerUserCredentialRow, ...]:
|
||||
return unwrap(
|
||||
self.proxy.transport.get(
|
||||
f"/v1/mcp/server/{server_id}/user-credentials",
|
||||
headers=self.proxy.transport.master,
|
||||
params=NoBody(),
|
||||
response_type=McpServerUserCredentialListResponse,
|
||||
)
|
||||
).root
|
||||
|
||||
@step("Revoke the user's stored OAuth token for the MCP server")
|
||||
def revoke_user_token(self, server_id: str, headers: AuthHeaders) -> None:
|
||||
_ = unwrap(
|
||||
self.proxy.transport.delete(
|
||||
f"/v1/mcp/server/{server_id}/oauth-user-credential",
|
||||
headers=headers,
|
||||
json=NoBody(),
|
||||
response_type=McpOauthUserCredentialStatus,
|
||||
)
|
||||
)
|
||||
|
||||
@step("Send a /chat/completions request to {body.model} with an MCP server attached as a tool")
|
||||
def chat_with_mcp(self, headers: AuthHeaders, body: ChatBody) -> ChatResponse:
|
||||
"""POST /chat/completions carrying the LiteLLM key in `headers` (either
|
||||
ingress form) with an MCP server attached in `body.tools`. The gateway
|
||||
resolves the user from the key and lists/executes the server's tools
|
||||
with that user's stored upstream token."""
|
||||
return unwrap(
|
||||
self.proxy.transport.post(
|
||||
"/chat/completions",
|
||||
headers=headers,
|
||||
json=body,
|
||||
response_type=ChatResponse,
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
def build_chat_client(proxy: ProxyClient) -> ChatMcpClient:
|
||||
return ChatMcpClient(proxy=proxy)
|
||||
|
|
@ -1,197 +0,0 @@
|
|||
"""An owned, source-built OAuth gateway with cold restarts and credential observations.
|
||||
|
||||
Only this child process is restarted. Its database and SSO client survive while
|
||||
its process-local caches do not; Redis is deliberately absent from its config.
|
||||
The optional live edge measures headers without recording credentials or bodies.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
import subprocess
|
||||
import sys
|
||||
import threading
|
||||
import time
|
||||
from collections.abc import Callable, Mapping
|
||||
from contextlib import ExitStack
|
||||
from dataclasses import dataclass, field
|
||||
from pathlib import Path
|
||||
from typing import Final
|
||||
|
||||
import psycopg
|
||||
from e2e_config import INHERITED_ENV_PREFIXES, available_port
|
||||
from e2e_http import NoBody
|
||||
from e2e_metadata import step
|
||||
from idp import Keycloak, stop_process_group
|
||||
from proxy_client import ProxyClient, build_proxy_client
|
||||
from psycopg.rows import class_row
|
||||
from pydantic import BaseModel, SecretStr, ValidationError
|
||||
|
||||
|
||||
class StoredOAuth(BaseModel):
|
||||
type: str
|
||||
access_token: SecretStr
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class CredentialRow:
|
||||
credential_b64: str = field(repr=False)
|
||||
|
||||
|
||||
@step("Read the user's stored OAuth credential for the MCP server from the database and decrypt it")
|
||||
def stored_oauth(user_id: str, server_id: str) -> StoredOAuth:
|
||||
"""Read the encrypted credential because management APIs omit the plaintext token."""
|
||||
from litellm.proxy.common_utils.encrypt_decrypt_utils import decrypt_value_helper
|
||||
|
||||
with psycopg.Connection[CredentialRow].connect(
|
||||
os.environ["DATABASE_URL"], row_factory=class_row(CredentialRow)
|
||||
) as conn:
|
||||
row: Final = conn.execute(
|
||||
'SELECT credential_b64 FROM "LiteLLM_MCPUserCredentials" WHERE user_id = %s AND server_id = %s',
|
||||
(user_id, server_id),
|
||||
).fetchone()
|
||||
assert row is not None, "canonical user/server has no persisted credential"
|
||||
plaintext: Final = decrypt_value_helper(
|
||||
row.credential_b64, "e2e_mcp_oauth", exception_type="debug", return_original_value=False
|
||||
)
|
||||
assert plaintext is not None, "persisted credential must decrypt with the gateway salt"
|
||||
assert plaintext != row.credential_b64, "persisted credential must be encrypted"
|
||||
try:
|
||||
credential: Final = StoredOAuth.model_validate_json(plaintext)
|
||||
except ValidationError:
|
||||
raise AssertionError("decrypted credential is not an OAuth payload") from None
|
||||
assert credential.type == "oauth2"
|
||||
assert bool(credential.access_token.get_secret_value()), "stored upstream token is empty"
|
||||
return credential
|
||||
|
||||
|
||||
class RpcMethod(BaseModel):
|
||||
method: str = ""
|
||||
|
||||
|
||||
@dataclass(slots=True)
|
||||
class OAuthObservation:
|
||||
gateway_token: str = field(default="", repr=False)
|
||||
_seen: tuple[tuple[str, str, bool], ...] = field(default=(), init=False, repr=False)
|
||||
_lock: threading.Lock = field(default_factory=threading.Lock, init=False, repr=False)
|
||||
|
||||
def observe(self, url: str, headers: Mapping[str, str], body: bytes | None) -> None:
|
||||
if body is None or not url.endswith("/mcp"):
|
||||
return
|
||||
try:
|
||||
operation: Final = RpcMethod.model_validate_json(body).method
|
||||
except ValidationError:
|
||||
return
|
||||
if operation not in ("tools/list", "tools/call"):
|
||||
return
|
||||
received: Final = headers.get("authorization", "")
|
||||
gateway_leaked: Final = any(
|
||||
value in (self.gateway_token, f"Bearer {self.gateway_token}") for value in headers.values()
|
||||
)
|
||||
with self._lock:
|
||||
self._seen = (*self._seen, (operation, received, gateway_leaked))
|
||||
|
||||
def assert_forwarded(self, expected: StoredOAuth) -> None:
|
||||
with self._lock:
|
||||
snapshot: Final = self._seen
|
||||
self._seen = ()
|
||||
assert {item[0] for item in snapshot} == {"tools/list", "tools/call"}, "missing upstream observations"
|
||||
expected_header: Final = f"Bearer {expected.access_token.get_secret_value()}"
|
||||
assert all(item[1] == expected_header for item in snapshot), "upstream bearer did not match the stored token"
|
||||
assert all(not item[2] for item in snapshot), "gateway bearer leaked to the upstream"
|
||||
|
||||
|
||||
@dataclass(slots=True)
|
||||
class OAuthGateway:
|
||||
base_url: str
|
||||
proxy: ProxyClient
|
||||
_environment: Mapping[str, str] = field(repr=False)
|
||||
_command: tuple[str, ...] = field(repr=False)
|
||||
_log_path: Path
|
||||
_child: subprocess.Popen[bytes] | None = field(default=None, init=False, repr=False)
|
||||
|
||||
@step("Start the separate LiteLLM proxy and wait for /health/liveliness")
|
||||
def start(self) -> None:
|
||||
with self._log_path.open("ab") as log:
|
||||
self._child = subprocess.Popen(
|
||||
self._command,
|
||||
env=self._environment,
|
||||
stdout=log,
|
||||
stderr=log,
|
||||
start_new_session=True,
|
||||
)
|
||||
deadline: Final = time.monotonic() + 120
|
||||
while time.monotonic() < deadline:
|
||||
assert self._child.poll() is None, "owned OAuth gateway exited; inspect its private log"
|
||||
result = self.proxy.transport.probe("/health/liveliness", params=NoBody())
|
||||
if result.status_code == 200:
|
||||
return
|
||||
time.sleep(0.5)
|
||||
raise AssertionError("owned OAuth gateway did not become ready")
|
||||
|
||||
@step("Stop the separate LiteLLM proxy")
|
||||
def stop(self) -> None:
|
||||
if self._child is not None:
|
||||
stop_process_group(self._child)
|
||||
assert self._child.poll() is not None, "old gateway process is still alive"
|
||||
|
||||
@step("Restart the separate LiteLLM proxy process so its in-memory caches start empty")
|
||||
def restart(self) -> None:
|
||||
assert self._child is not None
|
||||
previous: Final = self._child.pid
|
||||
self.stop()
|
||||
self.start()
|
||||
assert self._child.pid != previous, "gateway restart did not create a new process"
|
||||
|
||||
|
||||
@step("Start a separate LiteLLM proxy from source with JWT auth against Keycloak")
|
||||
def owned_gateway(idp: Keycloak, directory: Path, cleanup: ExitStack) -> OAuthGateway:
|
||||
for name in ("DATABASE_URL", "LITELLM_LICENSE", "LITELLM_SALT_KEY", "LITELLM_MASTER_KEY"):
|
||||
assert os.environ.get(name), f"{name} is required for the owned OAuth gateway"
|
||||
port: Final = available_port()
|
||||
base_url: Final = f"http://127.0.0.1:{port}"
|
||||
|
||||
def defer(callback: Callable[[], object]) -> None:
|
||||
cleanup.callback(callback)
|
||||
|
||||
browser: Final = idp.browser_client(callback_url=f"{base_url}/sso/callback", defer=defer)
|
||||
config: Final = directory / "oauth-gateway.yaml"
|
||||
config.write_text(
|
||||
"model_list: []\n"
|
||||
"general_settings:\n"
|
||||
" master_key: os.environ/LITELLM_MASTER_KEY\n"
|
||||
" database_url: os.environ/DATABASE_URL\n"
|
||||
" enable_jwt_auth: true\n"
|
||||
" litellm_jwtauth:\n"
|
||||
" user_id_jwt_field: sub\n"
|
||||
" user_email_jwt_field: email\n"
|
||||
" team_ids_jwt_field: groups\n"
|
||||
" user_id_upsert: true\n"
|
||||
)
|
||||
environment: Final = {
|
||||
**{key: value for key, value in os.environ.items() if not key.startswith(INHERITED_ENV_PREFIXES)},
|
||||
**browser.environment(idp.discovery()),
|
||||
"PROXY_BASE_URL": base_url,
|
||||
"JWT_PUBLIC_KEY_URL": idp.jwks_url,
|
||||
"JWT_ISSUER": idp.issuer,
|
||||
"JWT_AUDIENCE": "litellm-e2e",
|
||||
"DISABLE_SCHEMA_UPDATE": "true",
|
||||
"STORE_MODEL_IN_DB": "True",
|
||||
"PYTHONPATH": str(Path(__file__).resolve().parents[3]),
|
||||
}
|
||||
gateway: Final = OAuthGateway(
|
||||
base_url=base_url,
|
||||
proxy=build_proxy_client(
|
||||
base_url=base_url,
|
||||
control_plane_base_url=base_url,
|
||||
replica_urls=(base_url,),
|
||||
control_replica_urls=(base_url,),
|
||||
master_key=os.environ["LITELLM_MASTER_KEY"],
|
||||
),
|
||||
_environment=environment,
|
||||
_command=(sys.executable, "-m", "litellm.proxy.proxy_cli", "--config", str(config), "--port", str(port)),
|
||||
_log_path=directory / "oauth-gateway.log",
|
||||
)
|
||||
cleanup.callback(gateway.stop)
|
||||
gateway.start()
|
||||
return gateway
|
||||
|
|
@ -1,218 +0,0 @@
|
|||
"""On-demand e2e: a chat completion drives a gateway-managed OAuth MCP server.
|
||||
|
||||
The real end-user flow for MCP over an OAuth server: a user registers a Linear
|
||||
authorization_code server, authorizes it once so the gateway stores their
|
||||
upstream token, then sends a normal /chat/completions request with the Linear
|
||||
MCP attached. The gateway resolves the user from the LiteLLM key, lists Linear's
|
||||
tools with the stored per-user token, lets the model call one, executes it
|
||||
upstream with that token, and returns the answer. This is proven against the
|
||||
real Linear MCP server (mcp.linear.app) and a real Anthropic model, once per
|
||||
documented ingress header (x-litellm-api-key and Authorization).
|
||||
|
||||
The authorize dance is seeded through the mcp SDK's OAuthClientProvider; the one
|
||||
step Linear cannot auto-approve is the human consent, so it is captured once out
|
||||
of band (mcp/linear_session_capture.py) into a saved browser session and a
|
||||
headless Chromium clicks Approve every run. The test therefore skips unless
|
||||
E2E_LINEAR_STORAGE_STATE points at that session, so it never runs on the per-PR
|
||||
CI path; it is a nightly/on-demand real-server smoke test.
|
||||
|
||||
Fail-before-fix: without the stored per-user token the gateway lists no Linear
|
||||
tools, so mcp_list_tools comes back empty, nothing is called, and the
|
||||
assertions fail; a served, called, non-empty Linear tool proves the gateway
|
||||
pulled and used the user's token.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
|
||||
import pytest
|
||||
|
||||
from e2e_config import CHEAP_ANTHROPIC_MODEL, LINEAR_MCP_URL, LINEAR_STORAGE_STATE, unique_marker
|
||||
from e2e_http import AuthHeaders
|
||||
from e2e_metadata import Capability, Domain, Mode, Provider, Route, Subject, meta
|
||||
from lifecycle import ResourceManager
|
||||
from models import ChatBody, ChatMessage, KeyGenerateBody, McpChatTool, McpServerCreateBody, ObjectPermission
|
||||
from proxy_client import ProxyClient
|
||||
|
||||
pytest.importorskip("mcp", reason="mcp SDK not installed; run `uv sync --inexact --group e2e-dev`")
|
||||
pytest.importorskip(
|
||||
"playwright.async_api",
|
||||
reason="playwright not installed; run `uv pip install playwright` and `playwright install chromium`",
|
||||
)
|
||||
|
||||
from oauth_chat_client import ChatMcpClient, build_chat_client # noqa: E402 # imports follow the importorskip guards
|
||||
|
||||
pytestmark = [
|
||||
pytest.mark.e2e,
|
||||
pytest.mark.skipif(
|
||||
not LINEAR_STORAGE_STATE or not os.path.exists(LINEAR_STORAGE_STATE),
|
||||
reason="set E2E_LINEAR_STORAGE_STATE to a Linear session captured via mcp/linear_session_capture.py",
|
||||
),
|
||||
]
|
||||
|
||||
# Pinned from a live dance during verification (never guessed); the gateway
|
||||
# prefixes every upstream tool name with the server alias. list_teams is a
|
||||
# read-only Linear tool that takes no arguments and returns the caller's teams.
|
||||
LINEAR_READONLY_TOOL = "list_teams"
|
||||
LINEAR_PROMPT = "Use the list_teams tool to list my Linear teams, then reply with the name of one of them."
|
||||
|
||||
|
||||
@pytest.fixture(scope="session")
|
||||
def chat_client(proxy: ProxyClient) -> ChatMcpClient:
|
||||
return build_chat_client(proxy)
|
||||
|
||||
|
||||
class TestMcpChatCompletionOauth:
|
||||
"""A scoped internal-user key on a real Linear authorization_code server,
|
||||
used through /chat/completions once per ingress header: the gateway pulls
|
||||
the user's stored upstream token, lists and executes Linear's tools during
|
||||
the completion, and returns the answer."""
|
||||
|
||||
@pytest.mark.covers("mcp.list_tools.oauth.succeeds")
|
||||
@pytest.mark.covers("mcp.call_tool.oauth.succeeds")
|
||||
@meta(
|
||||
Subject(
|
||||
domain=Domain.MCP,
|
||||
route=Route.CHAT_COMPLETIONS,
|
||||
providers=(Provider.ANTHROPIC,),
|
||||
models=(CHEAP_ANTHROPIC_MODEL,),
|
||||
capabilities=(Capability.FUNCTION_CALLING,),
|
||||
mode=Mode.NONSTREAM,
|
||||
)
|
||||
)
|
||||
def test_chat_completion_uses_linear_with_x_litellm_api_key_header(
|
||||
self, chat_client: ChatMcpClient, resources: ResourceManager
|
||||
) -> None:
|
||||
marker = unique_marker()
|
||||
alias = f"e2elinear{marker}"
|
||||
created = chat_client.create_server(
|
||||
McpServerCreateBody(
|
||||
alias=alias,
|
||||
url=LINEAR_MCP_URL,
|
||||
allow_all_keys=False,
|
||||
auth_type="oauth2",
|
||||
oauth2_flow="authorization_code",
|
||||
)
|
||||
)
|
||||
resources.defer(lambda: chat_client.delete_server(created.server_id))
|
||||
|
||||
stored = chat_client.server_info(created.server_id)
|
||||
assert stored.auth_type == "oauth2"
|
||||
assert stored.oauth2_flow == "authorization_code"
|
||||
assert stored.allow_all_keys is False
|
||||
|
||||
key = chat_client.proxy.generate_key(
|
||||
KeyGenerateBody(
|
||||
user_id="e2e-test-user",
|
||||
object_permission=ObjectPermission(mcp_servers=[created.server_id]),
|
||||
)
|
||||
)
|
||||
resources.defer(lambda: chat_client.proxy.delete_key(key))
|
||||
|
||||
seeded = chat_client.seed_user_token(alias, key, LINEAR_STORAGE_STATE)
|
||||
assert f"{alias}-{LINEAR_READONLY_TOOL}" in seeded, (
|
||||
f"the authorize dance listed {seeded}, expected it to include {alias}-{LINEAR_READONLY_TOOL}"
|
||||
)
|
||||
|
||||
response = chat_client.chat_with_mcp(
|
||||
AuthHeaders.model_validate({"x-litellm-api-key": f"Bearer {key}"}),
|
||||
ChatBody(
|
||||
model=CHEAP_ANTHROPIC_MODEL,
|
||||
messages=[ChatMessage(role="user", content=LINEAR_PROMPT)],
|
||||
tools=[
|
||||
McpChatTool(
|
||||
server_url=f"litellm_proxy/mcp/{alias}",
|
||||
server_label=alias,
|
||||
require_approval="never",
|
||||
)
|
||||
],
|
||||
),
|
||||
)
|
||||
|
||||
message = response.choices[0].message
|
||||
assert message is not None and message.content, f"completion returned no answer: {response}"
|
||||
meta = message.provider_specific_fields
|
||||
assert meta is not None, f"no MCP metadata on the completion: {response}"
|
||||
listed = {t.function.name for t in (meta.mcp_list_tools or []) if t.function}
|
||||
assert f"{alias}-{LINEAR_READONLY_TOOL}" in listed, (
|
||||
f"the gateway listed {sorted(listed)}, expected the stored token to surface {alias}-{LINEAR_READONLY_TOOL}"
|
||||
)
|
||||
results = [r for r in (meta.mcp_call_results or []) if r.name == f"{alias}-{LINEAR_READONLY_TOOL}"]
|
||||
assert results and results[0].result, (
|
||||
f"Linear tool {alias}-{LINEAR_READONLY_TOOL} was not executed with a result: {meta.mcp_call_results}"
|
||||
)
|
||||
|
||||
@pytest.mark.covers("mcp.list_tools.oauth.succeeds")
|
||||
@pytest.mark.covers("mcp.call_tool.oauth.succeeds")
|
||||
@meta(
|
||||
Subject(
|
||||
domain=Domain.MCP,
|
||||
route=Route.CHAT_COMPLETIONS,
|
||||
providers=(Provider.ANTHROPIC,),
|
||||
models=(CHEAP_ANTHROPIC_MODEL,),
|
||||
capabilities=(Capability.FUNCTION_CALLING,),
|
||||
mode=Mode.NONSTREAM,
|
||||
)
|
||||
)
|
||||
def test_chat_completion_uses_linear_with_authorization_bearer_header(
|
||||
self, chat_client: ChatMcpClient, resources: ResourceManager
|
||||
) -> None:
|
||||
marker = unique_marker()
|
||||
alias = f"e2elinear{marker}"
|
||||
created = chat_client.create_server(
|
||||
McpServerCreateBody(
|
||||
alias=alias,
|
||||
url=LINEAR_MCP_URL,
|
||||
allow_all_keys=False,
|
||||
auth_type="oauth2",
|
||||
oauth2_flow="authorization_code",
|
||||
)
|
||||
)
|
||||
resources.defer(lambda: chat_client.delete_server(created.server_id))
|
||||
|
||||
stored = chat_client.server_info(created.server_id)
|
||||
assert stored.auth_type == "oauth2"
|
||||
assert stored.oauth2_flow == "authorization_code"
|
||||
assert stored.allow_all_keys is False
|
||||
|
||||
key = chat_client.proxy.generate_key(
|
||||
KeyGenerateBody(
|
||||
user_id="e2e-test-user",
|
||||
object_permission=ObjectPermission(mcp_servers=[created.server_id]),
|
||||
)
|
||||
)
|
||||
resources.defer(lambda: chat_client.proxy.delete_key(key))
|
||||
|
||||
seeded = chat_client.seed_user_token(alias, key, LINEAR_STORAGE_STATE)
|
||||
assert f"{alias}-{LINEAR_READONLY_TOOL}" in seeded, (
|
||||
f"the authorize dance listed {seeded}, expected it to include {alias}-{LINEAR_READONLY_TOOL}"
|
||||
)
|
||||
|
||||
response = chat_client.chat_with_mcp(
|
||||
AuthHeaders.model_validate({"authorization": f"Bearer {key}"}),
|
||||
ChatBody(
|
||||
model=CHEAP_ANTHROPIC_MODEL,
|
||||
messages=[ChatMessage(role="user", content=LINEAR_PROMPT)],
|
||||
tools=[
|
||||
McpChatTool(
|
||||
server_url=f"litellm_proxy/mcp/{alias}",
|
||||
server_label=alias,
|
||||
require_approval="never",
|
||||
)
|
||||
],
|
||||
),
|
||||
)
|
||||
|
||||
message = response.choices[0].message
|
||||
assert message is not None and message.content, f"completion returned no answer: {response}"
|
||||
meta = message.provider_specific_fields
|
||||
assert meta is not None, f"no MCP metadata on the completion: {response}"
|
||||
listed = {t.function.name for t in (meta.mcp_list_tools or []) if t.function}
|
||||
assert f"{alias}-{LINEAR_READONLY_TOOL}" in listed, (
|
||||
f"the gateway listed {sorted(listed)}, expected the stored token to surface {alias}-{LINEAR_READONLY_TOOL}"
|
||||
)
|
||||
results = [r for r in (meta.mcp_call_results or []) if r.name == f"{alias}-{LINEAR_READONLY_TOOL}"]
|
||||
assert results and results[0].result, (
|
||||
f"Linear tool {alias}-{LINEAR_READONLY_TOOL} was not executed with a result: {meta.mcp_call_results}"
|
||||
)
|
||||
|
|
@ -1,216 +0,0 @@
|
|||
"""Real OAuth consent, immediate MCP operations and cold-restart persistence.
|
||||
|
||||
Aggregate SSO uses the SDK's normal authentication. The per-server variant is
|
||||
explicitly a configured two-header client, not an Authorization-only OAuth host.
|
||||
The observed variants forward to the same real Linear upstream and compare its
|
||||
bearer at the forwarding boundary; direct variants retain unmodified discovery.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
from collections.abc import Iterator
|
||||
from contextlib import ExitStack
|
||||
from pathlib import Path
|
||||
from types import MappingProxyType
|
||||
from typing import Final, Literal
|
||||
|
||||
import pytest
|
||||
from e2e_config import LINEAR_MCP_URL, LINEAR_READONLY_TOOL, LINEAR_STORAGE_STATE, unique_marker
|
||||
from e2e_http import AuthHeaders, NoBody, get_external, unwrap
|
||||
from e2e_metadata import Domain, Route, Subject, meta
|
||||
from idp import Identity, Keycloak
|
||||
from lifecycle import ResourceManager
|
||||
from models import (
|
||||
McpOauthCredentials,
|
||||
McpServerCreateBody,
|
||||
ObjectPermission,
|
||||
TeamMemberAddBody,
|
||||
TeamMemberEntry,
|
||||
TeamUpdateBody,
|
||||
)
|
||||
from oauth_chat_client import ChatMcpClient, InMemoryTokenStorage, OauthToolRun, build_chat_client
|
||||
from oauth_gateway import OAuthGateway, OAuthObservation, owned_gateway, stored_oauth
|
||||
from provider_edge import LiveEdge, start_provider_edge
|
||||
from proxy_client import ProxyClient
|
||||
from pydantic import BaseModel, ValidationError
|
||||
|
||||
pytestmark = [pytest.mark.e2e, pytest.mark.mcp_oauth_live, pytest.mark.provider_live]
|
||||
|
||||
|
||||
class OAuthMetadata(BaseModel):
|
||||
issuer: str
|
||||
authorization_endpoint: str
|
||||
token_endpoint: str
|
||||
registration_endpoint: str
|
||||
|
||||
|
||||
class LinearTeam(BaseModel):
|
||||
id: str
|
||||
name: str
|
||||
|
||||
|
||||
class LinearTeams(BaseModel):
|
||||
teams: tuple[LinearTeam, ...]
|
||||
|
||||
|
||||
def assert_tool_result(run: OauthToolRun, tool: str) -> None:
|
||||
assert tool in run.tools
|
||||
assert run.is_error is False
|
||||
try:
|
||||
result: Final = LinearTeams.model_validate_json(run.text)
|
||||
except ValidationError:
|
||||
raise AssertionError("list_teams did not return the expected teams payload") from None
|
||||
assert result.teams, "the test workspace must contain at least one team"
|
||||
assert all(team.id and team.name for team in result.teams), "team results must contain identifiers and names"
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def oauth_gateway(idp: Keycloak, tmp_path_factory: pytest.TempPathFactory) -> Iterator[OAuthGateway]:
|
||||
assert LINEAR_STORAGE_STATE and Path(LINEAR_STORAGE_STATE).is_file(), (
|
||||
"E2E_LINEAR_STORAGE_STATE must name a captured Linear login; see mcp/linear_session_capture.py"
|
||||
)
|
||||
assert os.environ.get("E2E_FIXTURE_MODE", "live") == "live", "OAuth acceptance cannot use replay"
|
||||
with ExitStack() as cleanup:
|
||||
yield owned_gateway(idp, tmp_path_factory.mktemp("mcp-oauth"), cleanup)
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def proxy(oauth_gateway: OAuthGateway) -> ProxyClient:
|
||||
return oauth_gateway.proxy
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def client(proxy: ProxyClient) -> ChatMcpClient:
|
||||
return build_chat_client(proxy)
|
||||
|
||||
|
||||
class TestMcpOauthHappyPath:
|
||||
@pytest.mark.covers("mcp.list_tools.oauth.succeeds")
|
||||
@pytest.mark.covers("mcp.call_tool.oauth.succeeds")
|
||||
@pytest.mark.covers("mcp.call_tool.oauth.persists_across_processes")
|
||||
@meta(
|
||||
Subject(
|
||||
domain=Domain.MCP,
|
||||
route=Route.MCP,
|
||||
)
|
||||
)
|
||||
@pytest.mark.parametrize("route", ("aggregate_sso", "explicit_header_jwt"))
|
||||
@pytest.mark.parametrize("observed", (False, True), ids=("direct", "observed"))
|
||||
def test_consent_list_call_and_cold_restart(
|
||||
self,
|
||||
client: ChatMcpClient,
|
||||
resources: ResourceManager,
|
||||
jwt_identity: Identity,
|
||||
idp: Keycloak,
|
||||
oauth_gateway: OAuthGateway,
|
||||
route: Literal["aggregate_sso", "explicit_header_jwt"],
|
||||
observed: bool,
|
||||
) -> None:
|
||||
alias: Final = f"e2elinear{unique_marker()}"
|
||||
tool: Final = f"{alias}-{LINEAR_READONLY_TOOL}"
|
||||
token: Final = idp.access_token(jwt_identity)
|
||||
observation: Final = OAuthObservation(gateway_token=token)
|
||||
edge: Final = (
|
||||
start_provider_edge(
|
||||
LiveEdge(observe_request=observation.observe),
|
||||
mounts=MappingProxyType(
|
||||
{"linear": "https://mcp.linear.app", ".well-known": "https://mcp.linear.app/.well-known"}
|
||||
),
|
||||
)
|
||||
if observed
|
||||
else None
|
||||
)
|
||||
if edge is not None:
|
||||
resources.defer(edge.shutdown)
|
||||
metadata: Final = (
|
||||
unwrap(
|
||||
get_external(
|
||||
"https://mcp.linear.app/.well-known/oauth-authorization-server",
|
||||
response_type=OAuthMetadata,
|
||||
)
|
||||
)
|
||||
if observed
|
||||
else None
|
||||
)
|
||||
created: Final = client.create_server(
|
||||
McpServerCreateBody(
|
||||
alias=alias,
|
||||
server_name=alias,
|
||||
url=f"{edge.edge.api_base('linear')}/mcp" if edge is not None else LINEAR_MCP_URL,
|
||||
transport="http",
|
||||
allow_all_keys=False,
|
||||
auth_type="oauth2",
|
||||
oauth2_flow="authorization_code",
|
||||
per_server_oauth_discovery=route == "explicit_header_jwt",
|
||||
issuer=metadata.issuer if metadata else None,
|
||||
authorization_url=metadata.authorization_endpoint if metadata else None,
|
||||
token_url=metadata.token_endpoint if metadata else None,
|
||||
registration_url=metadata.registration_endpoint if metadata else None,
|
||||
credentials=McpOauthCredentials(upstream_resource=LINEAR_MCP_URL) if observed else None,
|
||||
)
|
||||
)
|
||||
resources.defer(lambda: client.delete_server(created.server_id))
|
||||
assert client.server_user_credentials(created.server_id) == (), (
|
||||
"scenario must start without upstream credentials"
|
||||
)
|
||||
unwrap(
|
||||
client.proxy.transport.post(
|
||||
"/team/member_add",
|
||||
headers=client.proxy.transport.master,
|
||||
json=TeamMemberAddBody(
|
||||
team_id=jwt_identity.group, member=TeamMemberEntry(user_id=jwt_identity.user_id, role="user")
|
||||
),
|
||||
response_type=NoBody,
|
||||
)
|
||||
)
|
||||
client.proxy.update_team(
|
||||
TeamUpdateBody(
|
||||
team_id=jwt_identity.group,
|
||||
object_permission=ObjectPermission(mcp_servers=[created.server_id]),
|
||||
)
|
||||
)
|
||||
headers: Final = {"x-litellm-api-key": f"Bearer {token}"} if route == "explicit_header_jwt" else {}
|
||||
resources.defer(
|
||||
lambda: client.revoke_user_token(
|
||||
created.server_id,
|
||||
AuthHeaders(authorization=f"Bearer {idp.access_token(jwt_identity)}"),
|
||||
)
|
||||
)
|
||||
identity: Final = jwt_identity if route == "aggregate_sso" else None
|
||||
first: Final = client.list_and_call(
|
||||
alias,
|
||||
headers,
|
||||
InMemoryTokenStorage(),
|
||||
LINEAR_STORAGE_STATE,
|
||||
tool,
|
||||
{},
|
||||
base_url=oauth_gateway.base_url,
|
||||
identity=identity,
|
||||
)
|
||||
assert_tool_result(first, tool)
|
||||
credentials: Final = client.server_user_credentials(created.server_id)
|
||||
assert len(credentials) == 1
|
||||
assert credentials[0].user_id == jwt_identity.user_id
|
||||
assert credentials[0].credential_type == "oauth2"
|
||||
first_stored_oauth: Final = stored_oauth(jwt_identity.user_id, created.server_id)
|
||||
if observed:
|
||||
observation.assert_forwarded(first_stored_oauth)
|
||||
oauth_gateway.restart()
|
||||
fresh_token: Final = idp.access_token(jwt_identity)
|
||||
observation.gateway_token = fresh_token
|
||||
second: Final = client.list_and_call(
|
||||
alias,
|
||||
{"Authorization": f"Bearer {fresh_token}"} if identity is None else {},
|
||||
InMemoryTokenStorage(),
|
||||
LINEAR_STORAGE_STATE if identity is not None else None,
|
||||
tool,
|
||||
{},
|
||||
base_url=oauth_gateway.base_url,
|
||||
identity=identity,
|
||||
allow_upstream_consent=False,
|
||||
)
|
||||
assert_tool_result(second, tool)
|
||||
second_stored_oauth: Final = stored_oauth(jwt_identity.user_id, created.server_id)
|
||||
if observed:
|
||||
observation.assert_forwarded(second_stored_oauth)
|
||||
|
|
@ -5,7 +5,8 @@ embedding deployment, and turns on `litellm_settings.mcp_semantic_tool_filter` t
|
|||
/config/update. A chat completion that attaches the server and asks to search logs must report,
|
||||
in the x-litellm-semantic-filter headers, that the catalog was cut to at most top_k tools, all
|
||||
of them tools the key lists. Requires DD_API_KEY + DD_APP_KEY and the proxy's
|
||||
OPENAI_API_KEY for the embedding deployment.
|
||||
OPENAI_API_KEY for the embedding deployment. The filter is proxy-wide, so the test is
|
||||
marked `proxy_wide_config` and runs on a stack the main suite is not driving.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
|
@ -23,7 +24,7 @@ from mcp_client import McpClient
|
|||
from models import ChatBody, ChatMessage, LiteLLMParamsBody, McpChatTool
|
||||
from pydantic import BaseModel
|
||||
|
||||
pytestmark = [pytest.mark.e2e, pytest.mark.quiet_stack]
|
||||
pytestmark = [pytest.mark.e2e, pytest.mark.proxy_wide_config]
|
||||
|
||||
TOP_K: Final = 3
|
||||
|
||||
|
|
|
|||
|
|
@ -730,28 +730,14 @@ class McpInfo(BaseModel):
|
|||
logo_url: str | None = None
|
||||
|
||||
|
||||
class McpOauthCredentials(BaseModel):
|
||||
upstream_resource: str
|
||||
|
||||
|
||||
class McpServerCreateBody(BaseModel):
|
||||
"""POST /v1/mcp/server. For a gateway-managed OAuth server, `auth_type` is
|
||||
`oauth2` and `oauth2_flow` is `authorization_code`; the upstream endpoints
|
||||
are discovered and registered via DCR when left unset. `allow_all_keys`
|
||||
false scopes the server to keys granted it through object_permission."""
|
||||
"""POST /v1/mcp/server. `allow_all_keys` false scopes the server to keys
|
||||
granted it through object_permission."""
|
||||
|
||||
alias: str
|
||||
url: str
|
||||
transport: str = "http"
|
||||
allow_all_keys: bool = True
|
||||
auth_type: str | None = None
|
||||
oauth2_flow: Literal["client_credentials", "authorization_code"] | None = None
|
||||
per_server_oauth_discovery: bool | None = None
|
||||
issuer: str | None = None
|
||||
authorization_url: str | None = None
|
||||
token_url: str | None = None
|
||||
registration_url: str | None = None
|
||||
credentials: McpOauthCredentials | None = Field(default=None, repr=False)
|
||||
server_name: str | None = None
|
||||
description: str | None = None
|
||||
mcp_info: McpInfo | None = None
|
||||
|
|
@ -791,26 +777,6 @@ class McpServerListResponse(RootModel[list[McpServerRow]]):
|
|||
"""GET /v1/mcp/server answers with a bare array of servers."""
|
||||
|
||||
|
||||
class McpServerUserCredentialRow(BaseModel):
|
||||
user_id: str
|
||||
credential_type: Literal["oauth2", "byok"]
|
||||
expires_at: str | None = None
|
||||
connected_at: str | None = None
|
||||
updated_at: str
|
||||
|
||||
|
||||
class McpServerUserCredentialListResponse(RootModel[tuple[McpServerUserCredentialRow, ...]]):
|
||||
"""GET /v1/mcp/server/{server_id}/user-credentials answers with a bare array."""
|
||||
|
||||
|
||||
class McpOauthUserCredentialStatus(BaseModel):
|
||||
server_id: str
|
||||
has_credential: bool
|
||||
expires_at: str | None = None
|
||||
is_expired: bool = False
|
||||
connected_at: str | None = None
|
||||
|
||||
|
||||
class ToolsetTool(BaseModel):
|
||||
server_id: str
|
||||
tool_name: str
|
||||
|
|
@ -1295,27 +1261,6 @@ class CostMap(RootModel[dict[str, CostMapEntry]]):
|
|||
pass
|
||||
|
||||
|
||||
class ConfigPatchResponse(BaseModel):
|
||||
status: str
|
||||
values: dict[str, float | dict[str, float]]
|
||||
|
||||
|
||||
class CostDiscountConfig(RootModel[dict[str, float]]):
|
||||
pass
|
||||
|
||||
|
||||
class CostMarginConfig(RootModel[dict[str, float | dict[str, float]]]):
|
||||
pass
|
||||
|
||||
|
||||
class CostDiscountConfigResponse(BaseModel):
|
||||
values: dict[str, float]
|
||||
|
||||
|
||||
class CostMarginConfigResponse(BaseModel):
|
||||
values: dict[str, float | dict[str, float]]
|
||||
|
||||
|
||||
class BedrockGuardrailParams(BaseModel):
|
||||
guardrail: Literal["bedrock"] = "bedrock"
|
||||
mode: Literal["pre_call"] = "pre_call"
|
||||
|
|
|
|||
|
|
@ -32,7 +32,7 @@ from setup.podcpu import gateway_cpu_sampler
|
|||
from setup.provision import MOCK_PROVIDERS, build_setup, guardrail_url, mock_url
|
||||
from setup.run import run_plan
|
||||
|
||||
pytestmark = [pytest.mark.perf, pytest.mark.quiet_stack]
|
||||
pytestmark = pytest.mark.perf
|
||||
_SPEND_SAMPLE: Final = 40
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -1,17 +1,19 @@
|
|||
"""Live e2e: a few hundred requests that fail before any provider answers must not
|
||||
grow the proxy's resident memory past a loose ceiling once the proxy is warm.
|
||||
grow the proxy's resident memory past a loose ceiling once the proxy is warm. It reads
|
||||
the proxy's own RSS, so it runs in the perf lane, where nothing else is hitting the stack.
|
||||
|
||||
The regression this guards shipped in v1.100.0: every retry breadcrumb copied the
|
||||
whole request and the copies nested into one router-global list, so a proxy under
|
||||
retry-heavy failing traffic grew until it was OOM-killed. The traffic here has that
|
||||
shape: a model group whose deployments refuse at the socket (an unreachable base
|
||||
URL) with cooldown_time 0 so the router keeps retrying them, per-request retries,
|
||||
and a fallback group that refuses the same way, each request carrying a long chat
|
||||
URL, with a literal key since the perf stack carries no provider keys) with
|
||||
cooldown_time 0 so the router keeps retrying them, per-request retries, and a
|
||||
fallback group that refuses the same way, each request carrying a long chat
|
||||
transcript so every whole-request copy costs hundreds of containers instead of a
|
||||
handful. Under the stack's cooldown policy a
|
||||
deployment that fails a handful of times in a row is benched (a bad-credential 401
|
||||
included), the router answers "No deployments available" without retrying, and the
|
||||
retry loop that leaks stops running; cooldown_time 0 keeps it running.
|
||||
handful. Under the stack's cooldown policy a deployment that fails a handful of
|
||||
times in a row is benched (a bad-credential 401 included), the router answers "No
|
||||
deployments available" without retrying, and the retry loop that leaks stops running;
|
||||
cooldown_time 0 keeps it running.
|
||||
|
||||
Four identical batches run back to back. The first two are the warmup that grows
|
||||
the proxy's caches and allocator arenas, the last two are the measured window, and
|
||||
|
|
@ -35,19 +37,11 @@ single load-balanced address reaches the workers behind it one answer at a time,
|
|||
a worker that answered only one checkpoint would otherwise drop out of the
|
||||
comparison, which is where a leaking worker could hide.
|
||||
|
||||
RSS alone is a coarse gauge: on the release stack (spend logs storing prompts,
|
||||
json logs, prometheus and otel callbacks) the same v1.100.0 breadcrumbs grew RSS
|
||||
by only about 15 MB per 300 failing requests, while every failing request's stored
|
||||
request snapshot carried a copy of the request per failed attempt, over 100 KB on
|
||||
the first call and a couple of MB once the copies nested, against tens of KB with
|
||||
the fix. So the first check sends one failing request before the phases, reads its
|
||||
spend log back through /spend/logs, and holds the stored request body to a fixed
|
||||
size budget: the deterministic catch for a breadcrumb that copies the whole
|
||||
request. It runs before the phases because the leaking writer drops its own rows
|
||||
under the phases' traffic (a queue budget hit, a recursion limit on the nested
|
||||
copies), which would turn the size check into a missing-row check.
|
||||
The stored-request half of that regression (every failing request's spend log carrying a
|
||||
copy of the request per failed attempt) is a deterministic size check, so it lives in
|
||||
tests/integration/spend/test_failing_request_stored_snapshot.py.
|
||||
|
||||
A second, cheaper check holds the idle footprint: every worker's RSS as the harness
|
||||
A cheaper check holds the idle footprint: every worker's RSS as the harness
|
||||
read it at collection time, before this pytest process sent any traffic (see
|
||||
conftest.pytest_collection_finish), must sit under a fixed budget. On the
|
||||
release gate that is a fresh stack right after its readiness gate, one worker per
|
||||
|
|
@ -74,7 +68,6 @@ from typing import Final
|
|||
|
||||
import pytest
|
||||
|
||||
from complexity_router_client import ComplexityRouterClient
|
||||
from e2e_config import (
|
||||
MEMORY_CONCURRENCY,
|
||||
MEMORY_IDLE_RSS_BUDGET_MB,
|
||||
|
|
@ -83,7 +76,6 @@ from e2e_config import (
|
|||
MEMORY_RSS_CEILING_MB,
|
||||
MEMORY_RSS_SAMPLE_INTERVAL_SECONDS,
|
||||
MEMORY_RSS_SETTLE_SAMPLES,
|
||||
MEMORY_STORED_REQUEST_BUDGET_KB,
|
||||
MEMORY_TRANSCRIPT_TURNS,
|
||||
PROXY_REPLICA_URLS,
|
||||
unique_marker,
|
||||
|
|
@ -91,11 +83,12 @@ from e2e_config import (
|
|||
from e2e_metadata import Domain, Mode, Provider, Subject, meta
|
||||
from lifecycle import ResourceManager
|
||||
from memory_readings import RssCapture, RssReading, WorkerKey, read_rss_everywhere
|
||||
from models import ChatMessage, RouterSettingsOverride, SpendLogRow
|
||||
from models import ChatMessage, LiteLLMParamsBody, ReliabilityChatBody, RouterSettingsOverride
|
||||
from proxy_client import ProxyClient
|
||||
from reliability_support import REAL_MODEL, chat_override, create_never_benched_refusing_deployment
|
||||
|
||||
pytestmark = [pytest.mark.e2e, pytest.mark.quiet_stack]
|
||||
pytestmark = [pytest.mark.e2e, pytest.mark.perf]
|
||||
|
||||
REFUSING_MODEL: Final = "openai/gpt-5.5"
|
||||
|
||||
DEPLOYMENTS_PER_GROUP: Final = 2
|
||||
RSS_SAMPLE_CAP: Final = 4 * MEMORY_RSS_SETTLE_SAMPLES
|
||||
|
|
@ -110,7 +103,6 @@ class FailedCall:
|
|||
status_code: int
|
||||
seconds: float
|
||||
body_head: str
|
||||
call_id: str | None
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
|
|
@ -123,8 +115,17 @@ class WorkerGrowth:
|
|||
return self.after.ram_usage_mb - self.warm.ram_usage_mb
|
||||
|
||||
|
||||
def _create_never_benched_refusing_deployment(proxy: ProxyClient, name: str) -> str:
|
||||
return proxy.create_model(
|
||||
name,
|
||||
LiteLLMParamsBody(
|
||||
model=REFUSING_MODEL, api_key="sk-unreachable", api_base="http://127.0.0.1:9/v1", cooldown_time=0
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
def _register_refusing_group(proxy: ProxyClient, resources: ResourceManager, name: str) -> None:
|
||||
for model_id in tuple(create_never_benched_refusing_deployment(proxy, name) for _ in range(DEPLOYMENTS_PER_GROUP)):
|
||||
for model_id in tuple(_create_never_benched_refusing_deployment(proxy, name) for _ in range(DEPLOYMENTS_PER_GROUP)):
|
||||
resources.defer(lambda model_id=model_id: proxy.delete_model(model_id))
|
||||
|
||||
|
||||
|
|
@ -141,10 +142,18 @@ TRANSCRIPT: Final = _transcript(MEMORY_TRANSCRIPT_TURNS)
|
|||
|
||||
def _fail_once(proxy: ProxyClient, key: str, model: str, override: RouterSettingsOverride) -> FailedCall:
|
||||
started: Final = time.perf_counter()
|
||||
resp: Final = chat_override(
|
||||
proxy, key, model, f"memory regression {unique_marker()}", override=override, history=TRANSCRIPT
|
||||
resp: Final = proxy.transport.send(
|
||||
"/chat/completions",
|
||||
headers=proxy.transport.bearer(key),
|
||||
json=ReliabilityChatBody(
|
||||
model=model,
|
||||
messages=[*TRANSCRIPT, ChatMessage(role="user", content=f"memory regression {unique_marker()}")],
|
||||
max_tokens=512,
|
||||
router_settings_override=override,
|
||||
cache={"no-cache": True},
|
||||
),
|
||||
)
|
||||
return FailedCall(resp.status_code, time.perf_counter() - started, resp.body[:300], resp.call_id)
|
||||
return FailedCall(resp.status_code, time.perf_counter() - started, resp.body[:300])
|
||||
|
||||
|
||||
def _fail_many(proxy: ProxyClient, key: str, model: str, override: RouterSettingsOverride) -> tuple[FailedCall, ...]:
|
||||
|
|
@ -259,24 +268,6 @@ def _assert_every_call_failed_through_fallback(calls: Sequence[FailedCall], fall
|
|||
)
|
||||
|
||||
|
||||
def _stored_request_kb(proxy: ProxyClient, call: FailedCall) -> float:
|
||||
assert call.call_id, (
|
||||
f"the failing call carried no x-litellm-call-id header, so its spend log cannot be read back: {call.body_head}"
|
||||
)
|
||||
rows: Final[Sequence[SpendLogRow]] = proxy.poll_logs_for_request_id(call.call_id)
|
||||
assert rows, (
|
||||
f"no spend log row appeared for failing call {call.call_id} within the poll window: either the stack "
|
||||
"writes no spend logs or its writer dropped the row, which the v1.100.0 one did once the stored "
|
||||
"request outgrew the writer's queue budget"
|
||||
)
|
||||
snapshot: Final = rows[0].proxy_server_request
|
||||
assert snapshot, (
|
||||
f"spend log {call.call_id} stored no request body, so the stack is not running with "
|
||||
"general_settings.store_prompts_in_spend_logs and the stored-request check would pass vacuously"
|
||||
)
|
||||
return len(json.dumps(snapshot).encode()) / 1024
|
||||
|
||||
|
||||
class TestReliabilityMemory:
|
||||
@pytest.mark.covers("reliability.perf.idle_memory.under_slo")
|
||||
@meta(
|
||||
|
|
@ -309,13 +300,13 @@ class TestReliabilityMemory:
|
|||
Subject(
|
||||
domain=Domain.ROUTING,
|
||||
providers=(Provider.OPENAI,),
|
||||
models=(REAL_MODEL,),
|
||||
models=(REFUSING_MODEL,),
|
||||
mode=Mode.NONSTREAM,
|
||||
)
|
||||
)
|
||||
def test_failing_requests_do_not_grow_rss_or_stored_request(
|
||||
def test_failing_requests_do_not_grow_rss(
|
||||
self,
|
||||
client: ComplexityRouterClient,
|
||||
proxy: ProxyClient,
|
||||
resources: ResourceManager,
|
||||
scoped_key: str,
|
||||
record_property: Callable[[str, object], None],
|
||||
|
|
@ -323,35 +314,25 @@ class TestReliabilityMemory:
|
|||
marker: Final = unique_marker()
|
||||
primary: Final = f"reliability-memory-{marker}"
|
||||
fallback: Final = f"reliability-memory-fb-{marker}"
|
||||
_register_refusing_group(client.proxy, resources, primary)
|
||||
_register_refusing_group(client.proxy, resources, fallback)
|
||||
_register_refusing_group(proxy, resources, primary)
|
||||
_register_refusing_group(proxy, resources, fallback)
|
||||
override: Final = RouterSettingsOverride(
|
||||
num_retries=MEMORY_RETRIES_PER_REQUEST, fallbacks=[{primary: [fallback]}]
|
||||
)
|
||||
|
||||
probe: Final = _fail_once(client.proxy, scoped_key, primary, override)
|
||||
_assert_every_call_failed_through_fallback((probe,), fallback)
|
||||
stored_kb: Final = _stored_request_kb(client.proxy, probe)
|
||||
assert stored_kb <= MEMORY_STORED_REQUEST_BUDGET_KB, (
|
||||
f"the spend log of one failing request stored a {stored_kb:.0f} KB request body, past the "
|
||||
f"{MEMORY_STORED_REQUEST_BUDGET_KB:.0f} KB budget for a {len(TRANSCRIPT)}-message transcript with "
|
||||
f"{MEMORY_RETRIES_PER_REQUEST} retries and a fallback; the retry breadcrumbs are copying the whole "
|
||||
f"request into the stored snapshot the way the v1.100.0 ones did"
|
||||
)
|
||||
|
||||
warm: Final = _settled_rss_after_batches(
|
||||
client.proxy,
|
||||
proxy,
|
||||
scoped_key,
|
||||
primary,
|
||||
override,
|
||||
fallback,
|
||||
record_property,
|
||||
_settled_rss_per_worker(client.proxy),
|
||||
_settled_rss_per_worker(proxy),
|
||||
1,
|
||||
WARMUP_BATCHES,
|
||||
)
|
||||
after: Final = _settled_rss_after_batches(
|
||||
client.proxy,
|
||||
proxy,
|
||||
scoped_key,
|
||||
primary,
|
||||
override,
|
||||
|
|
@ -11,7 +11,7 @@ from __future__ import annotations
|
|||
import os
|
||||
import time
|
||||
import warnings
|
||||
from collections.abc import Callable, Mapping
|
||||
from collections.abc import Callable, Mapping, Sequence
|
||||
from dataclasses import dataclass, field, replace
|
||||
from datetime import datetime
|
||||
from functools import reduce
|
||||
|
|
@ -155,6 +155,21 @@ class NotServableOn:
|
|||
last_result: Result[ModelsListResponse] | None
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class ModelRegistered:
|
||||
model_name: str
|
||||
model_id: str
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class ModelNotRegistered:
|
||||
model_name: str
|
||||
reason: str
|
||||
|
||||
|
||||
ModelRegistration = ModelRegistered | ModelNotRegistered
|
||||
|
||||
|
||||
def await_servable(
|
||||
list_models: Callable[[float], Result[ModelsListResponse]],
|
||||
*,
|
||||
|
|
@ -461,6 +476,7 @@ class ProxyClient:
|
|||
model_servable_db_sync_seconds: float = MODEL_SERVABLE_DB_SYNC_SECONDS
|
||||
model_servable_interval: float = MODEL_SERVABLE_INTERVAL
|
||||
model_servable_request_timeout: float = MODEL_SERVABLE_REQUEST_TIMEOUT
|
||||
settle: Callable[[float], None] = settle_propagation
|
||||
|
||||
def with_caller(self, caller: Caller) -> ProxyClient:
|
||||
return replace(self, caller=caller)
|
||||
|
|
@ -694,31 +710,59 @@ class ProxyClient:
|
|||
balancer address is configured (every request opens a fresh connection, so
|
||||
the caller's next request re-rolls), so waiting out PROPAGATION_TIMEOUT is
|
||||
what makes the model safe to use anywhere."""
|
||||
model_id = unwrap(
|
||||
self.transport.post(
|
||||
"/model/new",
|
||||
headers=self.management_headers(),
|
||||
json=body.model_copy(update={"litellm_params": route_cache_model(
|
||||
body.litellm_params, provider_edge_base,
|
||||
enabled=os.environ.get("E2E_PROVIDER_CACHE", "0") == "1" and not provider_live,
|
||||
mode=body.model_info.mode,
|
||||
)}),
|
||||
response_type=ModelNewResponse,
|
||||
)
|
||||
).model_id
|
||||
model_id = unwrap(self._post_model(body, provider_live=provider_live)).model_id
|
||||
written_at = time.monotonic()
|
||||
try:
|
||||
self._await_model_servable(body.model_name, listed_for)
|
||||
except BaseException:
|
||||
self.delete_model(model_id)
|
||||
raise
|
||||
settle_propagation(written_at)
|
||||
self.settle(written_at)
|
||||
return model_id
|
||||
|
||||
@step("Add a batch of deployments, then wait once for all of them to propagate")
|
||||
def register_models(
|
||||
self, bodies: Sequence[ModelNewBody], *, provider_live: bool = False
|
||||
) -> tuple[ModelRegistration, ...]:
|
||||
posted: Final = tuple((body.model_name, self._post_model(body, provider_live=provider_live)) for body in bodies)
|
||||
written_at: Final = time.monotonic()
|
||||
registrations: Final = tuple(self._servable_registration(name, result) for name, result in posted)
|
||||
self.settle(written_at)
|
||||
return registrations
|
||||
|
||||
def _post_model(self, body: ModelNewBody, *, provider_live: bool) -> Result[ModelNewResponse]:
|
||||
return self.transport.post(
|
||||
"/model/new",
|
||||
headers=self.management_headers(),
|
||||
json=body.model_copy(update={"litellm_params": route_cache_model(
|
||||
body.litellm_params, provider_edge_base,
|
||||
enabled=os.environ.get("E2E_PROVIDER_CACHE", "0") == "1" and not provider_live,
|
||||
mode=body.model_info.mode,
|
||||
)}),
|
||||
response_type=ModelNewResponse,
|
||||
)
|
||||
|
||||
def _servable_registration(self, model_name: str, posted: Result[ModelNewResponse]) -> ModelRegistration:
|
||||
if not isinstance(posted, Success):
|
||||
return ModelNotRegistered(model_name=model_name, reason=f"/model/new failed: {posted!r}")
|
||||
match self._servable_outcome(model_name, None):
|
||||
case Servable():
|
||||
return ModelRegistered(model_name=model_name, model_id=posted.data.model_id)
|
||||
case NotServableOn() as missing:
|
||||
self.delete_model(posted.data.model_id)
|
||||
return ModelNotRegistered(model_name=model_name, reason=self._not_servable_message(model_name, missing))
|
||||
|
||||
def _await_model_servable(self, model_name: str, listed_for: str | None = None) -> None:
|
||||
"""Block until every replica lists `model_name`, or fail at model_servable_timeout."""
|
||||
match self._servable_outcome(model_name, listed_for):
|
||||
case Servable():
|
||||
return
|
||||
case NotServableOn() as missing:
|
||||
raise AssertionError(self._not_servable_message(model_name, missing))
|
||||
|
||||
def _servable_outcome(self, model_name: str, listed_for: str | None) -> Servable | NotServableOn:
|
||||
headers: Final = self.management_headers(listed_for)
|
||||
outcome: Final = await_servable_everywhere(
|
||||
return await_servable_everywhere(
|
||||
{url: self._models_poller(transport, headers) for url, transport in self.replicas.items()},
|
||||
model_name=model_name,
|
||||
timeout=self.model_servable_timeout,
|
||||
|
|
@ -728,19 +772,15 @@ class ProxyClient:
|
|||
now=time.monotonic,
|
||||
sleep=time.sleep,
|
||||
)
|
||||
match outcome:
|
||||
case Servable():
|
||||
return
|
||||
case NotServableOn(replica=replica, last_result=last_result):
|
||||
raise AssertionError(
|
||||
servable_timeout_message(
|
||||
model_name=model_name,
|
||||
replica=replica,
|
||||
timeout=self.model_servable_timeout,
|
||||
db_sync_seconds=self.model_servable_db_sync_seconds,
|
||||
last_result=last_result,
|
||||
)
|
||||
)
|
||||
|
||||
def _not_servable_message(self, model_name: str, missing: NotServableOn) -> str:
|
||||
return servable_timeout_message(
|
||||
model_name=model_name,
|
||||
replica=missing.replica,
|
||||
timeout=self.model_servable_timeout,
|
||||
db_sync_seconds=self.model_servable_db_sync_seconds,
|
||||
last_result=missing.last_result,
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _models_poller(transport: Transport, headers: AuthHeaders) -> ModelsPoller:
|
||||
|
|
|
|||
|
|
@ -13,8 +13,6 @@ markers =
|
|||
prompt_caching_stack: needs a proxy running with router_settings.optional_pre_call_checks including prompt_caching; deselected unless E2E_PROMPT_CACHING_STACK is set
|
||||
cli_determinism: drives the real claude CLI for several seconds; deselected unless E2E_CLI_DETERMINISM is set
|
||||
redis_chaos: load test that pauses the proxy's Redis outright mid-run; needs a proxy booted from gateway/redis_chaos_ci_config.yml on the same host, and is deselected unless E2E_PERF_TESTS is set
|
||||
quiet_stack: measures the proxy itself, so it runs while no other test on this host is hitting the stack; every other test waits for it to finish
|
||||
mcp_oauth_live: real Linear OAuth consent via a captured browser session; deselected unless E2E_MCP_OAUTH_LIVE is set
|
||||
provider_edge_host: routes provider traffic through the pytest host's edge in every fixture mode, so the gateway must reach the pytest host; deselected unless E2E_PROVIDER_EDGE_HOST_REACHABLE is set
|
||||
owned_gateway: boots its own proxy from source against the stack's Postgres, so it needs DATABASE_URL on the pytest host; deselected unless E2E_OWNED_GATEWAY is set
|
||||
otel_v2: needs a proxy running with LITELLM_OTEL_V2=true; deselected unless E2E_OTEL_V2 is set
|
||||
|
|
@ -22,3 +20,4 @@ markers =
|
|||
secret_manager: needs a proxy booted from gateway/secret_manager_<system>_ci_config.yml against that live secret manager; deselected unless E2E_SECRET_MANAGER names the backend (see secret_manager/secret_backends.py)
|
||||
sagemaker_nova: needs a deployed SageMaker Nova endpoint the proxy's AWS credentials can invoke; deselected unless SAGEMAKER_NOVA_ENDPOINT names it
|
||||
sagemaker_nova2_lite: needs a deployed SageMaker Nova 2 Lite endpoint the proxy's AWS credentials can invoke; deselected unless SAGEMAKER_NOVA2_LITE_ENDPOINT names it
|
||||
proxy_wide_config: turns on a proxy-wide setting through /config/update, so it runs on a stack no other suite is driving; deselected unless E2E_PROXY_WIDE_CONFIG_STACK is set
|
||||
|
|
|
|||
|
|
@ -17,11 +17,6 @@ from models import (
|
|||
ChatBody,
|
||||
ChatMessage,
|
||||
ChatResponse,
|
||||
ConfigPatchResponse,
|
||||
CostDiscountConfig,
|
||||
CostDiscountConfigResponse,
|
||||
CostMarginConfig,
|
||||
CostMarginConfigResponse,
|
||||
GuardrailCreateBody,
|
||||
GuardrailCreateResponse,
|
||||
KeyDeleteBody,
|
||||
|
|
@ -31,14 +26,11 @@ from models import (
|
|||
)
|
||||
from spend_e2e_client import SpendClient
|
||||
|
||||
pytestmark = [pytest.mark.e2e, pytest.mark.quiet_stack]
|
||||
pytestmark = pytest.mark.e2e
|
||||
|
||||
BACKEND_MODEL: Final = "openai/gpt-6.1-sol"
|
||||
INPUT_RATE: Final = 0.00005
|
||||
OUTPUT_RATE: Final = 0.0001
|
||||
DISCOUNT: Final = 0.25
|
||||
MARGIN_PERCENT: Final = 0.1
|
||||
MARGIN_FIXED: Final = 0.0005
|
||||
|
||||
|
||||
def _register_bedrock_guardrail(
|
||||
|
|
@ -92,52 +84,6 @@ def _guarded_chat(client: SpendClient, key: str, model: str, name: str) -> Strea
|
|||
)
|
||||
|
||||
|
||||
def _set_discount(client: SpendClient, values: dict[str, float]) -> float:
|
||||
unwrap(
|
||||
client.proxy.transport.patch(
|
||||
"/config/cost_discount_config",
|
||||
headers=client.proxy.transport.master,
|
||||
json=CostDiscountConfig(values),
|
||||
response_type=ConfigPatchResponse,
|
||||
)
|
||||
)
|
||||
return time.monotonic()
|
||||
|
||||
|
||||
def _set_margin(client: SpendClient, values: dict[str, float | dict[str, float]]) -> float:
|
||||
unwrap(
|
||||
client.proxy.transport.patch(
|
||||
"/config/cost_margin_config",
|
||||
headers=client.proxy.transport.master,
|
||||
json=CostMarginConfig(values),
|
||||
response_type=ConfigPatchResponse,
|
||||
)
|
||||
)
|
||||
return time.monotonic()
|
||||
|
||||
|
||||
def _get_discount(client: SpendClient) -> dict[str, float]:
|
||||
return unwrap(
|
||||
client.proxy.transport.get(
|
||||
"/config/cost_discount_config",
|
||||
headers=client.proxy.transport.master,
|
||||
params=NoBody(),
|
||||
response_type=CostDiscountConfigResponse,
|
||||
)
|
||||
).values
|
||||
|
||||
|
||||
def _get_margin(client: SpendClient) -> dict[str, float | dict[str, float]]:
|
||||
return unwrap(
|
||||
client.proxy.transport.get(
|
||||
"/config/cost_margin_config",
|
||||
headers=client.proxy.transport.master,
|
||||
params=NoBody(),
|
||||
response_type=CostMarginConfigResponse,
|
||||
)
|
||||
).values
|
||||
|
||||
|
||||
def _register_model(client: SpendClient, resources: ResourceManager, prefix: str) -> str:
|
||||
model: Final = f"{prefix}-{unique_marker()}"
|
||||
model_id: Final = client.proxy.create_model(
|
||||
|
|
@ -199,100 +145,7 @@ def scoped_key(client: SpendClient, strict_resources: ResourceManager) -> str:
|
|||
return key
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def restored_pricing_config(client: SpendClient) -> Iterator[None]:
|
||||
discount: Final = _get_discount(client)
|
||||
margin: Final = _get_margin(client)
|
||||
try:
|
||||
_set_discount(client, {})
|
||||
reset_at: Final = _set_margin(client, {})
|
||||
settle_propagation(reset_at)
|
||||
yield
|
||||
finally:
|
||||
try:
|
||||
_set_discount(client, discount)
|
||||
finally:
|
||||
restored_at: Final = _set_margin(client, margin)
|
||||
settle_propagation(restored_at)
|
||||
|
||||
|
||||
class TestPricingConfigSpend:
|
||||
@pytest.mark.covers(
|
||||
"quota_management.spend_tracking.discount_config.logs_cost",
|
||||
exercised_on=["chat_completions"],
|
||||
)
|
||||
@meta(
|
||||
Subject(
|
||||
domain=Domain.SPEND_BUDGETS,
|
||||
route=Route.SPEND_REPORTING,
|
||||
providers=(Provider.OPENAI,),
|
||||
models=(BACKEND_MODEL,),
|
||||
mode=Mode.NONSTREAM,
|
||||
)
|
||||
)
|
||||
def test_configured_discount_reaches_persisted_spend_row(
|
||||
self,
|
||||
client: SpendClient,
|
||||
strict_resources: ResourceManager,
|
||||
scoped_key: str,
|
||||
restored_pricing_config: None,
|
||||
) -> None:
|
||||
discount_set_at: Final = _set_discount(client, {"openai": DISCOUNT})
|
||||
settle_propagation(discount_set_at)
|
||||
model: Final = _register_model(client, strict_resources, "discount-priced")
|
||||
chat: Final = unwrap(client.chat(scoped_key, model, f"reply with one word {unique_marker()}", max_tokens=16))
|
||||
assert chat.id and chat.usage
|
||||
assert chat.usage.prompt_tokens is not None and chat.usage.prompt_tokens > 0
|
||||
assert chat.usage.completion_tokens is not None and chat.usage.completion_tokens >= 0
|
||||
|
||||
row: Final = poll_cost_row(client.proxy, chat.id)
|
||||
assert row is not None
|
||||
base_cost: Final = _base_cost(row, chat.usage.prompt_tokens, chat.usage.completion_tokens)
|
||||
breakdown: Final = row.breakdown
|
||||
assert breakdown.discount_percent is not None and approx_equal(breakdown.discount_percent, DISCOUNT)
|
||||
assert breakdown.discount_amount is not None and approx_equal(breakdown.discount_amount, base_cost * DISCOUNT)
|
||||
assert row.spend is not None and approx_equal(row.spend, base_cost * (1 - DISCOUNT))
|
||||
|
||||
@pytest.mark.covers(
|
||||
"quota_management.spend_tracking.margin_config.logs_cost",
|
||||
exercised_on=["chat_completions"],
|
||||
)
|
||||
@meta(
|
||||
Subject(
|
||||
domain=Domain.SPEND_BUDGETS,
|
||||
route=Route.SPEND_REPORTING,
|
||||
providers=(Provider.OPENAI,),
|
||||
models=(BACKEND_MODEL,),
|
||||
mode=Mode.NONSTREAM,
|
||||
)
|
||||
)
|
||||
def test_configured_margin_reaches_persisted_spend_row(
|
||||
self,
|
||||
client: SpendClient,
|
||||
strict_resources: ResourceManager,
|
||||
scoped_key: str,
|
||||
restored_pricing_config: None,
|
||||
) -> None:
|
||||
margin_set_at: Final = _set_margin(
|
||||
client, {"openai": {"percentage": MARGIN_PERCENT, "fixed_amount": MARGIN_FIXED}}
|
||||
)
|
||||
settle_propagation(margin_set_at)
|
||||
model: Final = _register_model(client, strict_resources, "margin-priced")
|
||||
chat: Final = unwrap(client.chat(scoped_key, model, f"reply with one word {unique_marker()}", max_tokens=16))
|
||||
assert chat.id and chat.usage
|
||||
assert chat.usage.prompt_tokens is not None and chat.usage.prompt_tokens > 0
|
||||
assert chat.usage.completion_tokens is not None and chat.usage.completion_tokens >= 0
|
||||
|
||||
row: Final = poll_cost_row(client.proxy, chat.id)
|
||||
assert row is not None
|
||||
base_cost: Final = _base_cost(row, chat.usage.prompt_tokens, chat.usage.completion_tokens)
|
||||
breakdown: Final = row.breakdown
|
||||
expected_margin: Final = base_cost * MARGIN_PERCENT + MARGIN_FIXED
|
||||
assert breakdown.margin_percent is not None and approx_equal(breakdown.margin_percent, MARGIN_PERCENT)
|
||||
assert breakdown.margin_fixed_amount is not None and approx_equal(breakdown.margin_fixed_amount, MARGIN_FIXED)
|
||||
assert breakdown.margin_total_amount is not None and approx_equal(breakdown.margin_total_amount, expected_margin)
|
||||
assert row.spend is not None and approx_equal(row.spend, base_cost + expected_margin)
|
||||
|
||||
class TestGuardrailCostSpend:
|
||||
@pytest.mark.covers(
|
||||
"quota_management.spend_tracking.guardrail_cost.logs_cost",
|
||||
exercised_on=["chat_completions"],
|
||||
|
|
@ -311,7 +164,6 @@ class TestPricingConfigSpend:
|
|||
client: SpendClient,
|
||||
strict_resources: ResourceManager,
|
||||
scoped_key: str,
|
||||
restored_pricing_config: None,
|
||||
) -> None:
|
||||
identifier: Final = os.environ["BEDROCK_GUARDRAIL_IDENTIFIER"]
|
||||
version: Final = os.environ["BEDROCK_GUARDRAIL_VERSION"]
|
||||
|
|
@ -1,45 +0,0 @@
|
|||
"""Cross-process reader/writer lock over the proxy stack every xdist worker shares.
|
||||
Every collected test holds it shared, marker or not, since the Claude Code cells and
|
||||
other unmarked suites drive the same stack; a `quiet_stack` test holds it exclusive,
|
||||
and the `gate` file makes a waiting exclusive holder win over readers that arrive
|
||||
after it."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import fcntl
|
||||
import hashlib
|
||||
import tempfile
|
||||
from collections.abc import Generator
|
||||
from contextlib import ExitStack, contextmanager
|
||||
from pathlib import Path
|
||||
from typing import Final
|
||||
|
||||
from e2e_config import PROXY_BASE_URL
|
||||
|
||||
STACK_DIGEST: Final = hashlib.sha256(PROXY_BASE_URL.encode()).hexdigest()[:12]
|
||||
LOCK_DIR: Final = Path(tempfile.gettempdir()) / f"litellm-e2e-stack-{STACK_DIGEST}"
|
||||
GATE_FILE: Final = LOCK_DIR / "gate"
|
||||
STACK_FILE: Final = LOCK_DIR / "stack"
|
||||
|
||||
|
||||
@contextmanager
|
||||
def _flock(path: Path, operation: int) -> Generator[None]:
|
||||
with path.open("a") as handle:
|
||||
fcntl.flock(handle, operation)
|
||||
try:
|
||||
yield
|
||||
finally:
|
||||
fcntl.flock(handle, fcntl.LOCK_UN)
|
||||
|
||||
|
||||
@contextmanager
|
||||
def stack_lock(exclusive: bool) -> Generator[None]:
|
||||
LOCK_DIR.mkdir(parents=True, exist_ok=True)
|
||||
if exclusive:
|
||||
with _flock(GATE_FILE, fcntl.LOCK_EX), _flock(STACK_FILE, fcntl.LOCK_EX):
|
||||
yield
|
||||
return
|
||||
with ExitStack() as held:
|
||||
with _flock(GATE_FILE, fcntl.LOCK_SH):
|
||||
held.enter_context(_flock(STACK_FILE, fcntl.LOCK_SH))
|
||||
yield
|
||||
|
|
@ -1,6 +1,6 @@
|
|||
# e2e harness tests
|
||||
|
||||
Tests of the harness under `tests/e2e/` (the transport, the clients, the fixture bundle and replay edge, the stack lock, the IdP launcher, the coverage collector, the JUnit properties and the Claude Code driver), not of the product. They live outside `tests/e2e/` because the Buildkite e2e run copies that folder into the runner image and runs every test in it, so a harness test in there counts as a product test in the nightly numbers. Nothing here needs a proxy, provider keys or the network
|
||||
Tests of the harness under `tests/e2e/` (the transport, the clients, the fixture bundle and replay edge, the IdP launcher, the coverage collector, the JUnit properties and the Claude Code driver), not of the product. They live outside `tests/e2e/` because the Buildkite e2e run copies that folder into the runner image and runs every test in it, so a harness test in there counts as a product test in the nightly numbers. Nothing here needs a proxy, provider keys or the network
|
||||
|
||||
The layout mirrors `tests/e2e/`: `test_e2e_http.py` covers `tests/e2e/e2e_http.py`, `logging/test_datadog_reader.py` covers `tests/e2e/logging/datadog_reader.py`, and `claude_code/` covers the driver, builder, probe and version resolver. Put a new harness test under the folder that mirrors the suite folder whose module it covers
|
||||
|
||||
|
|
|
|||
|
|
@ -12,10 +12,11 @@ injected clock, so nothing here monkeypatches anything.
|
|||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import time
|
||||
from builtins import ExceptionGroup
|
||||
from collections.abc import Callable, Generator, Iterable, Mapping
|
||||
from contextlib import contextmanager
|
||||
from dataclasses import dataclass
|
||||
from dataclasses import dataclass, replace
|
||||
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
|
||||
from itertools import chain, repeat
|
||||
from queue import SimpleQueue
|
||||
|
|
@ -40,7 +41,10 @@ from models import (
|
|||
LiteLLMParamsBody,
|
||||
McpServerCreateBody,
|
||||
McpServerUpdateBody,
|
||||
ModelDeleteBody,
|
||||
ModelInfoBody,
|
||||
ModelListEntry,
|
||||
ModelNewBody,
|
||||
ModelsListResponse,
|
||||
OrgNewBody,
|
||||
OrgUpdateBody,
|
||||
|
|
@ -59,6 +63,8 @@ from proxy_client import (
|
|||
ConvergeOutcome,
|
||||
CredentialKind,
|
||||
EverywhereConverged,
|
||||
ModelNotRegistered,
|
||||
ModelRegistered,
|
||||
ModelsPoller,
|
||||
NeverConvergedOn,
|
||||
NotConverged,
|
||||
|
|
@ -691,3 +697,85 @@ def test_partial_updates_preserve_explicit_null_at_the_http_boundary(operation:
|
|||
)
|
||||
assert json.loads(bodies.get_nowait()) == expected
|
||||
assert bodies.empty()
|
||||
|
||||
|
||||
@contextmanager
|
||||
def model_registry(
|
||||
listed: frozenset[str], refused: frozenset[str]
|
||||
) -> Generator[tuple[ProxyClient, SimpleQueue[float], SimpleQueue[str]]]:
|
||||
writes: Final[SimpleQueue[float]] = SimpleQueue()
|
||||
deleted: Final[SimpleQueue[str]] = SimpleQueue()
|
||||
|
||||
class Handler(BaseHTTPRequestHandler):
|
||||
def log_message(self, format: str, *args: object) -> None:
|
||||
pass
|
||||
|
||||
def reply(self, status: int, body: object) -> None:
|
||||
self.send_response(status)
|
||||
self.end_headers()
|
||||
self.wfile.write(json.dumps(body).encode())
|
||||
|
||||
def do_GET(self) -> None:
|
||||
self.reply(200, {"data": [{"id": name} for name in sorted(listed)]})
|
||||
|
||||
def do_POST(self) -> None:
|
||||
raw: Final = self.rfile.read(int(self.headers.get("Content-Length", "0")))
|
||||
if self.path == "/model/delete":
|
||||
deleted.put(ModelDeleteBody.model_validate_json(raw).id)
|
||||
self.reply(200, {})
|
||||
return
|
||||
writes.put(time.monotonic())
|
||||
name: Final = ModelNewBody.model_validate_json(raw).model_name
|
||||
if name in refused:
|
||||
self.reply(500, {"error": "refused"})
|
||||
return
|
||||
self.reply(200, {"model_id": f"id-{name}"})
|
||||
|
||||
server: Final = ThreadingHTTPServer(("127.0.0.1", 0), Handler)
|
||||
thread: Final = Thread(target=lambda: server.serve_forever(poll_interval=0.01), daemon=True)
|
||||
thread.start()
|
||||
url: Final = f"http://127.0.0.1:{server.server_port}"
|
||||
proxy: Final = build_proxy_client(
|
||||
base_url=url, control_plane_base_url=url, replica_urls=(url,), control_replica_urls=(url,), master_key="bootstrap"
|
||||
)
|
||||
try:
|
||||
yield replace(proxy, model_servable_timeout=0.3, model_servable_interval=0.01), writes, deleted
|
||||
finally:
|
||||
server.shutdown()
|
||||
server.server_close()
|
||||
thread.join(timeout=5)
|
||||
|
||||
|
||||
class TestRegisterModels:
|
||||
def test_a_batch_waits_once_after_its_last_write_and_reports_each_model(self) -> None:
|
||||
names: Final = ("first", "refused", "ghost", "last")
|
||||
settled: Final[SimpleQueue[float]] = SimpleQueue()
|
||||
with (
|
||||
model_registry(listed=frozenset({"first", "last"}), refused=frozenset({"refused"})) as (
|
||||
proxy,
|
||||
writes,
|
||||
deleted,
|
||||
),
|
||||
without_retries(),
|
||||
):
|
||||
registrations: Final = replace(proxy, settle=settled.put).register_models(
|
||||
tuple(
|
||||
ModelNewBody(
|
||||
model_name=name, litellm_params=LiteLLMParamsBody(model="synthetic"), model_info=ModelInfoBody()
|
||||
)
|
||||
for name in names
|
||||
)
|
||||
)
|
||||
last_write: Final = max(writes.get_nowait() for _ in names)
|
||||
assert writes.empty()
|
||||
assert settled.get_nowait() >= last_write
|
||||
assert settled.empty(), "a batch must wait out propagation once, not once per model"
|
||||
assert tuple((type(r), r.model_name) for r in registrations) == (
|
||||
(ModelRegistered, "first"),
|
||||
(ModelNotRegistered, "refused"),
|
||||
(ModelNotRegistered, "ghost"),
|
||||
(ModelRegistered, "last"),
|
||||
)
|
||||
assert tuple(r.model_id for r in registrations if isinstance(r, ModelRegistered)) == ("id-first", "id-last")
|
||||
assert deleted.get_nowait() == "id-ghost"
|
||||
assert deleted.empty()
|
||||
|
|
|
|||
|
|
@ -1,117 +0,0 @@
|
|||
"""Cross-process behavior of the stack lock: readers share it, an exclusive holder waits for
|
||||
every reader and keeps them out, and a reader arriving behind a waiting exclusive holder
|
||||
queues behind it instead of starving it."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import fcntl
|
||||
import inspect
|
||||
import os
|
||||
import subprocess
|
||||
import sys
|
||||
import time
|
||||
from contextlib import ExitStack
|
||||
from pathlib import Path
|
||||
from typing import Final
|
||||
|
||||
import pytest
|
||||
import stack_lock
|
||||
|
||||
HARNESS_DIR: Final = Path(inspect.getfile(stack_lock)).resolve().parent
|
||||
DEADLINE_SECONDS: Final = 30.0
|
||||
SETTLE_SECONDS: Final = 0.5
|
||||
HOLDER_SCRIPT: Final = """
|
||||
import sys, time
|
||||
from pathlib import Path
|
||||
from stack_lock import stack_lock
|
||||
name, mode, release_path, log_path = sys.argv[1:]
|
||||
|
||||
|
||||
def record(event):
|
||||
with Path(log_path).open("a") as log:
|
||||
log.write(f"{name} {event}\\n")
|
||||
|
||||
|
||||
record("waiting")
|
||||
with stack_lock(exclusive=mode == "exclusive"):
|
||||
record("enter")
|
||||
while not Path(release_path).exists():
|
||||
time.sleep(0.02)
|
||||
record("exit")
|
||||
"""
|
||||
|
||||
|
||||
def _events(log_path: Path) -> tuple[str, ...]:
|
||||
return tuple(log_path.read_text().splitlines()) if log_path.exists() else ()
|
||||
|
||||
|
||||
def _wait_for_event(log_path: Path, event: str) -> None:
|
||||
deadline: Final = time.monotonic() + DEADLINE_SECONDS
|
||||
while event not in _events(log_path):
|
||||
if time.monotonic() > deadline:
|
||||
pytest.fail(f"{event!r} never appeared; events so far: {_events(log_path)}")
|
||||
time.sleep(0.02)
|
||||
|
||||
|
||||
def _wait_until_gate_is_held_exclusively(gate_path: Path) -> None:
|
||||
deadline: Final = time.monotonic() + DEADLINE_SECONDS
|
||||
with gate_path.open("a") as handle:
|
||||
while True:
|
||||
try:
|
||||
fcntl.flock(handle, fcntl.LOCK_SH | fcntl.LOCK_NB)
|
||||
except BlockingIOError:
|
||||
return
|
||||
fcntl.flock(handle, fcntl.LOCK_UN)
|
||||
if time.monotonic() > deadline:
|
||||
pytest.fail("no exclusive holder ever took the gate")
|
||||
time.sleep(0.02)
|
||||
|
||||
|
||||
def _start_holder(held: ExitStack, tmp_path: Path, name: str, mode: str) -> subprocess.Popen[bytes]:
|
||||
holder: Final = held.enter_context(
|
||||
subprocess.Popen(
|
||||
(
|
||||
sys.executable,
|
||||
"-P",
|
||||
"-c",
|
||||
HOLDER_SCRIPT,
|
||||
name,
|
||||
mode,
|
||||
str(tmp_path / f"release-{name}"),
|
||||
str(tmp_path / "events"),
|
||||
),
|
||||
cwd=HARNESS_DIR,
|
||||
env={**os.environ, "TMPDIR": str(tmp_path), "PYTHONPATH": str(HARNESS_DIR)},
|
||||
)
|
||||
)
|
||||
held.callback(holder.kill)
|
||||
return holder
|
||||
|
||||
|
||||
def test_readers_share_exclusive_waits_and_a_waiting_exclusive_beats_later_readers(tmp_path: Path) -> None:
|
||||
lock_dir: Final = tmp_path / f"litellm-e2e-stack-{stack_lock.STACK_DIGEST}"
|
||||
lock_dir.mkdir()
|
||||
log_path: Final = tmp_path / "events"
|
||||
with ExitStack() as held:
|
||||
first_reader: Final = _start_holder(held, tmp_path, "A", "shared")
|
||||
_wait_for_event(log_path, "A enter")
|
||||
second_reader: Final = _start_holder(held, tmp_path, "R", "shared")
|
||||
_wait_for_event(log_path, "R enter")
|
||||
(tmp_path / "release-R").touch()
|
||||
_wait_for_event(log_path, "R exit")
|
||||
writer: Final = _start_holder(held, tmp_path, "W", "exclusive")
|
||||
_wait_until_gate_is_held_exclusively(lock_dir / "gate")
|
||||
late_reader: Final = _start_holder(held, tmp_path, "B", "shared")
|
||||
_wait_for_event(log_path, "B waiting")
|
||||
time.sleep(SETTLE_SECONDS)
|
||||
(tmp_path / "release-A").touch()
|
||||
_wait_for_event(log_path, "W enter")
|
||||
(tmp_path / "release-W").touch()
|
||||
_wait_for_event(log_path, "B enter")
|
||||
(tmp_path / "release-B").touch()
|
||||
for holder in (first_reader, second_reader, writer, late_reader):
|
||||
assert holder.wait(timeout=DEADLINE_SECONDS) == 0
|
||||
events: Final = _events(log_path)
|
||||
assert events.index("R enter") < events.index("A exit")
|
||||
assert events.index("W enter") > events.index("A exit")
|
||||
assert events.index("B enter") > events.index("W exit")
|
||||
|
|
@ -697,6 +697,54 @@ def test_per_user_authorization_code_with_pkce_binds_the_token_to_the_authorizin
|
|||
)
|
||||
|
||||
|
||||
@pytest.mark.timeout(600)
|
||||
def test_per_user_authorization_code_token_is_read_back_from_the_database_by_a_fresh_proxy_without_redis(
|
||||
gateway: Gateway, tmp_path: Path
|
||||
) -> None:
|
||||
with mcp_peer() as peer, oauth_server() as auth, gateway.scenario() as scenario:
|
||||
alias: Final = "cold" + uuid.uuid4().hex[:8]
|
||||
identity: Final = _register_oauth(
|
||||
scenario,
|
||||
peer,
|
||||
auth,
|
||||
alias,
|
||||
auth_type="oauth2",
|
||||
oauth2_flow="authorization_code",
|
||||
credentials={"client_id": "cold-client", "client_secret": "cold-secret", "scopes": ["tools.call"]},
|
||||
)
|
||||
owner_id: Final = scenario.user()
|
||||
owner: Final = scenario.key(user_id=owner_id, object_permission={"mcp_servers": [identity]})
|
||||
registered: Final = gateway.client.post(
|
||||
f"/{alias}/register", json={"redirect_uris": [CLIENT_REDIRECT], "client_name": "integration"}
|
||||
)
|
||||
assert registered.status_code in (200, 201), registered.text
|
||||
client_id: Final = str(_response_object(registered)["client_id"])
|
||||
pkce: Final = _Pkce(secrets.token_urlsafe(32))
|
||||
code, _ = _authorize_through_gateway(gateway, auth, alias, owner, client_id, pkce)
|
||||
redeemed: Final = _redeem(gateway, alias, owner, client_id, code, pkce)
|
||||
assert redeemed.status_code == 200, redeemed.text
|
||||
token: Final = _issued_token(redeemed.json())
|
||||
stored: Final = read_rows(
|
||||
'SELECT credential_b64 FROM "LiteLLM_MCPUserCredentials" WHERE user_id = %s AND server_id = %s',
|
||||
(owner_id, identity),
|
||||
)
|
||||
assert len(stored) == 1, stored
|
||||
stored_bytes: Final = base64.urlsafe_b64decode(str(stored[0]["credential_b64"]))
|
||||
assert token.encode() not in stored_bytes, "stored token is plaintext"
|
||||
config: Final = tmp_path / "no_redis.yaml"
|
||||
config.write_text(
|
||||
"general_settings:\n"
|
||||
" master_key: os.environ/LITELLM_MASTER_KEY\n"
|
||||
" database_url: os.environ/DATABASE_URL\n"
|
||||
" store_model_in_db: true\n"
|
||||
)
|
||||
peer.drain()
|
||||
with owned_proxy(gateway, tmp_path, {}, config=config, remove_environment=("REDIS_HOST", "REDIS_PORT")) as cold:
|
||||
called: Final = call_tool(cold, owner, identity, f"{alias}-add", ADD)
|
||||
assert called.status_code == 200, called.text
|
||||
assert _authorizations(peer) == (f"Bearer {token}".encode(),)
|
||||
|
||||
|
||||
def test_authorization_request_without_pkce_is_refused_before_reaching_the_authorization_server(
|
||||
gateway: Gateway,
|
||||
) -> None:
|
||||
|
|
|
|||
|
|
@ -7,7 +7,7 @@ import httpx
|
|||
import pytest
|
||||
from integration._support.client import Gateway, eventually, object_value, string_value
|
||||
from integration._support.database import read_rows
|
||||
from pydantic import BaseModel, JsonValue
|
||||
from pydantic import BaseModel, ConfigDict, JsonValue
|
||||
|
||||
|
||||
class CostEstimate(BaseModel):
|
||||
|
|
@ -110,6 +110,25 @@ def _assert_logged_spends(expectations: tuple[tuple[str, float], ...]) -> None:
|
|||
assert actual == pytest.approx(dict(expectations)), rows
|
||||
|
||||
|
||||
class LoggedBreakdown(BaseModel):
|
||||
model_config = ConfigDict(frozen=True)
|
||||
|
||||
original_cost: float
|
||||
discount_percent: float | None = None
|
||||
discount_amount: float | None = None
|
||||
margin_percent: float | None = None
|
||||
margin_fixed_amount: float | None = None
|
||||
margin_total_amount: float | None = None
|
||||
|
||||
|
||||
def _logged_breakdown(request_id: str) -> LoggedBreakdown:
|
||||
rows: Final = read_rows('SELECT metadata FROM "LiteLLM_SpendLogs" WHERE request_id = %s', (request_id,))
|
||||
assert len(rows) == 1, rows
|
||||
metadata: Final = rows[0]["metadata"]
|
||||
parsed: Final = object_value(json.loads(metadata) if isinstance(metadata, str) else metadata)
|
||||
return LoggedBreakdown.model_validate(parsed["cost_breakdown"])
|
||||
|
||||
|
||||
def _estimated(gateway: Gateway, model: str) -> CostEstimate:
|
||||
response: Final = gateway.request(
|
||||
"POST", "/cost/estimate", {"model": model, "input_tokens": 20, "output_tokens": 20}
|
||||
|
|
@ -236,6 +255,17 @@ def test_discount_and_margin_reprice_a_request_and_survive_rejection(gateway: Ga
|
|||
(streamed_global_request_id, 0.12),
|
||||
)
|
||||
)
|
||||
discounted: Final = _logged_breakdown(discounted_request_id)
|
||||
assert discounted.original_cost == pytest.approx(0.06), discounted
|
||||
assert discounted.discount_percent == pytest.approx(0.05), discounted
|
||||
assert discounted.discount_amount == pytest.approx(0.06 * 0.05), discounted
|
||||
margined: Final = _logged_breakdown(provider_margin_request_id)
|
||||
assert margined.original_cost == pytest.approx(0.06), margined
|
||||
assert margined.discount_percent == pytest.approx(0.05), margined
|
||||
assert margined.discount_amount == pytest.approx(0.06 * 0.05), margined
|
||||
assert margined.margin_percent == pytest.approx(0.1), margined
|
||||
assert margined.margin_fixed_amount == pytest.approx(0.001), margined
|
||||
assert margined.margin_total_amount == pytest.approx(0.057 * 0.1 + 0.001), margined
|
||||
|
||||
rejections: Final = (
|
||||
(
|
||||
|
|
|
|||
|
|
@ -0,0 +1,84 @@
|
|||
import json
|
||||
import uuid
|
||||
from collections.abc import Iterator
|
||||
from pathlib import Path
|
||||
from typing import Final
|
||||
|
||||
import pytest
|
||||
import yaml
|
||||
from integration._support.client import Gateway, eventually, gateway_from_environment, object_value
|
||||
from integration._support.process import owned_proxy
|
||||
from pydantic import JsonValue
|
||||
|
||||
UNREACHABLE_BASE: Final = "http://127.0.0.1:9/v1"
|
||||
TRANSCRIPT_TURNS: Final = 40
|
||||
RETRIES: Final = 2
|
||||
STORED_REQUEST_BUDGET_KB: Final = 64
|
||||
|
||||
|
||||
def _config_storing_prompts(directory: Path) -> Path:
|
||||
config: Final = object_value(yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text()))
|
||||
general_settings: Final = object_value(config["general_settings"])
|
||||
path: Final = directory / "store-prompts.yaml"
|
||||
path.write_text(
|
||||
yaml.safe_dump({**config, "general_settings": {**general_settings, "store_prompts_in_spend_logs": True}})
|
||||
)
|
||||
return path
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def gateway(tmp_path_factory: pytest.TempPathFactory) -> Iterator[Gateway]:
|
||||
with gateway_from_environment() as parent:
|
||||
directory: Final = tmp_path_factory.mktemp("failing_request_stored_snapshot")
|
||||
with owned_proxy(parent, directory, {}, config=_config_storing_prompts(directory)) as owned:
|
||||
yield owned
|
||||
|
||||
|
||||
def _transcript(marker: str) -> list[JsonValue]:
|
||||
turns: Final[list[JsonValue]] = [
|
||||
{"role": role, "content": f"turn {turn // 2} {role}"}
|
||||
for turn, role in enumerate(("user", "assistant") * TRANSCRIPT_TURNS)
|
||||
]
|
||||
return [*turns, {"role": "user", "content": marker}]
|
||||
|
||||
|
||||
def _stored_request(gateway: Gateway, call_id: str) -> JsonValue:
|
||||
rows: Final = eventually(
|
||||
lambda: gateway.rows('SELECT proxy_server_request FROM "LiteLLM_SpendLogs" WHERE request_id = %s', (call_id,)),
|
||||
lambda values: len(values) == 1,
|
||||
seconds=40,
|
||||
)
|
||||
stored: Final = rows[0]["proxy_server_request"]
|
||||
return json.loads(stored) if isinstance(stored, str) else stored
|
||||
|
||||
|
||||
def test_failing_request_with_retries_and_fallback_stores_one_copy_of_the_request(gateway: Gateway) -> None:
|
||||
with gateway.scenario() as scenario:
|
||||
primary: Final = f"stored-snapshot-{uuid.uuid4().hex}"
|
||||
fallback: Final = f"stored-snapshot-fb-{uuid.uuid4().hex}"
|
||||
for name in (primary, primary, fallback, fallback):
|
||||
scenario.model(name=name, api_base=UNREACHABLE_BASE, cooldown_time=0)
|
||||
key: Final = scenario.key(models=[primary, fallback])
|
||||
marker: Final = f"stored snapshot {uuid.uuid4().hex}"
|
||||
|
||||
response: Final = gateway.request(
|
||||
"POST",
|
||||
"/v1/chat/completions",
|
||||
{
|
||||
"model": primary,
|
||||
"max_tokens": 16,
|
||||
"messages": _transcript(marker),
|
||||
"router_settings_override": {"num_retries": RETRIES, "fallbacks": [{primary: [fallback]}]},
|
||||
},
|
||||
key=key,
|
||||
)
|
||||
assert response.status_code != 200, response.text
|
||||
assert fallback in response.text, response.text
|
||||
|
||||
serialized: Final = json.dumps(_stored_request(gateway, response.headers["x-litellm-call-id"]))
|
||||
stored_kb: Final = len(serialized.encode()) / 1024
|
||||
assert marker in serialized, serialized[:400]
|
||||
assert stored_kb <= STORED_REQUEST_BUDGET_KB, (
|
||||
f"one failing request stored a {stored_kb:.0f} KB request snapshot for a "
|
||||
f"{2 * TRANSCRIPT_TURNS + 1}-message transcript with {RETRIES} retries and a fallback"
|
||||
)
|
||||
|
|
@ -1,222 +0,0 @@
|
|||
"""Unit tests for the Friendli transform in
|
||||
`.github/scripts/auto_update_price_and_context_window_file.py`."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import importlib.util
|
||||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
import pytest
|
||||
|
||||
SCRIPT_PATH = (
|
||||
Path(__file__).resolve().parents[2]
|
||||
/ ".github"
|
||||
/ "scripts"
|
||||
/ "auto_update_price_and_context_window_file.py"
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def sync_module():
|
||||
spec = importlib.util.spec_from_file_location(
|
||||
"auto_update_price_and_context_window_file", SCRIPT_PATH
|
||||
)
|
||||
assert spec and spec.loader, f"Could not load spec for {SCRIPT_PATH}"
|
||||
module = importlib.util.module_from_spec(spec)
|
||||
sys.modules["auto_update_price_and_context_window_file"] = module
|
||||
spec.loader.exec_module(module)
|
||||
return module
|
||||
|
||||
|
||||
def _reasoning_model(**overrides: object) -> dict:
|
||||
model = {
|
||||
"id": "zai-org/GLM-Test",
|
||||
"base_model": "zhipuai/glm-test",
|
||||
"context_length": 1048576,
|
||||
"max_completion_tokens": 131072,
|
||||
"pricing": {"input": "0.00000015", "output": "0.0000005", "input_cache_read": "0.00000003"},
|
||||
"reasoning": True,
|
||||
"reasoning_options": [{"type": "effort", "values": ["max", "high", "low"]}],
|
||||
"functionality": {
|
||||
"tool_call": True,
|
||||
"parallel_tool_call": True,
|
||||
"structured_output": True,
|
||||
"system_messages": True,
|
||||
"tool_choice": True,
|
||||
},
|
||||
"input_modalities": ["text", "image", "video"],
|
||||
"mode": "chat",
|
||||
}
|
||||
model.update(overrides)
|
||||
return model
|
||||
|
||||
|
||||
def test_transform_emits_declared_effort_levels_in_canonical_order(sync_module):
|
||||
entry = sync_module.transform_friendli_data([_reasoning_model()], {})[
|
||||
"friendliai/zai-org/GLM-Test"
|
||||
]
|
||||
assert entry["supports_reasoning"] is True
|
||||
assert entry["reasoning_effort_levels"] == ["low", "high", "max"]
|
||||
assert not any(k.endswith("_reasoning_effort") for k in entry)
|
||||
|
||||
|
||||
def test_transform_reasoning_model_without_effort_options_declares_empty_levels(sync_module):
|
||||
model = _reasoning_model(reasoning_options=[{"type": "budget_tokens", "values": []}])
|
||||
entry = sync_module.transform_friendli_data([model], {})["friendliai/zai-org/GLM-Test"]
|
||||
assert entry["reasoning_effort_levels"] == []
|
||||
|
||||
|
||||
def test_transform_non_reasoning_model_declares_no_levels(sync_module):
|
||||
model = _reasoning_model(reasoning=False, reasoning_options=[])
|
||||
entry = sync_module.transform_friendli_data([model], {})["friendliai/zai-org/GLM-Test"]
|
||||
assert entry["supports_reasoning"] is False
|
||||
assert "reasoning_effort_levels" not in entry
|
||||
|
||||
|
||||
def test_transform_max_tokens_mirrors_output_cap_not_context(sync_module):
|
||||
entry = sync_module.transform_friendli_data([_reasoning_model()], {})[
|
||||
"friendliai/zai-org/GLM-Test"
|
||||
]
|
||||
assert entry["max_input_tokens"] == 1048576
|
||||
assert entry["max_output_tokens"] == 131072
|
||||
assert entry["max_tokens"] == entry["max_output_tokens"]
|
||||
|
||||
|
||||
def test_transform_prompt_caching_follows_cache_pricing(sync_module):
|
||||
cached = sync_module.transform_friendli_data([_reasoning_model()], {})[
|
||||
"friendliai/zai-org/GLM-Test"
|
||||
]
|
||||
assert cached["supports_prompt_caching"] is True
|
||||
assert cached["cache_read_input_token_cost"] == 3e-08
|
||||
|
||||
uncached_model = _reasoning_model(pricing={"input": "0.00000014", "output": "0.0000004"})
|
||||
uncached = sync_module.transform_friendli_data([uncached_model], {})[
|
||||
"friendliai/zai-org/GLM-Test"
|
||||
]
|
||||
assert uncached["supports_prompt_caching"] is False
|
||||
assert "cache_read_input_token_cost" not in uncached
|
||||
|
||||
|
||||
def test_transform_modalities_set_vision_image_and_video_flags(sync_module):
|
||||
entry = sync_module.transform_friendli_data([_reasoning_model()], {})[
|
||||
"friendliai/zai-org/GLM-Test"
|
||||
]
|
||||
assert entry["supports_vision"] is True
|
||||
assert entry["supports_image_input"] is True
|
||||
assert entry["supports_video_input"] is True
|
||||
|
||||
text_only = _reasoning_model(input_modalities=["text"])
|
||||
entry_text = sync_module.transform_friendli_data([text_only], {})[
|
||||
"friendliai/zai-org/GLM-Test"
|
||||
]
|
||||
assert entry_text["supports_vision"] is False
|
||||
assert entry_text["supports_image_input"] is False
|
||||
assert entry_text["supports_video_input"] is False
|
||||
|
||||
|
||||
def test_transform_skips_rows_without_valid_token_prices_so_priced_local_entries_survive(sync_module):
|
||||
local = {
|
||||
"friendliai/zai-org/GLM-Test": {
|
||||
"litellm_provider": "friendliai",
|
||||
"input_cost_per_token": 1.5e-07,
|
||||
"output_cost_per_token": 5e-07,
|
||||
}
|
||||
}
|
||||
unpriced_rows = [
|
||||
_reasoning_model(pricing={}),
|
||||
_reasoning_model(pricing=None),
|
||||
_reasoning_model(pricing={"input": "0.00000015"}),
|
||||
_reasoning_model(pricing={"output": "0.0000005"}),
|
||||
_reasoning_model(pricing={"input": "not-a-number", "output": "0.0000005"}),
|
||||
_reasoning_model(pricing={"input": "-0.00000015", "output": "0.0000005"}),
|
||||
_reasoning_model(pricing={"input": "inf", "output": "0.0000005"}),
|
||||
_reasoning_model(pricing={"input": "nan", "output": "0.0000005"}),
|
||||
]
|
||||
remote = sync_module.transform_friendli_data(unpriced_rows, local)
|
||||
assert remote == {}
|
||||
sync_module.sync_local_data_with_remote(local, remote, replace_keys=frozenset(remote))
|
||||
assert local["friendliai/zai-org/GLM-Test"]["input_cost_per_token"] == 1.5e-07
|
||||
assert local["friendliai/zai-org/GLM-Test"]["output_cost_per_token"] == 5e-07
|
||||
|
||||
|
||||
def test_transform_keeps_zero_priced_rows(sync_module):
|
||||
free_model = _reasoning_model(pricing={"input": "0", "output": "0"})
|
||||
entry = sync_module.transform_friendli_data([free_model], {})["friendliai/zai-org/GLM-Test"]
|
||||
assert entry["input_cost_per_token"] == 0.0
|
||||
assert entry["output_cost_per_token"] == 0.0
|
||||
|
||||
|
||||
def test_transforms_survive_failed_fetch(sync_module):
|
||||
assert sync_module.transform_friendli_data(None, {}) == {}
|
||||
assert sync_module.transform_friendli_data([], {}) == {}
|
||||
assert sync_module.transform_openrouter_data(None) == {}
|
||||
assert sync_module.transform_vercel_ai_gateway_data(None) == {}
|
||||
|
||||
|
||||
def test_vercel_transform_skips_rows_without_token_pricing_or_limits(sync_module):
|
||||
rows = [
|
||||
{
|
||||
"id": "wan-video",
|
||||
"pricing": {"video_duration_pricing": [{"resolution": "720p", "cost_per_second": "0.1"}]},
|
||||
},
|
||||
{
|
||||
"id": "qwen3-embedding",
|
||||
"context_window": 32768,
|
||||
"max_tokens": 32768,
|
||||
"pricing": {"input": "0.00000001"},
|
||||
},
|
||||
{
|
||||
"id": "no-limits-chat",
|
||||
"pricing": {"input": "0.000001", "output": "0.000002"},
|
||||
},
|
||||
{
|
||||
"id": "good-chat",
|
||||
"context_window": 128000,
|
||||
"max_tokens": 8192,
|
||||
"pricing": {"input": "0.000001", "output": "0.000002"},
|
||||
},
|
||||
]
|
||||
transformed = sync_module.transform_vercel_ai_gateway_data(rows)
|
||||
assert list(transformed) == ["vercel_ai_gateway/good-chat"]
|
||||
assert transformed["vercel_ai_gateway/good-chat"]["input_cost_per_token"] == 1e-06
|
||||
assert transformed["vercel_ai_gateway/good-chat"]["output_cost_per_token"] == 2e-06
|
||||
|
||||
|
||||
def test_sync_replaces_friendli_entries_so_dropped_cache_pricing_does_not_survive(sync_module):
|
||||
local = {
|
||||
"friendliai/zai-org/GLM-Test": {
|
||||
"litellm_provider": "friendliai",
|
||||
"cache_read_input_token_cost": 3e-08,
|
||||
"supports_prompt_caching": True,
|
||||
}
|
||||
}
|
||||
uncached_model = _reasoning_model(pricing={"input": "0.00000014", "output": "0.0000004"})
|
||||
remote = sync_module.transform_friendli_data([uncached_model], local)
|
||||
sync_module.sync_local_data_with_remote(local, remote, replace_keys=frozenset(remote))
|
||||
synced = local["friendliai/zai-org/GLM-Test"]
|
||||
assert "cache_read_input_token_cost" not in synced
|
||||
assert synced["supports_prompt_caching"] is False
|
||||
|
||||
|
||||
def test_sync_still_merges_entries_outside_replace_keys(sync_module):
|
||||
local = {"openrouter/some-model": {"input_cost_per_token": 1e-06, "supports_vision": True}}
|
||||
remote = {"openrouter/some-model": {"input_cost_per_token": 2e-06}}
|
||||
sync_module.sync_local_data_with_remote(local, remote)
|
||||
assert local["openrouter/some-model"] == {"input_cost_per_token": 2e-06, "supports_vision": True}
|
||||
|
||||
|
||||
def test_transform_inherits_allowlisted_keys_from_base_model_entry(sync_module):
|
||||
local = {
|
||||
"zhipuai/glm-test": {
|
||||
"supports_pdf_input": True,
|
||||
"supports_assistant_prefill": True,
|
||||
"input_cost_per_token": 9e-06,
|
||||
}
|
||||
}
|
||||
entry = sync_module.transform_friendli_data([_reasoning_model()], local)[
|
||||
"friendliai/zai-org/GLM-Test"
|
||||
]
|
||||
assert entry["supports_pdf_input"] is True
|
||||
assert entry["supports_assistant_prefill"] is True
|
||||
assert entry["input_cost_per_token"] == 1.5e-07
|
||||
|
|
@ -52,7 +52,7 @@ CI = [".github/workflows/test-unit.yml"]
|
|||
("mcp-dependencies", ["pyproject.toml"], "run"),
|
||||
("mcp-dependencies", ["uv.lock"], "run"),
|
||||
("mcp-dependencies", ["litellm/experimental_mcp_client/client.py"], "run"),
|
||||
("mcp-dependencies", ["tests/e2e/mcp/oauth_chat_client.py"], "run"),
|
||||
("mcp-dependencies", ["tests/e2e/mcp/datadog_mcp.py"], "run"),
|
||||
("mcp-dependencies", ["litellm-proxy-extras/pyproject.toml"], "run"),
|
||||
("mcp-dependencies", ["scripts/check_mcp_sdk_install.py"], "run"),
|
||||
("mcp-dependencies", [".github/workflows/test-mcp-dependency-resolution.yml"], "run"),
|
||||
|
|
|
|||
|
|
@ -1,5 +1,6 @@
|
|||
from __future__ import annotations
|
||||
|
||||
import fnmatch
|
||||
import re
|
||||
import shlex
|
||||
import subprocess
|
||||
|
|
@ -7,6 +8,7 @@ from pathlib import Path
|
|||
from typing import Final
|
||||
|
||||
import pytest
|
||||
import yaml
|
||||
|
||||
WORKFLOW: Final = Path(__file__).resolve().parents[2] / ".github" / "workflows" / "test-linting.yml"
|
||||
DIFF_GATE: Final = re.compile(r'git diff --name-only --diff-filter=\w+ "\$GATE_BASE_SHA" HEAD -- (.+?) \|')
|
||||
|
|
@ -100,3 +102,101 @@ def test_e2e_harness_gate_fires_on_any_e2e_or_harness_file_its_dependency_manife
|
|||
(trigger, "tests/e2e/ui/spec.ts", "tests/e2e/ui/pkg/page.py", "elsewhere/pyproject.toml", "tests/e2e_other/module.py"),
|
||||
)
|
||||
assert selected == frozenset({trigger})
|
||||
|
||||
|
||||
UI_API_TYPES_DETECT: Final[str] = next(
|
||||
step["run"]
|
||||
for step in yaml.safe_load(WORKFLOW.read_text())["jobs"]["ui-api-types"]["steps"]
|
||||
if step.get("id") == "changes"
|
||||
)
|
||||
PRE_COMMIT_LINT: Final = WORKFLOW.parents[2] / "scripts" / "pre_commit_lint.sh"
|
||||
TERRAFORM_PROVIDER_WORKFLOW: Final = WORKFLOW.with_name("test-terraform-provider.yml")
|
||||
|
||||
|
||||
def _ui_api_types_relevant(tmp_path: Path, changed: str) -> bool:
|
||||
_git(tmp_path, "init", "-q", "-b", "main")
|
||||
_git(tmp_path, "config", "user.email", "t@t")
|
||||
_git(tmp_path, "config", "user.name", "t")
|
||||
_git(tmp_path, "commit", "-q", "--allow-empty", "-m", "base")
|
||||
_git(tmp_path, "switch", "-q", "-c", "pr")
|
||||
target: Final = tmp_path / changed
|
||||
target.parent.mkdir(parents=True, exist_ok=True)
|
||||
target.write_text("x = 1\n")
|
||||
_git(tmp_path, "add", "-A")
|
||||
_git(tmp_path, "commit", "-qm", "change")
|
||||
_git(tmp_path, "switch", "-q", "main")
|
||||
_git(tmp_path, "commit", "-q", "--allow-empty", "-m", "main moved")
|
||||
_git(tmp_path, "merge", "-q", "--no-ff", "--no-edit", "pr")
|
||||
output: Final = tmp_path / "github_output"
|
||||
subprocess.run(
|
||||
["bash", "-c", UI_API_TYPES_DETECT],
|
||||
cwd=tmp_path,
|
||||
check=True,
|
||||
capture_output=True,
|
||||
env={"PATH": "/usr/bin:/bin", "GITHUB_OUTPUT": str(output)},
|
||||
)
|
||||
return output.read_text().strip() == "relevant=true"
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"changed",
|
||||
(
|
||||
"litellm/proxy/proxy_server.py",
|
||||
"litellm/types/router.py",
|
||||
"litellm/tracing/generated/models.py",
|
||||
"litellm/router_strategy/complexity_router/config.py",
|
||||
"litellm/models/credentials.py",
|
||||
"enterprise/litellm_enterprise/types/proxy/audit_logging_endpoints.py",
|
||||
"pyproject.toml",
|
||||
"uv.lock",
|
||||
"ui/litellm-dashboard/scripts/gen-api-types.mjs",
|
||||
"ui/litellm-dashboard/package-lock.json",
|
||||
"ui/litellm-dashboard/src/lib/http/schema.d.ts",
|
||||
),
|
||||
)
|
||||
def test_ui_api_types_check_runs_on_a_pull_request_touching_any_spec_input(tmp_path: Path, changed: str) -> None:
|
||||
assert _ui_api_types_relevant(tmp_path, changed)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"changed",
|
||||
(
|
||||
"litellm-rust/crates/traces/src/lib.rs",
|
||||
"litellm-proxy-extras/litellm_proxy_extras/utils.py",
|
||||
"tests/unit/test_router.py",
|
||||
"ui/litellm-dashboard/src/app/page.tsx",
|
||||
"docs/my-website/docs/proxy/config.md",
|
||||
"litellm-dashboard/pyproject.toml",
|
||||
),
|
||||
)
|
||||
def test_ui_api_types_check_skips_a_pull_request_that_cannot_change_the_spec(tmp_path: Path, changed: str) -> None:
|
||||
assert not _ui_api_types_relevant(tmp_path, changed)
|
||||
|
||||
|
||||
def _first_single_quoted_capture(pattern: str, text: str) -> str:
|
||||
match: Final = re.search(pattern, text, re.MULTILINE)
|
||||
assert match is not None, pattern
|
||||
return match.group(1)
|
||||
|
||||
|
||||
def _tracked_paths_matching(extended_regex: str) -> frozenset[str]:
|
||||
tracked: Final = _git(WORKFLOW.parents[2], "ls-files")
|
||||
selected: Final = subprocess.run(["grep", "-E", extended_regex], input=tracked, capture_output=True, text=True)
|
||||
return frozenset(selected.stdout.splitlines())
|
||||
|
||||
|
||||
def test_make_check_runs_the_api_type_sync_on_the_same_tracked_files_as_ci() -> None:
|
||||
ci_trigger: Final = _first_single_quoted_capture(r"grep -Eq '([^']+)'", UI_API_TYPES_DETECT)
|
||||
make_check_trigger: Final = _first_single_quoted_capture(r"^spec_pattern='([^']+)'$", PRE_COMMIT_LINT.read_text())
|
||||
ci_selected: Final = _tracked_paths_matching(ci_trigger)
|
||||
assert "pyproject.toml" in ci_selected
|
||||
assert _tracked_paths_matching(make_check_trigger) == ci_selected - {".github/workflows/test-linting.yml"}
|
||||
|
||||
|
||||
def test_a_pull_request_changing_any_route_file_runs_the_terraform_provider_endpoint_audit() -> None:
|
||||
triggers: Final = yaml.safe_load(TERRAFORM_PROVIDER_WORKFLOW.read_text())[True]["pull_request"]["paths"]
|
||||
route_files: Final = _git(
|
||||
WORKFLOW.parents[2], "grep", "-l", "-E", r"APIRouter\(", "--", "*.py", ":(exclude)tests/**"
|
||||
).splitlines()
|
||||
assert route_files
|
||||
assert [name for name in route_files if not any(fnmatch.fnmatch(name, pattern) for pattern in triggers)] == []
|
||||
|
|
|
|||
|
|
@ -427,6 +427,34 @@ def test_run_ends_with_a_summary_of_ran_and_skipped_blocks(tmp_path: Path) -> No
|
|||
assert "check: FAIL" not in proc.stdout
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"spec_input",
|
||||
(
|
||||
"litellm/tracing/generated/models.py",
|
||||
"litellm/models/credentials.py",
|
||||
"enterprise/litellm_enterprise/types/proxy/audit_logging_endpoints.py",
|
||||
"pyproject.toml",
|
||||
"uv.lock",
|
||||
),
|
||||
)
|
||||
def test_a_staged_spec_input_outside_proxy_and_types_runs_the_api_type_sync(tmp_path: Path, spec_input: str) -> None:
|
||||
repo, bin_dir = _sandbox(tmp_path)
|
||||
_commit_all(repo, "base")
|
||||
_stage_file(repo, spec_input, "x = 1\n")
|
||||
proc = _run(repo, bin_dir, {})
|
||||
assert proc.returncode == 0, proc.stdout + proc.stderr
|
||||
assert "ran: dashboard API-type sync (npm run gen:api)" in proc.stdout
|
||||
|
||||
|
||||
def test_a_staged_rust_crate_change_skips_the_api_type_sync(tmp_path: Path) -> None:
|
||||
repo, bin_dir = _sandbox(tmp_path)
|
||||
_commit_all(repo, "base")
|
||||
_stage_file(repo, "litellm-rust/crates/traces/src/lib.rs", "fn main() {}\n")
|
||||
proc = _run(repo, bin_dir, {})
|
||||
assert proc.returncode == 0, proc.stdout + proc.stderr
|
||||
assert "skipped: dashboard API-type sync (npm run gen:api)" in proc.stdout
|
||||
|
||||
|
||||
def test_staged_files_matching_no_check_print_an_explicit_noop_note_and_nonempty_log(tmp_path: Path) -> None:
|
||||
repo, bin_dir = _sandbox(tmp_path)
|
||||
_commit_all(repo, "base")
|
||||
|
|
|
|||
|
|
@ -1,378 +0,0 @@
|
|||
import importlib.util
|
||||
import json
|
||||
from pathlib import Path
|
||||
from types import MappingProxyType
|
||||
|
||||
import pytest
|
||||
|
||||
ROOT = Path(__file__).resolve().parents[2]
|
||||
SCRIPT = ROOT / "scripts" / "sync_together_ai_models.py"
|
||||
FIXTURES = Path(__file__).resolve().parent / "fixtures" / "together_ai_sync"
|
||||
|
||||
_spec = importlib.util.spec_from_file_location("sync_together_ai_models", SCRIPT)
|
||||
assert _spec is not None and _spec.loader is not None
|
||||
sync = importlib.util.module_from_spec(_spec)
|
||||
_spec.loader.exec_module(sync)
|
||||
|
||||
RECORDED_CATALOG = sync.load_catalog(FIXTURES.joinpath("models_serverless.json").read_bytes())
|
||||
RECORDED_DOC = sync.parse_deprecations(FIXTURES.joinpath("deprecations.md").read_text())
|
||||
|
||||
|
||||
def _doc(removal_dates: dict[str, str], redirects: dict[str, str] | None = None) -> object:
|
||||
return sync.DeprecationDoc(
|
||||
removal_dates=MappingProxyType(removal_dates),
|
||||
redirects=MappingProxyType(redirects or {}),
|
||||
)
|
||||
|
||||
|
||||
def _chat_model(model_id: str, ctx: int = 4096, price: float = 1.0, cached: float | None = None) -> object:
|
||||
return sync.CatalogModel(
|
||||
id=model_id,
|
||||
type="chat",
|
||||
context_length=ctx,
|
||||
pricing=sync.CatalogPricing(input=price, output=price, cached_input=cached),
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("per_million", "expected"),
|
||||
[
|
||||
(3, 3e-06),
|
||||
(15, 1.5e-05),
|
||||
(1.4, 1.4e-06),
|
||||
(0.25999999999999995, 2.6e-07),
|
||||
(0.060000000000000005, 6e-08),
|
||||
(1.0399999999999998, 1.04e-06),
|
||||
(0, 0.0),
|
||||
],
|
||||
)
|
||||
def test_per_token_normalizes_float_artifacts(per_million: float, expected: float) -> None:
|
||||
assert sync.per_token(per_million) == expected
|
||||
|
||||
|
||||
def test_parse_deprecations_recorded_fixture() -> None:
|
||||
assert dict(RECORDED_DOC.redirects) == {
|
||||
"mistralai/Mistral-7B-Instruct-v0.3": "mistralai/Ministral-3-14B-Instruct-2512",
|
||||
"Kimi-K2": "Kimi-K2-0905",
|
||||
"DeepSeek-V3": "DeepSeek-V3.1",
|
||||
"DeepSeek-V3-0324": "DeepSeek-V3.1",
|
||||
"DeepSeek-R1": "DeepSeek-R1-0528",
|
||||
}
|
||||
assert len(RECORDED_DOC.removal_dates) == 208
|
||||
assert RECORDED_DOC.removal_dates["google/gemma-3n-E4B-it"] == "2026-08-04"
|
||||
|
||||
|
||||
def test_parse_deprecations_duplicate_rows_keep_most_recent_date() -> None:
|
||||
assert RECORDED_DOC.removal_dates["Qwen/Qwen3-235B-A22B-Thinking-2507"] == "2026-04-16"
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"markdown",
|
||||
[
|
||||
"# Deprecations\n\nNothing here anymore.\n",
|
||||
"\n## Active model redirects\n\n| A | B |\n| --- | --- |\n| `x` | `y` |\n\n## Something else\n",
|
||||
"\n## Deprecation history\n\n### Inference\n\n| Date | Model | R |\n| --- | --- | --- |\n| 2026-01-01 | `m` | No |\n",
|
||||
],
|
||||
)
|
||||
def test_parse_deprecations_raises_when_a_table_parses_empty(markdown: str) -> None:
|
||||
with pytest.raises(sync.SyncError):
|
||||
sync.parse_deprecations(markdown)
|
||||
|
||||
|
||||
def test_load_catalog_raises_on_shape_change() -> None:
|
||||
with pytest.raises(sync.SyncError):
|
||||
sync.load_catalog(b'[{"id": "x", "type": "chat"}]')
|
||||
|
||||
|
||||
def test_load_catalog_raises_when_no_token_models_remain() -> None:
|
||||
only_video = json.dumps([{"id": "v", "type": "video", "pricing": {"input": 0, "output": 0}}]).encode()
|
||||
with pytest.raises(sync.SyncError):
|
||||
sync.load_catalog(only_video)
|
||||
|
||||
|
||||
def test_recorded_catalog_counts() -> None:
|
||||
assert len(RECORDED_CATALOG) == 102
|
||||
assert sum(1 for model in RECORDED_CATALOG if model.type in sync.TYPE_TO_MODE) == 26
|
||||
assert sum(1 for model in RECORDED_CATALOG if model.pricing.cached_input) == 13
|
||||
|
||||
|
||||
def test_added_chat_model_matches_reviewed_registry_shape() -> None:
|
||||
outcome = sync.compute_sync({}, RECORDED_CATALOG, RECORDED_DOC)
|
||||
assert len(outcome.added) == 26
|
||||
assert not outcome.deprecated
|
||||
assert outcome.cost_map["together_ai/moonshotai/Kimi-K3"] == {
|
||||
"cache_read_input_token_cost": 3e-07,
|
||||
"input_cost_per_token": 3e-06,
|
||||
"litellm_provider": "together_ai",
|
||||
"max_input_tokens": 1048576,
|
||||
"max_tokens": 1048576,
|
||||
"mode": "chat",
|
||||
"output_cost_per_token": 1.5e-05,
|
||||
"source": "https://docs.together.ai/docs/serverless-models",
|
||||
"supports_function_calling": True,
|
||||
"supports_parallel_function_calling": True,
|
||||
"supports_prompt_caching": True,
|
||||
"supports_reasoning": True,
|
||||
"supports_response_schema": True,
|
||||
"supports_tool_choice": True,
|
||||
"supports_vision": True,
|
||||
}
|
||||
|
||||
|
||||
def test_added_embedding_model_has_no_output_token_cap() -> None:
|
||||
outcome = sync.compute_sync({}, RECORDED_CATALOG, RECORDED_DOC)
|
||||
assert outcome.cost_map["together_ai/intfloat/multilingual-e5-large-instruct"] == {
|
||||
"input_cost_per_token": 2e-08,
|
||||
"litellm_provider": "together_ai",
|
||||
"max_input_tokens": 514,
|
||||
"max_tokens": 514,
|
||||
"mode": "embedding",
|
||||
"output_cost_per_token": 2e-08,
|
||||
"output_vector_size": 1024,
|
||||
"source": "https://docs.together.ai/docs/serverless-models",
|
||||
}
|
||||
|
||||
|
||||
def test_moderation_type_maps_to_chat_mode() -> None:
|
||||
outcome = sync.compute_sync({}, RECORDED_CATALOG, RECORDED_DOC)
|
||||
guard = outcome.cost_map["together_ai/meta-llama/Llama-Guard-4-12B"]
|
||||
assert guard["mode"] == "chat"
|
||||
assert "max_output_tokens" not in guard
|
||||
|
||||
|
||||
def test_output_ceiling_comes_from_the_rule_never_from_context_length() -> None:
|
||||
glm = next(model for model in RECORDED_CATALOG if model.id == "zai-org/GLM-5.2")
|
||||
fresh = sync.compute_sync({}, [_chat_model("acme/unreviewed", ctx=1048576), glm], _doc({"x": "2026-01-01"}))
|
||||
unreviewed = fresh.cost_map["together_ai/acme/unreviewed"]
|
||||
assert "max_output_tokens" not in unreviewed
|
||||
assert (unreviewed["max_input_tokens"], unreviewed["max_tokens"]) == (1048576, 1048576)
|
||||
reviewed = fresh.cost_map["together_ai/zai-org/GLM-5.2"]
|
||||
assert (reviewed["max_input_tokens"], reviewed["max_output_tokens"], reviewed["max_tokens"]) == (
|
||||
1048575,
|
||||
128000,
|
||||
128000,
|
||||
)
|
||||
inflated = {
|
||||
"together_ai/zai-org/GLM-5.2": {
|
||||
"input_cost_per_token": 1.4e-06,
|
||||
"litellm_provider": "together_ai",
|
||||
"max_input_tokens": 1048575,
|
||||
"max_output_tokens": 1048575,
|
||||
"max_tokens": 1048575,
|
||||
"mode": "chat",
|
||||
"output_cost_per_token": 4.4e-06,
|
||||
}
|
||||
}
|
||||
corrected = sync.compute_sync(inflated, [glm], _doc({"x": "2026-01-01"}))
|
||||
assert corrected.cost_map["together_ai/zai-org/GLM-5.2"]["max_output_tokens"] == 128000
|
||||
assert any("max_output_tokens: 1048575 -> 128000" in line for line in corrected.updated)
|
||||
|
||||
|
||||
def test_docs_removed_but_live_model_stays_live_with_warning() -> None:
|
||||
outcome = sync.compute_sync({}, RECORDED_CATALOG, RECORDED_DOC)
|
||||
gemma = outcome.cost_map["together_ai/google/gemma-3n-E4B-it"]
|
||||
assert "deprecation_date" not in gemma
|
||||
assert any("gemma-3n-E4B-it" in warning and "2026-08-04" in warning for warning in outcome.warnings)
|
||||
|
||||
|
||||
def test_price_change_updates_api_fields_and_keeps_curated_ones() -> None:
|
||||
registry = {
|
||||
"together_ai/acme/chat-1": {
|
||||
"input_cost_per_token": 9e-07,
|
||||
"litellm_provider": "together_ai",
|
||||
"max_input_tokens": 4096,
|
||||
"max_output_tokens": 2048,
|
||||
"max_tokens": 2048,
|
||||
"mode": "chat",
|
||||
"output_cost_per_token": 9e-07,
|
||||
"supports_audio_input": True,
|
||||
}
|
||||
}
|
||||
outcome = sync.compute_sync(registry, [_chat_model("acme/chat-1", ctx=8192, price=2.0)], _doc({"x": "2026-01-01"}))
|
||||
entry = outcome.cost_map["together_ai/acme/chat-1"]
|
||||
assert entry["input_cost_per_token"] == 2e-06
|
||||
assert entry["max_input_tokens"] == 8192
|
||||
assert entry["max_output_tokens"] == 2048
|
||||
assert entry["supports_audio_input"] is True
|
||||
assert len(outcome.updated) == 1
|
||||
assert "input_cost_per_token" in outcome.updated[0]
|
||||
|
||||
|
||||
def test_cached_input_appearing_and_disappearing() -> None:
|
||||
registry = {
|
||||
"together_ai/acme/chat-1": {
|
||||
"cache_read_input_token_cost": 1e-07,
|
||||
"input_cost_per_token": 1e-06,
|
||||
"litellm_provider": "together_ai",
|
||||
"max_input_tokens": 4096,
|
||||
"mode": "chat",
|
||||
"output_cost_per_token": 1e-06,
|
||||
"supports_prompt_caching": True,
|
||||
},
|
||||
"together_ai/acme/chat-2": {
|
||||
"input_cost_per_token": 1e-06,
|
||||
"litellm_provider": "together_ai",
|
||||
"max_input_tokens": 4096,
|
||||
"mode": "chat",
|
||||
"output_cost_per_token": 1e-06,
|
||||
},
|
||||
}
|
||||
catalog = [_chat_model("acme/chat-1"), _chat_model("acme/chat-2", cached=0.25999999999999995)]
|
||||
outcome = sync.compute_sync(registry, catalog, _doc({"x": "2026-01-01"}))
|
||||
assert "cache_read_input_token_cost" not in outcome.cost_map["together_ai/acme/chat-1"]
|
||||
assert "supports_prompt_caching" not in outcome.cost_map["together_ai/acme/chat-1"]
|
||||
assert outcome.cost_map["together_ai/acme/chat-2"]["cache_read_input_token_cost"] == 2.6e-07
|
||||
assert outcome.cost_map["together_ai/acme/chat-2"]["supports_prompt_caching"] is True
|
||||
|
||||
|
||||
def test_capability_rule_backfills_existing_entry() -> None:
|
||||
registry = {
|
||||
"together_ai/moonshotai/Kimi-K3": {
|
||||
"input_cost_per_token": 3e-06,
|
||||
"litellm_provider": "together_ai",
|
||||
"max_input_tokens": 1048576,
|
||||
"mode": "chat",
|
||||
"output_cost_per_token": 1.5e-05,
|
||||
}
|
||||
}
|
||||
kimi = next(model for model in RECORDED_CATALOG if model.id == "moonshotai/Kimi-K3")
|
||||
outcome = sync.compute_sync(registry, [kimi], _doc({"x": "2026-01-01"}))
|
||||
assert outcome.cost_map["together_ai/moonshotai/Kimi-K3"]["supports_reasoning"] is True
|
||||
assert any("supports_reasoning" in line for line in outcome.updated)
|
||||
|
||||
|
||||
def test_disappeared_model_gets_docs_date_and_is_never_deleted() -> None:
|
||||
registry = {
|
||||
"together_ai/acme/gone": {
|
||||
"input_cost_per_token": 1e-06,
|
||||
"litellm_provider": "together_ai",
|
||||
"mode": "chat",
|
||||
"output_cost_per_token": 1e-06,
|
||||
}
|
||||
}
|
||||
outcome = sync.compute_sync(registry, [_chat_model("acme/alive")], _doc({"acme/gone": "2026-07-01"}))
|
||||
assert outcome.cost_map["together_ai/acme/gone"]["deprecation_date"] == "2026-07-01"
|
||||
assert outcome.deprecated == ("together_ai/acme/gone: deprecation_date",)
|
||||
|
||||
|
||||
def test_disappeared_model_without_docs_date_warns_instead() -> None:
|
||||
registry = {
|
||||
"together_ai/acme/gone": {
|
||||
"input_cost_per_token": 1e-06,
|
||||
"litellm_provider": "together_ai",
|
||||
"mode": "chat",
|
||||
"output_cost_per_token": 1e-06,
|
||||
}
|
||||
}
|
||||
outcome = sync.compute_sync(registry, [_chat_model("acme/alive")], _doc({"other": "2026-07-01"}))
|
||||
assert "deprecation_date" not in outcome.cost_map["together_ai/acme/gone"]
|
||||
assert not outcome.deprecated
|
||||
assert any("acme/gone" in warning and "human" in warning for warning in outcome.warnings)
|
||||
|
||||
|
||||
def test_curated_deprecation_date_is_never_overwritten() -> None:
|
||||
registry = {
|
||||
"together_ai/acme/gone": {
|
||||
"deprecation_date": "2026-06-15",
|
||||
"input_cost_per_token": 1e-06,
|
||||
"litellm_provider": "together_ai",
|
||||
"mode": "chat",
|
||||
"output_cost_per_token": 1e-06,
|
||||
}
|
||||
}
|
||||
outcome = sync.compute_sync(registry, [_chat_model("acme/alive")], _doc({"acme/gone": "2026-07-01"}))
|
||||
assert outcome.cost_map["together_ai/acme/gone"]["deprecation_date"] == "2026-06-15"
|
||||
assert any("2026-06-15" in warning and "2026-07-01" in warning for warning in outcome.warnings)
|
||||
|
||||
|
||||
def test_redirect_chain_resolves_to_final_live_model() -> None:
|
||||
doc = _doc({"acme/a": "2026-01-01"}, redirects={"acme/a": "acme/b", "acme/b": "acme/c"})
|
||||
live = frozenset({"acme/c"})
|
||||
assert sync.resolve_successor("acme/a", doc, live) == "acme/c"
|
||||
|
||||
|
||||
def test_redirect_dead_end_yields_no_successor() -> None:
|
||||
doc = _doc({"acme/a": "2026-01-01"}, redirects={"acme/a": "acme/b"})
|
||||
assert sync.resolve_successor("acme/a", doc, frozenset({"acme/other"})) is None
|
||||
|
||||
|
||||
def test_redirect_short_names_resolve_by_unique_suffix() -> None:
|
||||
doc = _doc({"moonshotai/Kimi-K2": "2026-01-01"}, redirects={"Kimi-K2": "Kimi-K2-0905"})
|
||||
live = frozenset({"moonshotai/Kimi-K2-0905"})
|
||||
assert sync.resolve_successor("moonshotai/Kimi-K2", doc, live) == "moonshotai/Kimi-K2-0905"
|
||||
|
||||
|
||||
def test_successor_written_only_when_not_curated() -> None:
|
||||
registry = {
|
||||
"together_ai/acme/a": {
|
||||
"input_cost_per_token": 1e-06,
|
||||
"litellm_provider": "together_ai",
|
||||
"mode": "chat",
|
||||
"output_cost_per_token": 1e-06,
|
||||
},
|
||||
"together_ai/acme/b": {
|
||||
"input_cost_per_token": 1e-06,
|
||||
"litellm_provider": "together_ai",
|
||||
"metadata": {"successor": "together_ai/acme/curated"},
|
||||
"mode": "chat",
|
||||
"output_cost_per_token": 1e-06,
|
||||
},
|
||||
}
|
||||
doc = _doc({"acme/a": "2026-01-01", "acme/b": "2026-01-01"}, redirects={"acme/a": "acme/c", "acme/b": "acme/c"})
|
||||
outcome = sync.compute_sync(registry, [_chat_model("acme/c")], doc)
|
||||
assert outcome.cost_map["together_ai/acme/a"]["metadata"] == {"successor": "together_ai/acme/c"}
|
||||
assert outcome.cost_map["together_ai/acme/b"]["metadata"] == {"successor": "together_ai/acme/curated"}
|
||||
assert any("acme/curated" in warning for warning in outcome.warnings)
|
||||
|
||||
|
||||
def test_reappearance_clears_deprecation_date() -> None:
|
||||
registry = {
|
||||
"together_ai/acme/back": {
|
||||
"deprecation_date": "2026-05-01",
|
||||
"input_cost_per_token": 1e-06,
|
||||
"litellm_provider": "together_ai",
|
||||
"max_input_tokens": 4096,
|
||||
"mode": "chat",
|
||||
"output_cost_per_token": 1e-06,
|
||||
}
|
||||
}
|
||||
outcome = sync.compute_sync(registry, [_chat_model("acme/back")], _doc({"x": "2026-01-01"}))
|
||||
assert "deprecation_date" not in outcome.cost_map["together_ai/acme/back"]
|
||||
assert outcome.reappeared == ("together_ai/acme/back",)
|
||||
|
||||
|
||||
def test_new_chat_model_without_rule_is_flagged() -> None:
|
||||
outcome = sync.compute_sync({}, [_chat_model("acme/unreviewed")], _doc({"x": "2026-01-01"}))
|
||||
assert any("acme/unreviewed" in warning and "capability rule" in warning for warning in outcome.warnings)
|
||||
|
||||
|
||||
def test_new_keys_land_at_the_end_of_the_provider_block() -> None:
|
||||
registry = {
|
||||
"aaa": {"mode": "chat"},
|
||||
"together_ai/acme/old": {
|
||||
"input_cost_per_token": 1e-06,
|
||||
"litellm_provider": "together_ai",
|
||||
"mode": "chat",
|
||||
"output_cost_per_token": 1e-06,
|
||||
},
|
||||
"zzz": {"mode": "chat"},
|
||||
}
|
||||
outcome = sync.compute_sync(registry, [_chat_model("acme/old"), _chat_model("acme/new")], _doc({"x": "2026-01-01"}))
|
||||
assert list(outcome.cost_map) == ["aaa", "together_ai/acme/old", "together_ai/acme/new", "zzz"]
|
||||
|
||||
|
||||
def test_sync_is_idempotent_over_the_repo_cost_map() -> None:
|
||||
cost_map = json.loads((ROOT / "model_prices_and_context_window.json").read_text())
|
||||
first = sync.compute_sync(cost_map, RECORDED_CATALOG, RECORDED_DOC)
|
||||
second = sync.compute_sync(first.cost_map, RECORDED_CATALOG, RECORDED_DOC)
|
||||
assert not second.has_changes
|
||||
assert second.cost_map == first.cost_map
|
||||
|
||||
|
||||
def test_pr_body_lists_every_section_and_the_skipped_types() -> None:
|
||||
outcome = sync.compute_sync({}, RECORDED_CATALOG, RECORDED_DOC)
|
||||
body = sync.render_pr_body(outcome)
|
||||
assert "### Added (26)" in body
|
||||
assert "### Warnings needing a human call" in body
|
||||
assert "image (29)" in body
|
||||
assert "video (38)" in body
|
||||
|
|
@ -4,7 +4,7 @@ Never put LiteLLM tokens or API keys in `localStorage`. `localStorage` survives
|
|||
|
||||
When you fix lint violations that are grandfathered in `eslint-suppressions.json`, run `eslint . --prune-suppressions` and commit the updated baseline so the gate ratchets down instead of leaving a stale suppression
|
||||
|
||||
`src/lib/http/schema.d.ts` is generated from the proxy's OpenAPI spec. Never hand-edit it. After changing a backend route or response model that the dashboard consumes, run `npm run gen:api` and commit the result (CI `lint / ui-api-types` enforces this)
|
||||
`src/lib/http/schema.d.ts` is generated from the proxy's OpenAPI spec. Never hand-edit it. After changing any backend route or request or response model, or bumping a dependency, run `npm run gen:api` and commit the result, since the spec includes every route, not only the ones the dashboard calls (CI `lint / ui-api-types` enforces this)
|
||||
|
||||
Tests come in three tiers, named by the standard definitions. `Foo.test.tsx` is a unit test: one module, collaborators replaced by doubles, no multi-component tree, and it should run in milliseconds. `Foo.integration.test.tsx` renders a real component tree with real children and only stubs the network boundary. It costs seconds per case, so it earns its place by proving wiring that a unit test cannot reach. Browser-level tests live in `tests/e2e/ui/` as Playwright specs against a live proxy
|
||||
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue