Merge remote-tracking branch 'origin/main' into litellm_aws_rotation_values

This commit is contained in:
Yuneng Jiang 2026-09-16 17:29:10 -07:00
commit 495b731bbc
No known key found for this signature in database
46 changed files with 2662 additions and 122 deletions

View file

@ -2983,7 +2983,7 @@ workflows:
name: integration-<< matrix.suite >>
matrix:
parameters:
suite: [management, accounting, providers]
suite: [management, accounting, database, providers]
filters:
branches:
only:

View file

@ -502,7 +502,12 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge):
elif key == "response_format":
text_format = self._transform_response_format_to_text_format(value)
if text_format:
responses_api_request["text"] = text_format
responses_api_request["text"] = self._merge_text(responses_api_request, text_format)
elif key == "verbosity":
responses_api_request["text"] = self._merge_text(
responses_api_request,
MappingProxyType({"verbosity": value}), # pyright: ignore[reportUnknownArgumentType] # untyped value
)
elif key == "tool_choice":
responses_api_request["tool_choice"] = self._normalize_tool_choice_for_responses_api(value)
elif key == "stream_options":
@ -518,6 +523,19 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge):
elif key == "web_search_options":
self._add_web_search_tool(responses_api_request, value)
@staticmethod
def _merge_text(
responses_api_request: "ResponsesAPIOptionalRequestParams", update: Mapping[str, object]
) -> "ResponseText":
existing: Final = cast( # cast-ok: text field is a ResponseText | dict[str, Any] | None union
"dict[str, object]",
dict(responses_api_request).get("text") or {}, # mutable-ok: one-shot merge seed
)
return cast( # cast-ok: merged mapping is a valid ResponseText shape
"ResponseText",
{**existing, **update}, # mutable-ok: one-shot merged payload
)
def _build_sanitized_litellm_params(self, litellm_params: dict) -> dict[str, object]:
"""Build sanitized litellm_params with merged metadata."""
responses_optional_param_keys: Final = set(ResponsesAPIOptionalRequestParams.__annotations__.keys())

View file

@ -22,6 +22,7 @@ from litellm.llms.bedrock_mantle.common_utils import (
BEDROCK_MANTLE_DEFAULT_REGION,
BedrockMantleAuthMixin,
)
from litellm.llms.openai.chat.gpt_5_transformation import is_gpt_reasoning_series_name
from litellm.secret_managers.main import get_secret_str
from litellm.types.llms.openai import AllMessageValues
from litellm.types.router import GenericLiteLLMParams
@ -108,13 +109,22 @@ class BedrockMantleChatConfig(BedrockMantleAuthMixin, OpenAILikeChatConfig):
def get_supported_openai_params(self, model: str) -> list:
base_params: Final = super().get_supported_openai_params(model)
extra_params: Final = tuple(
param
for param, supported in (
("verbosity", is_gpt_reasoning_series_name(model)),
("reasoning_effort", self._supports_reasoning(model)),
)
if supported and param not in base_params
)
return [*base_params, *extra_params] # mutable-ok: fresh list required by the inherited signature
def _supports_reasoning(self, model: str) -> bool:
try:
if litellm.supports_reasoning(model=model, custom_llm_provider=self.custom_llm_provider):
if "reasoning_effort" not in base_params:
base_params.append("reasoning_effort")
return litellm.supports_reasoning(model=model, custom_llm_provider=self.custom_llm_provider)
except Exception as e:
verbose_logger.debug("BedrockMantleChatConfig: error checking reasoning support: %s", e)
return base_params
return False
def get_model_response_iterator(
self,

View file

@ -512,7 +512,7 @@ def create_tool_function(
)
from litellm.proxy._experimental.mcp_server.outbound_credentials.result import Error, Ok
match validate_static_credential(auth_type, effective_headers, upstream_token_header):
match validate_static_credential(auth_type, effective_headers, upstream_token_header, headers or ()):
case Error(error):
raise_public(error)
case Ok():

View file

@ -13,7 +13,7 @@ from __future__ import annotations
import base64
import os
from collections.abc import Mapping
from collections.abc import Iterable, Mapping
from typing import TYPE_CHECKING, Final, Literal, NoReturn
from fastapi import HTTPException
@ -426,16 +426,19 @@ def validate_static_credential(
auth_type: MCPAuthType,
headers: Mapping[str, str],
upstream_token_header: str | None = None,
static_header_names: Iterable[str] = (),
) -> Result[None, CredError]:
if auth_type not in _STATIC_MODES:
return Ok(None)
default_slot: Final = "X-API-Key" if auth_type == MCPAuth.api_key else "Authorization"
admin_chosen_slots: Final = tuple(static_header_names) if auth_type == MCPAuth.api_key else ()
slots: Final = frozenset(
name.lower()
for name in (
upstream_token_header or default_slot,
default_slot,
"Authorization",
*admin_chosen_slots,
)
)
values: Final = tuple((name.lower(), value.strip()) for name, value in headers.items() if name.lower() in slots)
@ -448,7 +451,9 @@ async def prepare_mcp_client(server: MCPServer, client: MCPClient) -> MCPClient:
if server.auth_type not in _STATIC_MODES or client.transport_type == MCPTransport.stdio:
return client
request: Final = await client.prepare_request_auth()
match validate_static_credential(server.auth_type, request.headers, server.upstream_token_header):
match validate_static_credential(
server.auth_type, request.headers, server.upstream_token_header, server.static_headers or ()
):
case Error(error):
raise_public(error)
case Ok():

View file

@ -2,7 +2,7 @@
These tests exercise a running gateway, PostgreSQL and Redis with an owned local upstream. CircleCI owns this suite. Tests are grouped by behavior, with no automatic test retries or fallback to paid provider calls
Use `tests/integration/run.py management`, `accounting` or `providers` to run a selected group. Set `INTEGRATION_PROXY_URL`, `INTEGRATION_UPSTREAM_URL`, `INTEGRATION_MASTER_KEY` and `DATABASE_URL` to an isolated test deployment. The runner selects the new domain directories explicitly; the legacy OCI and sandbox selections remain separate
Use `tests/integration/run.py management`, `accounting`, `database` or `providers` to run a selected group. Set `INTEGRATION_PROXY_URL`, `INTEGRATION_UPSTREAM_URL`, `INTEGRATION_MASTER_KEY` and `DATABASE_URL` to an isolated test deployment. The runner selects the new domain directories explicitly; the legacy OCI and sandbox selections remain separate
Management also requires `INTEGRATION_PEER_URL`, `REDIS_HOST` and `REDIS_PORT`. CircleCI starts two directly addressed proxy processes sharing only that job's stores. The test-only CLI wrapper supplies enterprise route entitlement, following the existing behavior suite's convention. It does not qualify license validation; run it with one worker and no reload
@ -17,3 +17,11 @@ Define integration contract IDs and their canonical test nodes in `contracts.jso
Provider sentinels currently use the controlled server, not live recordings. The provider shard also runs the existing strict replay controls for changed requests, exhausted interactions, leftover interactions and no provider connection. Future recorded scenarios must use that replay-only implementation; missing recordings cannot fall back to a real provider. The observation endpoint is destructive and the current selection runs serially against one owned upstream
Fixtures must contain synthetic data only. Keep private incident records and source documents out of code, fixtures, logs and PR descriptions
Database cases own their temporary schemas, roles, constraints and proxy processes. They prove reader-versus-writer execution with PostgreSQL lock observations, exercise real transaction wait limits and verify rollback after a reached database failure
Accounting cases compare persisted input and output cost components against literal rates, including zero and default prices. Cache state models assert actual upstream calls, response identity and every persisted charge. Generated accounting tests have a 180-second test limit to accommodate the asynchronous spend writer; CircleCI keeps the whole shard capped at 11 minutes
Provider contracts exercise actual TCP requests with synthetic credentials and local protocol peers. The S3 verifier uses independently implemented equations, a published known-answer vector, a fixed signing clock and deliberately invalid signed requests. Bedrock cases clear ambient AWS credential sources and check the literal model path, loaded role references, STS requests and bearer-only behavior
Streaming checks send real HTTP transfer chunks, including one-byte partitions, fragmented tools, incomplete transfers and a cancellation barrier. They assert meaningful text, tool arguments, final usage and persisted cost. The Redis recovery case owns a separate database and Redis process, uses the supported one-second circuit-breaker recovery setting, waits for the real subscriber and verifies response data in Redis after restart. CircleCI reuses its existing Redis image for that extra process; it never pulls an image during tests

View file

@ -0,0 +1,114 @@
import os
import socket
import signal
import subprocess
import sys
import time
import uuid
from collections.abc import Iterator, Mapping
from contextlib import contextmanager
from pathlib import Path
from typing import Final
import httpx
import psutil
from integration._support.client import Gateway
def in_group(process: psutil.Process, group: int) -> bool:
try:
return os.getpgid(process.pid) == group
except ProcessLookupError:
return False
def group_members(group: int) -> tuple[psutil.Process, ...]:
return tuple(process for process in psutil.process_iter() if in_group(process, group))
def signal_group(group: int, action: int) -> None:
try:
os.killpg(group, action)
except ProcessLookupError:
pass
def stop_root_process(process: subprocess.Popen[bytes]) -> bool:
if process.poll() is not None:
return True
process.terminate()
try:
process.wait(timeout=30)
except subprocess.TimeoutExpired:
return False
return True
@contextmanager
def owned_proxy(gateway: Gateway, directory: Path, overrides: Mapping[str, str], *, config: Path | None = None, remove_environment: tuple[str, ...] = ()) -> Iterator[Gateway]:
with socket.socket() as reserve:
reserve.bind(("127.0.0.1", 0))
port: Final = reserve.getsockname()[1]
root: Final = Path(__file__).resolve().parents[3]
environment: Final = {
**{name: value for name, value in os.environ.items() if name not in remove_environment},
"LITELLM_MASTER_KEY": gateway.key,
"LITELLM_SALT_KEY": os.environ.get("LITELLM_SALT_KEY", "sk-integration-salt"),
"STORE_MODEL_IN_DB": "True",
**overrides,
}
output: Final = Path(os.environ.get("INTEGRATION_RESULTS_DIR", str(directory)))
output.mkdir(parents=True, exist_ok=True)
with (output / f"owned-proxy-{uuid.uuid4().hex}.log").open("w") as log:
process: Final = subprocess.Popen(
[
sys.executable,
"-m",
"integration._support.proxy",
"--config",
str(config or "tests/integration/proxy_config.yaml"),
"--host",
"127.0.0.1",
"--port",
str(port),
"--num_workers",
"1",
"--telemetry",
"False",
"--use_prisma_db_push",
"--enforce_prisma_migration_check",
],
cwd=root,
env=environment,
stdout=log,
stderr=subprocess.STDOUT,
start_new_session=True,
)
try:
with httpx.Client(base_url=f"http://127.0.0.1:{port}", timeout=15, trust_env=False) as client:
deadline: Final = time.monotonic() + 70
while True:
assert process.poll() is None, "Owned proxy exited before readiness"
try:
if client.get("/health/readiness", timeout=2).status_code == 200:
break
except httpx.TransportError:
pass
assert time.monotonic() < deadline, "Owned proxy readiness deadline exceeded"
time.sleep(0.1)
yield Gateway(client, gateway.key, gateway.upstream_url)
finally:
root_stopped: Final = stop_root_process(process)
residual: Final = group_members(process.pid)
if residual:
signal_group(process.pid, signal.SIGTERM)
psutil.wait_procs(residual, timeout=5)
remaining: Final = group_members(process.pid)
if remaining:
signal_group(process.pid, signal.SIGKILL)
psutil.wait_procs(remaining, timeout=3)
process.wait(timeout=3)
survivors: Final = group_members(process.pid)
assert not survivors, "Owned proxy child survived cleanup"
assert root_stopped and not remaining, "Owned proxy required forced cleanup"

View file

@ -0,0 +1,116 @@
import os
import shutil
import signal
import socket
import subprocess
import time
import uuid
from collections.abc import Iterator
from contextlib import contextmanager
from dataclasses import dataclass
from pathlib import Path
from typing import Final, TextIO
from redis import Redis
from redis.exceptions import ConnectionError as RedisConnectionError
@dataclass
class OwnedRedis:
host: str
port: int
command: tuple[str, ...]
log: TextIO
pid_file: str
process: subprocess.Popen | None = None
server_pid: int | None = None
def start(self) -> None:
assert self.process is None
self.process = subprocess.Popen(self.command, stdout=self.log, stderr=subprocess.STDOUT, start_new_session=True)
deadline: Final = time.monotonic() + 8
with Redis(host=self.host, port=self.port, socket_connect_timeout=0.2, socket_timeout=0.2) as client:
while True:
assert self.process.poll() is None, "Owned Redis exited before readiness"
try:
if client.ping():
actual: Final = int(client.info("server")["process_id"])
expected: Final = self.process.pid if self.command[0] != "docker" else int(subprocess.check_output(["docker", "exec", "redis-cache", "cat", self.pid_file], timeout=2))
assert actual == expected, "Redis readiness reached a different process"
self.server_pid = actual
return
except RedisConnectionError:
pass
assert time.monotonic() < deadline, "Owned Redis readiness deadline exceeded"
time.sleep(0.05)
def stop(self) -> None:
assert self.process is not None
failure = None
forced = False
try:
if self.process.poll() is None:
with Redis(host=self.host, port=self.port, socket_connect_timeout=1, socket_timeout=1) as client:
assert int(client.info("server")["process_id"]) == self.server_pid, "Redis ownership changed before shutdown"
client.shutdown(nosave=True)
except Exception as error:
failure = error
finally:
try:
self.process.wait(timeout=5)
except subprocess.TimeoutExpired:
forced = True
self.signal(signal.SIGTERM)
try:
self.process.wait(timeout=5)
except subprocess.TimeoutExpired:
self.signal(signal.SIGKILL)
self.process.wait(timeout=3)
self.process = None
self.server_pid = None
with Redis(host=self.host, port=self.port, socket_connect_timeout=0.2, socket_timeout=0.2) as client:
try:
client.ping()
except RedisConnectionError:
stopped = True
else:
stopped = False
assert stopped, "Owned Redis still serves after shutdown"
assert failure is None and not forced, f"Owned Redis required shutdown recovery: {failure!r}"
def signal(self, action: signal.Signals) -> None:
assert self.process is not None
if self.command[0] != "docker":
self.process.send_signal(action)
return
pid: Final = int(subprocess.check_output(["docker", "exec", "redis-cache", "cat", self.pid_file], timeout=2))
command: Final = subprocess.check_output(["docker", "exec", "redis-cache", "cat", f"/proc/{pid}/cmdline"], timeout=2)
assert self.pid_file.encode() in command, "Redis process ownership changed"
subprocess.run(["docker", "exec", "redis-cache", "kill", f"-{int(action)}", str(pid)], check=True, timeout=2)
@contextmanager
def owned_redis(directory: Path) -> Iterator[OwnedRedis]:
binary: Final = shutil.which("redis-server")
if binary:
with socket.socket() as reservation:
reservation.bind(("127.0.0.1", 0))
port = reservation.getsockname()[1]
host = "127.0.0.1"
prefix = (binary,)
else:
host = subprocess.check_output(["docker", "inspect", "--format", "{{range .NetworkSettings.Networks}}{{.IPAddress}}{{end}}", "redis-cache"], text=True).strip()
assert host, "CircleCI owned Redis container has no address"
port = 16379
prefix = ("docker", "exec", "redis-cache", "redis-server")
output: Final = Path(os.environ.get("INTEGRATION_RESULTS_DIR", str(directory)))
output.mkdir(parents=True, exist_ok=True)
with (output / "owned-redis-recovery.log").open("w") as log:
pid_file: Final = str(directory / "owned-redis.pid") if binary else f"/tmp/integration-redis-{uuid.uuid4().hex}.pid"
server: Final = OwnedRedis(host, port, (*prefix, "--port", str(port), "--set-proc-title", "no", "--pidfile", pid_file, "--bind", "0.0.0.0" if not binary else "127.0.0.1", "--protected-mode", "no", "--save", "", "--appendonly", "no"), log, pid_file)
try:
server.start()
yield server
finally:
if server.process is not None:
server.stop()

View file

@ -0,0 +1,25 @@
import hashlib
import hmac
from collections.abc import Mapping
from typing import Final
def encoded_path(value: str) -> str:
safe: Final = b"ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789-_.~/"
return "".join(chr(byte) if byte in safe else f"%{byte:02X}" for byte in value.encode("utf-8"))
def signature(
method: str, path: str, headers: Mapping[str, str], signed: str, body: bytes, secret: str, scope: str,
) -> tuple[str, str]:
"""AWS SigV4 equations, independent of botocore and LiteLLM's signer."""
canonical_headers: Final = "".join(name + ":" + " ".join(headers[name].split()) + "\n" for name in signed.split(";"))
canonical: Final = "\n".join((method, path, "", canonical_headers, signed, hashlib.sha256(body).hexdigest()))
canonical_hash: Final = hashlib.sha256(canonical.encode()).hexdigest()
date, region, service, terminator = scope.split("/")
assert terminator == "aws4_request"
key = ("AWS4" + secret).encode()
for part in (date, region, service, terminator):
key = hmac.new(key, part.encode(), hashlib.sha256).digest()
to_sign: Final = "\n".join(("AWS4-HMAC-SHA256", headers["x-amz-date"], scope, canonical_hash))
return canonical_hash, hmac.new(key, to_sign.encode(), hashlib.sha256).hexdigest()

View file

@ -31,6 +31,12 @@ INTERNAL_FIELDS: Final = frozenset(
)
def error_type(status: int) -> str:
if status == 429:
return "rate_limit_error"
return "invalid_request_error" if status < 500 else "server_error"
@dataclass(frozen=True, slots=True)
class Observation:
path: str
@ -66,7 +72,7 @@ class Provider:
status: Final = script.popleft()
if status != 200:
return JSONResponse(
{"error": {"message": "Controlled provider failure", "type": "api_error", "code": str(status)}},
{"error": {"message": "Controlled provider failure", "type": error_type(status), "code": str(status)}},
status_code=status,
)
return await chat_completions(request)

View file

@ -0,0 +1,113 @@
from __future__ import annotations
import threading
from collections.abc import Callable, Iterator, Mapping
from contextlib import contextmanager
from dataclasses import dataclass
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
from queue import SimpleQueue
from typing import Final
@dataclass(frozen=True, slots=True)
class Request:
method: str
target: str
headers: Mapping[str, str]
body: bytes
@dataclass(frozen=True, slots=True)
class Reply:
status: int = 200
body: bytes = b"{}"
content_type: str = "application/json"
chunks: tuple[bytes, ...] | None = None
abort_after: int | None = None
gate_after_first: threading.Event | None = None
@dataclass(frozen=True, slots=True)
class Wire:
url: str
received: SimpleQueue[Request]
disconnected: SimpleQueue[str]
def drain(self) -> tuple[Request, ...]:
return tuple(self.received.get_nowait() for _ in range(self.received.qsize()))
@contextmanager
def wire_server(respond: Callable[[Request], Reply]) -> Iterator[Wire]:
"""Owned TCP peer; requests traverse the real HTTP client and serialization."""
received: Final[SimpleQueue[Request]] = SimpleQueue()
errors: Final[SimpleQueue[Exception]] = SimpleQueue()
disconnected: Final[SimpleQueue[str]] = SimpleQueue()
class Handler(BaseHTTPRequestHandler):
protocol_version = "HTTP/1.1"
timeout = 5
def respond(self) -> None:
request: Final = Request(
self.command, self.path,
{name.lower(): value for name, value in self.headers.items()},
self.rfile.read(int(self.headers.get("content-length", "0"))),
)
received.put(request)
try:
reply = respond(request)
except Exception as error:
errors.put(error)
reply = Reply(status=500)
self.send_response(reply.status)
self.send_header("content-type", reply.content_type)
if reply.chunks is None:
self.send_header("content-length", str(len(reply.body)))
else:
self.send_header("transfer-encoding", "chunked")
self.send_header("connection", "close")
self.end_headers()
try:
if reply.chunks is None:
self.wfile.write(reply.body)
else:
for index, chunk in enumerate(reply.chunks):
if reply.abort_after == index:
break
self.wfile.write(b"%x\r\n%s\r\n" % (len(chunk), chunk))
self.wfile.flush()
if index == 0 and reply.gate_after_first is not None:
assert reply.gate_after_first.wait(timeout=5), "Stream barrier was never released"
else:
self.wfile.write(b"0\r\n\r\n")
self.wfile.flush()
except (BrokenPipeError, ConnectionResetError):
disconnected.put(request.target)
except Exception as error:
errors.put(error)
self.close_connection = True
do_POST = respond
do_PUT = respond
do_GET = respond
do_DELETE = respond
def log_message(self, format: str, *args: object) -> None:
pass
class OwnedHTTPServer(ThreadingHTTPServer):
daemon_threads = False
with OwnedHTTPServer(("127.0.0.1", 0), Handler) as server:
thread: Final = threading.Thread(target=server.serve_forever, kwargs={"poll_interval": 0.05})
thread.start()
try:
yield Wire(f"http://127.0.0.1:{server.server_port}", received, disconnected)
finally:
server.shutdown()
thread.join(timeout=6)
assert not thread.is_alive(), "Owned HTTP server survived cleanup"
server.server_close()
failure: Final = None if errors.empty() else errors.get_nowait()
assert failure is None, f"Owned HTTP peer failed: {failure!r}"

View file

@ -75,6 +75,81 @@
],
"tests/integration/authorization/test_warmed_policy.py::test_expiry_and_explicit_clear_reach_both_warmed_workers": [
"mgmt.key.update.expiry_changes_reach_warmed_workers"
],
"tests/integration/database/test_partition_transactions.py::test_real_partition_ddl_survives_witnessed_lock_and_is_idempotent": [
"other.database.partitions.lock_wait_outlives_transaction_default",
"other.database.partitions.repeat_preserves_rows"
],
"tests/integration/database/test_reader_writer_regeneration.py::test_key_regeneration_uses_writer_with_a_real_readonly_reader": [
"other.database.regeneration.writer_updates_dependent_grants"
],
"tests/integration/pricing/test_price_precedence.py::test_generated_zero_null_and_omitted_prices_follow_independent_arithmetic": [
"quota_management.spend_tracking.price_precedence.zero_and_default_rates"
],
"tests/integration/pricing/test_price_precedence.py::test_same_upstream_aliases_keep_distinct_prices_after_reload": [
"quota_management.spend_tracking.alias_prices.remain_independent_on_reload"
],
"tests/integration/spend/test_cache_and_quota.py::test_generated_cache_sequences_preserve_content_usage_and_zero_hit_cost": [
"quota_management.response_cache.generated_sequences_preserve_content_and_accounting"
],
"tests/integration/spend/test_cache_and_quota.py::test_key_budget_at_boundary_blocks_provider_then_explicit_reset_restores": [
"quota_management.budget.key.boundary_blocks_before_provider_and_reset_restores"
],
"tests/integration/spend/test_cache_and_quota.py::test_different_system_messages_do_not_share_a_cached_response": [
"quota_management.response_cache.system_messages_partition_cache_identity"
],
"tests/integration/database/test_transaction_atomicity.py::test_access_group_second_key_constraint_failure_rolls_back_all_writes": [
"other.database.access_group.failed_second_write_rolls_back_first"
],
"tests/integration/spend/test_cache_and_quota.py::test_repeated_hits_keep_response_identity_and_create_distinct_zero_cost_rows": [
"quota_management.response_cache.repeated_hits_preserve_identity_and_single_charge"
],
"tests/integration/providers/test_s3_wire.py::test_sigv4_verifier_matches_published_put_and_rejects_corruption": [
"other.provider_wire.s3.verifier_known_answer_and_negative_controls"
],
"tests/integration/providers/test_s3_wire.py::test_s3_sync_and_async_uploads_pass_independent_wire_verification": [
"other.provider_wire.s3.sync_async_reserved_keys_are_signed_and_accepted"
],
"tests/integration/providers/test_bedrock_auth_wire.py::test_bearer_only_sdk_sync_async_requests_do_not_require_aws_credentials": [
"other.provider_wire.bedrock.bearer_sdk_skips_credential_chain"
],
"tests/integration/providers/test_bedrock_auth_wire.py::test_bearer_environment_reference_loads_from_db_and_yaml_and_survives_reload": [
"other.provider_wire.bedrock.bearer_db_yaml_survives_reload"
],
"tests/integration/streaming/test_stream_contracts.py::test_generated_tcp_partitions_preserve_unicode_text_identity_and_final_usage": [
"other.streaming.byte_partitions.preserve_text_identity_and_usage"
],
"tests/integration/streaming/test_stream_contracts.py::test_fragmented_tool_names_and_arguments_keep_each_call_identity": [
"other.streaming.tools.fragmented_calls_keep_independent_arguments"
],
"tests/integration/streaming/test_stream_contracts.py::test_proxy_stream_usage_visibility_keeps_exact_persisted_charge": [
"other.streaming.usage.client_visibility_preserves_persisted_accounting"
],
"tests/integration/streaming/test_stream_contracts.py::test_truncated_http_stream_is_an_error_and_next_stream_succeeds": [
"other.streaming.failure.truncated_transport_raises_and_control_recovers"
],
"tests/integration/streaming/test_stream_contracts.py::test_client_cancellation_releases_the_actual_provider_connection": [
"other.streaming.cancellation.closes_actual_provider_connection"
],
"tests/integration/routing/test_observed_routing.py::test_retry_counts_and_public_errors_match_actual_provider_attempts": [
"other.routing.retries.several_attempts_reach_success_without_hidden_retries",
"other.routing.errors.nonretryable_and_exhausted_failures_remain_errors"
],
"tests/integration/routing/test_observed_routing.py::test_loaded_fallback_selects_expected_deployment_and_keeps_response_identity": [
"other.routing.fallback.loaded_configuration_selects_only_permitted_target"
],
"tests/integration/routing/test_observed_routing.py::test_saved_deployment_target_update_changes_wire_and_preserves_control": [
"other.routing.alias_update.persisted_target_changes_only_selected_route"
],
"tests/integration/providers/test_bedrock_role_configuration.py::test_role_reference_from_db_and_yaml_reaches_real_sts_http_and_bedrock": [
"other.provider_wire.bedrock.db_yaml_role_reference_reaches_sts_and_signed_request"
],
"tests/integration/routing/test_redis_recovery.py::test_owned_redis_outage_recovers_requests_and_real_response_cache": [
"other.routing.redis.owned_outage_recovers_serving_and_response_cache"
],
"tests/integration/providers/test_anthropic_wire.py::test_anthropic_tool_history_and_cache_tokens_keep_wire_and_accounting_contracts": [
"other.provider_wire.anthropic.tool_history_system_cache_and_internal_fields",
"quota_management.spend_tracking.cache_tokens.disjoint_classes_use_explicit_rates"
]
}
}

View file

@ -0,0 +1,98 @@
import asyncio
import os
import time
import uuid
from dataclasses import dataclass
from datetime import datetime, timezone
from typing import Final
from urllib.parse import parse_qsl, urlencode, urlsplit, urlunsplit
import psycopg
import pytest
from psycopg import sql
from prisma import Prisma
from integration._support.database import read_rows
from litellm.proxy.db.db_transaction_queue.spend_logs_partition_manager import SpendLogsPartitionManager
@dataclass(frozen=True)
class PartitionConnection:
db: Prisma
@pytest.mark.covers(
"other.database.partitions.lock_wait_outlives_transaction_default",
"other.database.partitions.repeat_preserves_rows",
)
async def test_real_partition_ddl_survives_witnessed_lock_and_is_idempotent() -> None:
schema: Final = f"integration_{uuid.uuid4().hex}"
url: Final = os.environ["DATABASE_URL"]
parsed: Final = urlsplit(url)
scoped_url: Final = urlunsplit(
parsed._replace(query=urlencode({**dict(parse_qsl(parsed.query)), "schema": schema}))
)
parent: Final = sql.Identifier(schema, "LiteLLM_SpendLogs")
with psycopg.connect(url, autocommit=True) as setup:
setup.execute(sql.SQL("CREATE SCHEMA {}").format(sql.Identifier(schema)))
try:
setup.execute(
sql.SQL(
'CREATE TABLE {} (request_id text, "startTime" timestamp NOT NULL) PARTITION BY RANGE ("startTime")'
).format(parent)
)
database: Final = Prisma(datasource={"url": scoped_url})
await database.connect()
try:
manager: Final = SpendLogsPartitionManager(interval="day", precreate_ahead=0)
with psycopg.connect(url) as blocker:
blocker.execute(sql.SQL("LOCK TABLE {} IN ACCESS SHARE MODE").format(parent))
blocker_pid: Final = blocker.info.backend_pid
operation: Final = asyncio.create_task(
manager.ensure_partitions(PartitionConnection(database), lambda: 7000)
)
wait_deadline: Final = time.monotonic() + 3
try:
while True:
witnesses: Final = read_rows(
"SELECT a.pid, extract(epoch FROM "
"clock_timestamp()-a.query_start)::double precision AS age "
"FROM pg_stat_activity a WHERE %s = ANY(pg_blocking_pids(a.pid)) "
"AND a.wait_event_type = 'Lock' AND a.query LIKE 'CREATE TABLE IF NOT EXISTS%%'",
(blocker_pid,),
)
if witnesses:
break
assert time.monotonic() < wait_deadline, "Partition DDL never reached the held lock"
await asyncio.sleep(0.02)
assert len(witnesses) == 1
held_at: Final = time.monotonic()
age: Final = float(witnesses[0]["age"])
await asyncio.sleep(max(0, 5.6 - age))
held_seconds: Final = age + time.monotonic() - held_at
assert held_seconds >= 5.5, f"Lock released before the transaction boundary: {held_seconds}"
assert not operation.done(), "DDL completed while its required lock was held"
except BaseException:
operation.cancel()
await asyncio.gather(operation, return_exceptions=True)
raise
finally:
blocker.rollback()
ensured: Final = await asyncio.wait_for(operation, timeout=5)
assert len(ensured) == 1, "Partition DDL failed after the permitted lock wait"
catalog: Final = read_rows(
"SELECT child.relname FROM pg_inherits i JOIN pg_class child ON child.oid=i.inhrelid "
"JOIN pg_class parent ON parent.oid=i.inhparent JOIN pg_namespace n ON n.oid=parent.relnamespace "
"WHERE n.nspname=%s AND parent.relname='LiteLLM_SpendLogs'",
(schema,),
)
assert catalog == [{"relname": ensured[0]}]
now: Final = datetime.now(timezone.utc).replace(tzinfo=None)
setup.execute(sql.SQL("INSERT INTO {} VALUES (%s, %s)").format(parent), ("retained", now))
assert await manager.ensure_partitions(PartitionConnection(database), lambda: 7000) == ensured
assert setup.execute(sql.SQL("SELECT request_id FROM {}").format(parent)).fetchall() == [("retained",)]
finally:
await database.disconnect()
finally:
setup.execute(sql.SQL("DROP SCHEMA {} CASCADE").format(sql.Identifier(schema)))
assert read_rows("SELECT nspname FROM pg_namespace WHERE nspname=%s", (schema,)) == []

View file

@ -0,0 +1,138 @@
import os
import uuid
from hashlib import sha256
from pathlib import Path
from concurrent.futures import ThreadPoolExecutor
from typing import Final
from urllib.parse import urlsplit, urlunsplit
import psycopg
import pytest
from psycopg import sql
from integration._support.client import Gateway, eventually, string_value
from integration._support.database import read_rows
from integration._support.process import owned_proxy
def delete_if_present(candidate: Gateway, key: str) -> None:
digest: Final = sha256(key.encode()).hexdigest()
if read_rows('SELECT token FROM "LiteLLM_VerificationToken" WHERE token=%s', (digest,)):
candidate.post("/key/delete", {"keys": [key]})
assert read_rows('SELECT token FROM "LiteLLM_VerificationToken" WHERE token=%s', (digest,)) == []
@pytest.mark.covers("other.database.regeneration.writer_updates_dependent_grants")
def test_key_regeneration_uses_writer_with_a_real_readonly_reader(gateway: Gateway, tmp_path: Path) -> None:
role: Final = f"integration_reader_{uuid.uuid4().hex}"
url: Final = os.environ["DATABASE_URL"]
parsed: Final = urlsplit(url)
reader_url: Final = urlunsplit(
parsed._replace(netloc=f"{role}:integration-reader-password@{parsed.hostname}:{parsed.port}")
)
with psycopg.connect(url, autocommit=True) as admin:
admin.execute(
sql.SQL("CREATE ROLE {} LOGIN PASSWORD 'integration-reader-password' NOSUPERUSER NOINHERIT").format(
sql.Identifier(role)
)
)
try:
admin.execute(sql.SQL("GRANT USAGE ON SCHEMA public TO {}").format(sql.Identifier(role)))
admin.execute(sql.SQL("GRANT SELECT ON ALL TABLES IN SCHEMA public TO {}").format(sql.Identifier(role)))
admin.execute(sql.SQL("ALTER ROLE {} SET default_transaction_read_only = on").format(sql.Identifier(role)))
with psycopg.connect(reader_url, autocommit=True) as reader:
assert reader.execute("SHOW transaction_read_only").fetchone() == ("on",)
with pytest.raises(psycopg.errors.ReadOnlySqlTransaction):
reader.execute('UPDATE "LiteLLM_VerificationToken" SET blocked = true WHERE false')
with owned_proxy(gateway, tmp_path, {"DATABASE_URL_READ_REPLICA": reader_url}) as candidate:
assert read_rows("SELECT pid FROM pg_stat_activity WHERE usename=%s", (role,)), (
"Candidate reader was never connected"
)
with gateway.scenario() as scenario:
model: Final = scenario.model()
outside: Final = scenario.model()
old: Final = string_value(candidate.post("/key/generate", {"models": [outside]})["key"])
new: Final = f"sk-integration-{uuid.uuid4().hex}"
scenario.cleanups.callback(delete_if_present, gateway, old)
scenario.cleanups.callback(delete_if_present, gateway, new)
old_hash: Final = sha256(old.encode()).hexdigest()
before: Final = candidate.request(
"POST",
"/v1/chat/completions",
{"model": model, "messages": [{"role": "user", "content": "no grant yet"}]},
key=old,
)
assert before.status_code == 403 and before.json()["error"]["type"] == "key_model_access_denied", (
before.text
)
response: Final = candidate.request(
"POST",
"/v1/access_group",
{
"access_group_name": f"integration-{uuid.uuid4().hex}",
"access_model_names": [model],
"assigned_key_ids": [old_hash],
},
)
assert response.status_code == 201, response.text
group: Final = string_value(response.json()["access_group_id"])
try:
with psycopg.connect(url) as blocker, ThreadPoolExecutor(max_workers=1) as executor:
blocker.execute('LOCK TABLE "LiteLLM_AccessGroupTable" IN ACCESS EXCLUSIVE MODE')
pending: Final = executor.submit(candidate.request, "GET", f"/v1/access_group/{group}")
try:
reached: Final = eventually(
lambda: read_rows(
"SELECT usename FROM pg_stat_activity WHERE %s=ANY(pg_blocking_pids(pid)) "
"AND usename=%s AND query LIKE 'SELECT%%'",
(blocker.info.backend_pid, role),
),
bool,
seconds=3,
)
assert reached == [{"usename": role}]
finally:
blocker.rollback()
selected: Final = pending.result(timeout=5)
assert selected.status_code == 200 and selected.json()["access_group_id"] == group, (
selected.text
)
assert candidate.chat(model, key=old)["usage"]["total_tokens"] == 40
regenerated: Final = candidate.post(
"/key/regenerate", {"key": old, "new_key": new, "grace_period": "0s"}
)
assert regenerated["key"] == new
new_hash: Final = sha256(new.encode()).hexdigest()
assert new != old
assert read_rows(
'SELECT assigned_key_ids FROM "LiteLLM_AccessGroupTable" WHERE access_group_id=%s', (group,)
) == [{"assigned_key_ids": [new_hash]}]
assert read_rows(
'SELECT token, access_group_ids FROM "LiteLLM_VerificationToken" WHERE token=ANY(%s)',
([old_hash, new_hash],),
) == [{"token": new_hash, "access_group_ids": [group]}]
assert candidate.chat(model, key=new)["usage"]["total_tokens"] == 40
assert candidate.chat(outside, key=new)["usage"]["total_tokens"] == 40
denied: Final = candidate.request(
"POST",
"/v1/chat/completions",
{"model": model, "messages": [{"role": "user", "content": "rotated key"}]},
key=old,
)
assert (
denied.status_code == 401 and denied.json()["error"]["type"] == "token_not_found_in_db"
), denied.text
finally:
deleted: Final = gateway.request("DELETE", f"/v1/access_group/{group}")
assert deleted.status_code == 204, deleted.text
assert (
read_rows(
'SELECT access_group_id FROM "LiteLLM_AccessGroupTable" WHERE access_group_id=%s',
(group,),
)
== []
)
finally:
admin.execute(sql.SQL("DROP OWNED BY {}").format(sql.Identifier(role)))
admin.execute(sql.SQL("DROP ROLE {}").format(sql.Identifier(role)))
assert read_rows("SELECT rolname FROM pg_roles WHERE rolname=%s", (role,)) == []

View file

@ -0,0 +1,125 @@
import os
import uuid
from contextlib import ExitStack
from hashlib import sha256
from typing import Final
import psycopg
import pytest
from psycopg import sql
from integration._support.client import Gateway
from integration._support.database import read_rows
@pytest.mark.covers("other.database.access_group.failed_second_write_rolls_back_first")
def test_access_group_second_key_constraint_failure_rolls_back_all_writes(gateway: Gateway) -> None:
with gateway.scenario() as scenario:
model: Final = scenario.model()
outside: Final = scenario.model()
keys: Final = (scenario.key(models=[outside]), scenario.key(models=[outside]))
tokens: Final = [sha256(key.encode()).hexdigest() for key in keys]
name: Final = f"integration-{uuid.uuid4().hex}"
constraint: Final = f"integration_reject_{uuid.uuid4().hex}"
witness: Final = constraint + "_seq"
check_function: Final = constraint + "_check"
body: Final = {"access_group_name": name, "access_model_names": [model], "assigned_key_ids": tokens}
def remove_partial_group() -> None:
for row in read_rows(
'SELECT access_group_id FROM "LiteLLM_AccessGroupTable" WHERE access_group_name=%s', (name,)
):
response: Final = gateway.request("DELETE", f"/v1/access_group/{row['access_group_id']}")
assert response.status_code == 204, response.text
assert (
read_rows('SELECT access_group_id FROM "LiteLLM_AccessGroupTable" WHERE access_group_name=%s', (name,))
== []
)
scenario.cleanups.callback(remove_partial_group)
before: Final = read_rows(
'SELECT token, access_group_ids FROM "LiteLLM_VerificationToken" WHERE token=ANY(%s) ORDER BY token',
(tokens,),
)
with psycopg.connect(os.environ["DATABASE_URL"], autocommit=True) as connection, ExitStack() as cleanup:
connection.execute(sql.SQL("CREATE SEQUENCE {}").format(sql.Identifier(witness)))
cleanup.callback(connection.execute, sql.SQL("DROP SEQUENCE {}").format(sql.Identifier(witness)))
connection.execute(
sql.SQL(
"CREATE FUNCTION {}(text[]) RETURNS boolean LANGUAGE plpgsql AS $$ BEGIN IF "
"cardinality($1)>0 THEN PERFORM nextval({}); RETURN false; END IF; RETURN true; END $$"
).format(sql.Identifier(check_function), sql.Literal(witness))
)
cleanup.callback(
connection.execute, sql.SQL("DROP FUNCTION {}(text[])").format(sql.Identifier(check_function))
)
connection.execute(
sql.SQL(
'ALTER TABLE "LiteLLM_VerificationToken" ADD '
"CONSTRAINT {} CHECK (token <> {} OR {}(access_group_ids))"
).format(sql.Identifier(constraint), sql.Literal(tokens[1]), sql.Identifier(check_function))
)
cleanup.callback(
connection.execute,
sql.SQL('ALTER TABLE "LiteLLM_VerificationToken" DROP CONSTRAINT {}').format(
sql.Identifier(constraint)
),
)
try:
assert connection.execute(
sql.SQL("SELECT is_called FROM {}").format(sql.Identifier(witness))
).fetchone() == (False,)
failed: Final = gateway.request("POST", "/v1/access_group", body)
assert failed.status_code == 500, failed.text
assert connection.execute(
sql.SQL("SELECT is_called FROM {}").format(sql.Identifier(witness))
).fetchone() == (True,)
assert (
read_rows(
'SELECT access_group_id FROM "LiteLLM_AccessGroupTable" WHERE access_group_name=%s', (name,)
)
== []
)
assert (
read_rows(
"SELECT token, access_group_ids FROM "
'"LiteLLM_VerificationToken" WHERE token=ANY(%s) ORDER BY token',
(tokens,),
)
== before
)
for key in keys:
denied: Final = gateway.request(
"POST",
"/v1/chat/completions",
{"model": model, "messages": [{"role": "user", "content": "rolled back grant"}]},
key=key,
)
assert denied.status_code == 403 and denied.json()["error"]["type"] == "key_model_access_denied", (
denied.text
)
finally:
cleanup.close()
created: Final = gateway.request("POST", "/v1/access_group", body)
assert created.status_code == 201, created.text
identity: Final = created.json()["access_group_id"]
try:
for key in keys:
assert gateway.chat(model, key=key)["usage"]["total_tokens"] == 40
finally:
deleted: Final = gateway.request("DELETE", f"/v1/access_group/{identity}")
assert deleted.status_code == 204, deleted.text
assert (
read_rows(
'SELECT access_group_id FROM "LiteLLM_AccessGroupTable" WHERE access_group_id=%s', (identity,)
)
== []
)
assert (
read_rows(
'SELECT token, access_group_ids FROM "LiteLLM_VerificationToken" WHERE token=ANY(%s) ORDER BY token',
(tokens,),
)
== before
)
assert read_rows("SELECT conname FROM pg_constraint WHERE conname=%s", (constraint,)) == []

View file

@ -0,0 +1,123 @@
import json
import uuid
from typing import Final
import pytest
from hypothesis import Phase, example, given, settings, strategies as st
from integration._support.client import Gateway, eventually, object_value
from integration._support.database import read_rows
@pytest.mark.covers("quota_management.spend_tracking.price_precedence.zero_and_default_rates")
@pytest.mark.timeout(180)
def test_generated_zero_null_and_omitted_prices_follow_independent_arithmetic(gateway: Gateway) -> None:
@settings(max_examples=20, deadline=None, database=None, phases=(Phase.explicit, Phase.generate, Phase.shrink))
@example(rates=(0, 0))
@example(rates=(1, 2))
@example(rates=("null", "null"))
@given(
rates=st.one_of(
st.sampled_from((("omitted", "omitted"), ("null", "null"))),
st.tuples(st.integers(0, 25), st.integers(0, 25)),
)
)
def check(rates: tuple[str | int, str | int]) -> None:
defaults: Final = rates[0] in ("omitted", "null")
assert defaults or (isinstance(rates[0], int) and isinstance(rates[1], int))
input_rate, output_rate = (
(0.00000015, 0.0000006) if defaults else (float(rates[0]) / 1_000_000, float(rates[1]) / 1_000_000)
)
parameters: Final = (
{}
if rates[0] == "omitted"
else {
"input_cost_per_token": None if rates[0] == "null" else input_rate,
"output_cost_per_token": None if rates[0] == "null" else output_rate,
}
)
with gateway.scenario() as scenario:
model: Final = scenario.model(**parameters)
response: Final = gateway.request(
"POST",
"/v1/chat/completions",
{
"model": model,
"messages": [{"role": "user", "content": f"independent price {uuid.uuid4().hex}"}],
},
)
assert response.status_code == 200, response.text
assert response.json()["usage"] == {"prompt_tokens": 20, "completion_tokens": 20, "total_tokens": 40}
expected: Final = 20 * input_rate + 20 * output_rate
if expected:
assert float(response.headers["x-litellm-response-cost"]) == pytest.approx(expected, rel=1e-6)
else:
assert response.headers.get("x-litellm-response-cost") in (None, "0", "0.0")
rows: Final = eventually(
lambda: read_rows(
"SELECT spend, metadata, prompt_tokens, "
'completion_tokens FROM "LiteLLM_SpendLogs" WHERE request_id=%s',
(response.json()["id"],),
),
lambda values: len(values) == 1,
seconds=70,
)
assert rows[0]["prompt_tokens"] == 20 and rows[0]["completion_tokens"] == 20
assert float(rows[0]["spend"]) == pytest.approx(expected, rel=1e-6)
metadata: Final = rows[0]["metadata"]
parsed: Final = json.loads(metadata) if isinstance(metadata, str) else object_value(metadata)
breakdown: Final = object_value(parsed["cost_breakdown"])
assert float(breakdown["input_cost"]) == pytest.approx(20 * input_rate, rel=1e-6)
assert float(breakdown["output_cost"]) == pytest.approx(20 * output_rate, rel=1e-6)
check()
@pytest.mark.covers("quota_management.spend_tracking.alias_prices.remain_independent_on_reload")
def test_same_upstream_aliases_keep_distinct_prices_after_reload(gateway: Gateway) -> None:
for order in (("free", "paid"), ("paid", "free")):
with gateway.scenario() as scenario:
rates: Final = {
"free": {"input_cost_per_token": 0, "output_cost_per_token": 0},
"paid": {"input_cost_per_token": 0.001, "output_cost_per_token": 0.003},
}
aliases: Final = {kind: scenario.model(**rates[kind]) for kind in order}
for generation in range(2):
for kind in order if generation == 0 else reversed(order):
model: Final = aliases[kind]
cost: Final = 0.08 if kind == "paid" else 0.0
response: Final = gateway.request(
"POST",
"/v1/chat/completions",
{
"model": model,
"messages": [{"role": "user", "content": f"alias price {model} {generation}"}],
},
)
assert response.status_code == 200, response.text
assert response.json()["usage"]["total_tokens"] == 40
if cost:
assert float(response.headers["x-litellm-response-cost"]) == pytest.approx(cost)
rows: Final = eventually(
lambda response=response: read_rows(
'SELECT spend, metadata FROM "LiteLLM_SpendLogs" WHERE request_id=%s',
(response.json()["id"],),
),
lambda values: len(values) == 1,
seconds=70,
)
assert float(rows[0]["spend"]) == pytest.approx(cost)
metadata: Final = rows[0]["metadata"]
parsed: Final = json.loads(metadata) if isinstance(metadata, str) else object_value(metadata)
breakdown: Final = object_value(parsed["cost_breakdown"])
assert float(breakdown["input_cost"]) == pytest.approx(20 * rates[kind]["input_cost_per_token"])
assert float(breakdown["output_cost"]) == pytest.approx(20 * rates[kind]["output_cost_per_token"])
if generation == 0:
entries: Final = gateway.get("/model/info")["data"]
target: Final = next(entry for entry in entries if entry["model_name"] == aliases["paid"])
changed: Final = gateway.request(
"PATCH",
f"/model/{target['model_info']['id']}/update",
{"model_info": {"description": "price reload"}},
)
assert changed.status_code == 200, changed.text

View file

@ -0,0 +1,61 @@
import json
import uuid
from typing import Final
import pytest
from integration._support.client import Gateway, eventually, object_value
from integration._support.database import read_rows
from integration._support.wire import Reply, Request, wire_server
@pytest.mark.covers("other.provider_wire.anthropic.tool_history_system_cache_and_internal_fields", "quota_management.spend_tracking.cache_tokens.disjoint_classes_use_explicit_rates")
def test_anthropic_tool_history_and_cache_tokens_keep_wire_and_accounting_contracts(gateway: Gateway) -> None:
identity: Final = "anthropic-wire-" + uuid.uuid4().hex
tool_schema: Final = {"type": "object", "properties": {"x": {"type": "integer"}, "y": {"type": "integer"}}, "required": ["x", "y"]}
def respond(request: Request) -> Reply:
assert request.method == "POST" and request.target == "/v1/messages"
assert request.headers["x-api-key"] == "synthetic-anthropic-key"
body: Final = json.loads(request.body)
assert body["model"] == "claude-sonnet-4-5-20250929"
assert body["system"] == [{"type": "text", "text": "synthetic policy", "cache_control": {"type": "ephemeral"}}]
assert body["tools"][0]["name"] == "add" and body["tools"][0]["input_schema"] == tool_schema
assert body["max_tokens"] == 16
assert not {"timeout", "stream_chunk_size", "litellm_params", "litellm_metadata", "rpm", "tpm"}.intersection(body)
messages: Final = body["messages"]
assert [message["role"] for message in messages] == ["user", "assistant", "user"]
assert messages[0]["content"] == [{"type": "text", "text": "first"}]
assert messages[1]["content"] == [{"type": "tool_use", "id": "history-call", "name": "add", "input": {"x": 1, "y": 2}}]
assert messages[2]["content"] == [{"type": "tool_result", "tool_use_id": "history-call", "content": "3"}, {"type": "text", "text": "next"}]
return Reply(body=json.dumps({"id": identity, "type": "message", "role": "assistant", "model": "claude-sonnet-4-5-20250929", "content": [{"type": "tool_use", "id": "next-call", "name": "add", "input": {"x": 3, "y": 4}}], "stop_reason": "tool_use", "stop_sequence": None, "usage": {"input_tokens": 10, "output_tokens": 4, "cache_read_input_tokens": 5, "cache_creation_input_tokens": 7}}).encode())
with wire_server(respond) as wire, gateway.scenario() as scenario:
model: Final = scenario.model(model="anthropic/claude-sonnet-4-5-20250929", api_base=wire.url, api_key="synthetic-anthropic-key", input_cost_per_token=0.001, output_cost_per_token=0.002, cache_read_input_token_cost=0.0001, cache_creation_input_token_cost=0.002)
response: Final = gateway.request("POST", "/v1/chat/completions", {
"model": model, "max_tokens": 16, "timeout": 5,
"messages": [
{"role": "system", "content": [{"type": "text", "text": "synthetic policy", "cache_control": {"type": "ephemeral"}}]},
{"role": "user", "content": "first"},
{"role": "assistant", "tool_calls": [{"id": "history-call", "type": "function", "function": {"name": "add", "arguments": '{"x":1,"y":2}'}}]},
{"role": "tool", "tool_call_id": "history-call", "content": "3"},
{"role": "user", "content": "next"},
],
"tools": [{"type": "function", "function": {"name": "add", "parameters": tool_schema}}],
})
assert response.status_code == 200, response.text
body: Final = response.json()
assert body["id"].startswith("chatcmpl-")
assert body["choices"][0]["finish_reason"] == "tool_calls"
tool: Final = body["choices"][0]["message"]["tool_calls"][0]
assert tool["id"] == "next-call" and tool["function"]["name"] == "add"
assert json.loads(tool["function"]["arguments"]) == {"x": 3, "y": 4}
assert body["usage"]["prompt_tokens"] == 22 and body["usage"]["completion_tokens"] == 4
assert len(wire.drain()) == 1
rows: Final = eventually(lambda: read_rows('SELECT spend, prompt_tokens, completion_tokens, metadata FROM "LiteLLM_SpendLogs" WHERE request_id=%s', (body["id"],)), lambda values: len(values) == 1, seconds=70)
assert float(rows[0]["spend"]) == pytest.approx(10 * 0.001 + 5 * 0.0001 + 7 * 0.002 + 4 * 0.002)
assert rows[0]["prompt_tokens"] == 22 and rows[0]["completion_tokens"] == 4
metadata: Final = rows[0]["metadata"]
parsed: Final = json.loads(metadata) if isinstance(metadata, str) else object_value(metadata)
assert parsed["cost_breakdown"]["input_cost"] == pytest.approx(0.0245)
assert parsed["cost_breakdown"]["output_cost"] == pytest.approx(0.008)

View file

@ -0,0 +1,99 @@
import asyncio
import json
import os
import uuid
from pathlib import Path
from typing import Final
import pytest
import yaml
from integration._support.client import Gateway
from integration._support.process import owned_proxy
from integration._support.wire import Reply, Request, wire_server
MODEL: Final = "bedrock/converse/anthropic.claude-3-haiku-20240307-v1:0"
TOKEN: Final = "synthetic-bedrock-bearer"
RESPONSE: Final = json.dumps({
"output": {"message": {"role": "assistant", "content": [{"text": "bedrock wire control"}]}},
"stopReason": "end_turn", "usage": {"inputTokens": 11, "outputTokens": 4, "totalTokens": 15},
"metrics": {"latencyMs": 1},
}).encode()
def bearer_peer(request: Request) -> Reply:
assert request.method == "POST" and request.target == "/model/anthropic.claude-3-haiku-20240307-v1%3A0/converse"
assert request.headers["authorization"] == f"Bearer {TOKEN}"
assert "x-amz-security-token" not in request.headers
body: Final = json.loads(request.body)
assert body["messages"] == [{"role": "user", "content": [{"text": "synthetic bearer request"}]}]
assert body["system"] == [{"text": "synthetic system"}]
assert body["inferenceConfig"]["maxTokens"] == 16
assert not {"timeout", "stream_chunk_size", "litellm_params", "litellm_metadata", "api_key"}.intersection(body)
return Reply(body=RESPONSE)
@pytest.mark.covers("other.provider_wire.bedrock.bearer_sdk_skips_credential_chain")
async def test_bearer_only_sdk_sync_async_requests_do_not_require_aws_credentials(monkeypatch: pytest.MonkeyPatch, tmp_path: Path) -> None:
import litellm
empty: Final = tmp_path / "empty-aws-config"
empty.write_text("")
for name in tuple(name for name in os.environ if name.startswith("AWS_")):
monkeypatch.delenv(name, raising=False)
for name, value in {"AWS_CONFIG_FILE": str(empty), "AWS_SHARED_CREDENTIALS_FILE": str(empty), "AWS_EC2_METADATA_DISABLED": "true", "LITELLM_RUST": "false"}.items():
monkeypatch.setenv(name, value)
with wire_server(bearer_peer) as wire:
with pytest.raises(litellm.APIConnectionError, match=r"config profile .* could not be found"):
await asyncio.to_thread(litellm.completion, model=MODEL, aws_profile_name="integration-profile-must-not-be-read", aws_region_name="us-east-1", aws_bedrock_runtime_endpoint=wire.url, messages=[{"role": "user", "content": "synthetic credential control"}], timeout=5, num_retries=0)
assert wire.drain() == ()
for source in ("argument", "environment"):
if source == "environment":
monkeypatch.setenv("AWS_BEARER_TOKEN_BEDROCK", TOKEN)
parameters: Final = {
"model": MODEL, "api_key": TOKEN if source == "argument" else None,
"aws_region_name": "us-east-1", "aws_profile_name": "integration-profile-must-not-be-read",
"aws_bedrock_runtime_endpoint": wire.url, "timeout": 5, "num_retries": 0,
"messages": [{"role": "system", "content": "synthetic system"}, {"role": "user", "content": "synthetic bearer request"}],
"max_tokens": 16,
}
for asynchronous in (False, True):
result: Final = await litellm.acompletion(**parameters) if asynchronous else await asyncio.to_thread(litellm.completion, **parameters)
assert result.choices[0].message.content == "bedrock wire control"
assert result.choices[0].finish_reason == "stop"
assert result.usage.prompt_tokens == 11 and result.usage.completion_tokens == 4
assert len(wire.drain()) == 1
@pytest.mark.covers("other.provider_wire.bedrock.bearer_db_yaml_survives_reload")
def test_bearer_environment_reference_loads_from_db_and_yaml_and_survives_reload(gateway: Gateway, tmp_path: Path) -> None:
empty: Final = tmp_path / "empty-aws-config"
empty.write_text("")
with wire_server(bearer_peer) as wire:
parameters: Final = {
"model": MODEL, "api_key": "os.environ/INTEGRATION_BEARER_TOKEN", "aws_region_name": "us-east-1",
"aws_profile_name": "integration-profile-must-not-be-read", "aws_bedrock_runtime_endpoint": wire.url,
}
alias: Final = f"integration-yaml-{uuid.uuid4().hex}"
configuration: Final = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text())
configuration["model_list"] = [{"model_name": alias, "litellm_params": parameters, "model_info": {"id": alias}}]
path: Final = tmp_path / "bedrock.yaml"
path.write_text(yaml.safe_dump(configuration))
overrides: Final = {"INTEGRATION_BEARER_TOKEN": TOKEN, "AWS_CONFIG_FILE": str(empty), "AWS_SHARED_CREDENTIALS_FILE": str(empty), "AWS_EC2_METADATA_DISABLED": "true", "LITELLM_RUST": "false"}
with owned_proxy(gateway, tmp_path, overrides, config=path, remove_environment=tuple(name for name in os.environ if name.startswith("AWS_"))) as candidate, candidate.scenario() as scenario:
database_model: Final = scenario.model(**parameters)
for generation in range(2):
for model in (alias, database_model):
response: Final = candidate.request("POST", "/v1/chat/completions", {
"model": model, "messages": [{"role": "system", "content": "synthetic system"}, {"role": "user", "content": "synthetic bearer request"}],
"max_tokens": 16, "cache": {"no-cache": True},
})
assert response.status_code == 200, response.text
assert response.json()["choices"][0]["message"]["content"] == "bedrock wire control"
assert response.json()["usage"]["total_tokens"] == 15
assert len(wire.drain()) == 1, f"Expected actual provider call after reload {generation}"
if generation == 0:
entries: Final = candidate.get("/model/info")["data"]
target: Final = next(entry for entry in entries if entry["model_name"] == database_model)
response: Final = candidate.request("PATCH", f"/model/{target['model_info']['id']}/update", {"model_info": {"description": "bearer reload"}})
assert response.status_code == 200, response.text

View file

@ -0,0 +1,75 @@
import json
import os
import uuid
from pathlib import Path
from typing import Final
from urllib.parse import parse_qs
import pytest
import yaml
from integration._support.client import Gateway
from integration._support.process import owned_proxy
from integration._support.wire import Reply, Request, wire_server
from integration.providers.test_bedrock_auth_wire import MODEL, RESPONSE
@pytest.mark.covers("other.provider_wire.bedrock.db_yaml_role_reference_reaches_sts_and_signed_request")
def test_role_reference_from_db_and_yaml_reaches_real_sts_http_and_bedrock(gateway: Gateway, tmp_path: Path) -> None:
role: Final = "arn:aws:iam::123456789012:role/integration-" + uuid.uuid4().hex
assumed_key: Final = "ASIAINTEGRATION000001"
assumed_token: Final = "synthetic-assumed-session-token"
def sts(request: Request) -> Reply:
parameters: Final = parse_qs(request.body.decode())
action: Final = parameters["Action"][0]
assert request.method == "POST" and action in {"GetCallerIdentity", "AssumeRole"}
if action == "GetCallerIdentity":
result = "<GetCallerIdentityResult><Arn>arn:aws:iam::123456789012:user/integration-source</Arn><UserId>integration-source</UserId><Account>123456789012</Account></GetCallerIdentityResult>"
else:
assert parameters["RoleArn"] == [role]
assert parameters["RoleSessionName"][0] in {"integration-yaml-session", "integration-db-session"}
result = f"<AssumeRoleResult><Credentials><AccessKeyId>{assumed_key}</AccessKeyId><SecretAccessKey>synthetic-assumed-secret-key-for-testing</SecretAccessKey><SessionToken>{assumed_token}</SessionToken><Expiration>2035-01-01T00:00:00Z</Expiration></Credentials><AssumedRoleUser><Arn>arn:aws:sts::123456789012:assumed-role/integration/session</Arn><AssumedRoleId>integration:session</AssumedRoleId></AssumedRoleUser><PackedPolicySize>0</PackedPolicySize></AssumeRoleResult>"
return Reply(content_type="text/xml", body=f'<{action}Response xmlns="https://sts.amazonaws.com/doc/2011-06-15/">{result}<ResponseMetadata><RequestId>synthetic-sts-request</RequestId></ResponseMetadata></{action}Response>'.encode())
def bedrock(request: Request) -> Reply:
assert request.method == "POST" and request.target == "/model/anthropic.claude-3-haiku-20240307-v1%3A0/converse"
assert f"Credential={assumed_key}/" in request.headers["authorization"]
assert request.headers["x-amz-security-token"] == assumed_token
assert json.loads(request.body)["messages"][0]["content"][0]["text"] == "synthetic role request"
return Reply(body=RESPONSE)
with wire_server(sts) as authority, wire_server(bedrock) as provider:
parameters: Final = {
"model": MODEL, "aws_region_name": "us-east-1", "aws_role_name": "os.environ/INTEGRATION_ROLE_ARN",
"aws_session_name": "integration-yaml-session", "aws_bedrock_runtime_endpoint": provider.url,
"aws_sts_endpoint": authority.url,
}
alias: Final = "integration-role-yaml-" + uuid.uuid4().hex
configuration: Final = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text())
configuration["model_list"] = [{"model_name": alias, "litellm_params": parameters, "model_info": {"id": alias}}]
path: Final = tmp_path / "roles.yaml"
path.write_text(yaml.safe_dump(configuration))
empty: Final = tmp_path / "empty-aws-config"
empty.write_text("")
overrides: Final = {
"INTEGRATION_ROLE_ARN": role, "AWS_ACCESS_KEY_ID": "AKIAINTEGRATION000001", "AWS_SECRET_ACCESS_KEY": "synthetic-source-secret-key-for-testing",
"AWS_CONFIG_FILE": str(empty), "AWS_SHARED_CREDENTIALS_FILE": str(empty), "AWS_EC2_METADATA_DISABLED": "true",
"AWS_ENDPOINT_URL_STS": authority.url, "AWS_DEFAULT_REGION": "us-east-1", "LITELLM_RUST": "false",
}
with owned_proxy(gateway, tmp_path, overrides, config=path, remove_environment=tuple(name for name in os.environ if name.startswith("AWS_"))) as candidate, candidate.scenario() as scenario:
database_model: Final = scenario.model(**{**parameters, "api_key": None, "aws_session_name": "integration-db-session"})
for generation in range(2):
for model in (alias, database_model):
response: Final = candidate.request("POST", "/v1/chat/completions", {"model": model, "messages": [{"role": "user", "content": "synthetic role request"}], "cache": {"no-cache": True}})
assert response.status_code == 200, response.text
assert response.json()["choices"][0]["message"]["content"] == "bedrock wire control"
assert response.json()["usage"]["total_tokens"] == 15
assert len(provider.drain()) == 1
if generation == 0:
target: Final = next(entry for entry in candidate.get("/model/info")["data"] if entry["model_name"] == database_model)
response: Final = candidate.request("PATCH", f"/model/{target['model_info']['id']}/update", {"model_info": {"description": "role reload"}})
assert response.status_code == 200, response.text
assumed: Final = tuple(parse_qs(request.body.decode()) for request in authority.drain() if parse_qs(request.body.decode())["Action"] == ["AssumeRole"])
assert {entry["RoleSessionName"][0] for entry in assumed} == {"integration-yaml-session", "integration-db-session"}
assert all(entry["RoleArn"] == [role] for entry in assumed)

View file

@ -0,0 +1,111 @@
import asyncio
import base64
import hashlib
import hmac
import json
from datetime import datetime
from typing import Final
import httpx
import pytest
from integration._support.sigv4 import encoded_path, signature
from integration._support.wire import Reply, Request, wire_server
ACCESS: Final = "AKIAIOSFODNN7EXAMPLE"
SECRET: Final = "wJalrXUtnFEMI/K7MDENG/bPxRfiCYEXAMPLEKEY"
@pytest.mark.covers("other.provider_wire.s3.verifier_known_answer_and_negative_controls")
def test_sigv4_verifier_matches_published_put_and_rejects_corruption() -> None:
# Public AWS example credentials and PUT vector, not an active account:
# https://docs.aws.amazon.com/AmazonS3/latest/developerguide/sig-v4-header-based-auth.html
headers: Final = {
"date": "Fri, 24 May 2013 00:00:00 GMT", "host": "examplebucket.s3.amazonaws.com",
"x-amz-content-sha256": "44ce7dd67c959e0d3524ffac1771dfbba87d2b6b4b4e99e42034a8b803f8b072",
"x-amz-date": "20130524T000000Z", "x-amz-storage-class": "REDUCED_REDUNDANCY",
}
signed: Final = "date;host;x-amz-content-sha256;x-amz-date;x-amz-storage-class"
expected: Final = (
"9e0e90d9c76de8fa5b200d8c849cd5b8dc7a3be3951ddb7f6a76b4158342019d",
"98ad721746da40c64f1a55b78f14c238d841ea1380cd77a1b5971af0ece108bd",
)
actual: Final = signature("PUT", "/test%24file.text", headers, signed, b"Welcome to Amazon S3.", SECRET, "20130524/us-east-1/s3/aws4_request")
assert actual == expected
assert signature("PUT", "/test$file.text", headers, signed, b"Welcome to Amazon S3.", SECRET, "20130524/us-east-1/s3/aws4_request") != expected
assert encoded_path("/bucket/a=b+c/d e/雪.json") == "/bucket/a%3Db%2Bc/d%20e/%E9%9B%AA.json"
@pytest.mark.covers("other.provider_wire.s3.sync_async_reserved_keys_are_signed_and_accepted")
async def test_s3_sync_and_async_uploads_pass_independent_wire_verification(monkeypatch: pytest.MonkeyPatch) -> None:
from litellm.integrations.s3_v2 import S3Logger
from litellm.types.integrations.s3_v2 import s3BatchLoggingElement
monkeypatch.setattr("botocore.auth.get_current_datetime", lambda: datetime(2026, 9, 14))
payload: Final = {"id": "synthetic-event", "content": "synthetic snow 雪"}
expected_path = ""
def verify(request: Request) -> Reply:
if request.method != "PUT" or request.target != expected_path:
return Reply(status=403)
try:
authorization: Final = request.headers.get("authorization", "")
assert authorization.startswith("AWS4-HMAC-SHA256 ")
fields: Final = dict(part.split("=", 1) for part in authorization.removeprefix("AWS4-HMAC-SHA256 ").split(", "))
access, scope = fields["Credential"].split("/", 1)
assert access == ACCESS and scope == "20260914/us-east-1/s3/aws4_request"
assert request.headers["x-amz-date"] == "20260914T000000Z"
signed: Final = fields["SignedHeaders"].split(";")
assert signed == sorted(set(signed))
assert {"host", "content-md5", "x-amz-date"}.issubset(signed)
assert {name for name in request.headers if name.startswith("x-amz-") and name != "x-amz-content-sha256"}.issubset(signed)
assert request.headers["content-md5"] == base64.b64encode(hashlib.md5(request.body, usedforsecurity=False).digest()).decode()
assert request.headers["x-amz-content-sha256"] == hashlib.sha256(request.body).hexdigest()
expected: Final = signature("PUT", request.target, request.headers, fields["SignedHeaders"], request.body, SECRET, scope)[1]
return Reply(status=200 if hmac.compare_digest(expected, fields["Signature"]) else 403)
except (AssertionError, KeyError, ValueError):
return Reply(status=403)
with wire_server(verify) as wire:
prior: Final = asyncio.all_tasks()
logger: Final = S3Logger(s3_bucket_name="integration-bucket", s3_region_name="us-east-1", s3_endpoint_url=wire.url,
s3_aws_access_key_id=ACCESS, s3_aws_secret_access_key=SECRET, s3_callback_params_override={})
owned: Final = asyncio.all_tasks() - prior
assert len(owned) == 1
try:
for mode in ("sync", "async"):
for key in ("plain.json", "a=b+c/d e/雪.json", "percent%2Fplus+.json"):
expected_path = encoded_path(f"/integration-bucket/{key}")
element: Final = s3BatchLoggingElement(payload=payload, s3_object_key=key, s3_object_download_filename="event.json")
if mode == "sync":
await asyncio.to_thread(logger.upload_data_to_s3, element)
else:
await logger.async_upload_data_to_s3(element)
requests: Final = wire.drain()
assert len(requests) == 1, "Upload must be accepted on its first actual PUT"
request: Final = requests[0]
assert request.target == expected_path
assert json.loads(request.body) == payload
assert verify(request).status == 200
with httpx.Client(timeout=5, trust_env=False) as client:
corrupt: Final = {**request.headers, "authorization": request.headers["authorization"][:-1] + ("0" if request.headers["authorization"][-1] != "0" else "1")}
assert client.put(wire.url + expected_path, content=request.body, headers=corrupt).status_code == 403
assert client.put(wire.url + expected_path + "-wrong", content=request.body, headers=request.headers).status_code == 403
assert client.put(wire.url + expected_path, content=request.body + b" ", headers={name: value for name, value in request.headers.items() if name != "content-length"}).status_code == 403
fields: Final = dict(part.split("=", 1) for part in request.headers["authorization"].removeprefix("AWS4-HMAC-SHA256 ").split(", "))
for signed, scope, md5 in (
(fields["SignedHeaders"].replace("host;", ""), "20260914/us-east-1/s3/aws4_request", request.headers["content-md5"]),
(fields["SignedHeaders"], "20260914/us-west-2/s3/aws4_request", request.headers["content-md5"]),
(fields["SignedHeaders"], "20260914/us-east-1/s3/aws4_request", "AAAAAAAAAAAAAAAAAAAAAA=="),
):
candidate_headers: Final = {**request.headers, "content-md5": md5}
digest: Final = signature("PUT", request.target, candidate_headers, signed, request.body, SECRET, scope)[1]
candidate_headers["authorization"] = f"AWS4-HMAC-SHA256 Credential={ACCESS}/{scope}, SignedHeaders={signed}, Signature={digest}"
assert client.put(wire.url + expected_path, content=request.body, headers=candidate_headers).status_code == 403
assert len(wire.drain()) == 6
finally:
for task in owned:
task.cancel()
await asyncio.gather(*owned, return_exceptions=True)
assert all(task.done() for task in owned)

View file

@ -0,0 +1,98 @@
import json
import uuid
from pathlib import Path
from typing import Final
import httpx
import pytest
import yaml
from integration._support.client import Gateway, object_value
from integration._support.wire import Reply, Request, wire_server
@pytest.mark.covers("other.routing.retries.several_attempts_reach_success_without_hidden_retries", "other.routing.errors.nonretryable_and_exhausted_failures_remain_errors")
def test_retry_counts_and_public_errors_match_actual_provider_attempts(gateway: Gateway) -> None:
with httpx.Client(base_url=gateway.upstream_url, timeout=5, trust_env=False) as upstream, gateway.scenario() as scenario:
original: Final = object_value(gateway.get("/router/settings")["current_values"])["num_retries"]
provider_model: Final = "errors-" + uuid.uuid4().hex
model: Final = scenario.model(model=f"openai/{provider_model}", input_cost_per_token=0, output_cost_per_token=0)
def remove() -> None:
response: Final = upstream.delete(f"/__scripts/{provider_model}")
assert response.status_code in (200, 404)
assert upstream.get(f"/__scripts/{provider_model}").status_code == 404
scenario.cleanups.callback(remove)
try:
for index, (retries, statuses, status, attempts) in enumerate(((2, [500, 500, 200], 200, 3), (2, [400, 200], 400, 1), (1, [429, 429, 200], 429, 2), (1, [500, 500, 200], 500, 2))):
gateway.post("/config/update", {"router_settings": {"num_retries": retries}})
assert object_value(gateway.get("/router/settings")["current_values"])["num_retries"] == retries
upstream.post(f"/__scripts/{provider_model}", json={"statuses": statuses}).raise_for_status()
upstream.get("/__observations").raise_for_status()
response: Final = gateway.request("POST", "/v1/chat/completions", {"model": model, "messages": [{"role": "user", "content": f"{provider_model} {index}"}]})
assert response.status_code == status, response.text
requests: Final = upstream.get("/__observations").json()["requests"]
assert len(requests) == attempts
assert all(request["body"]["model"] == provider_model for request in requests)
assert upstream.get(f"/__scripts/{provider_model}").json()["remaining"] == statuses[attempts:]
if status == 200:
assert response.json()["usage"]["total_tokens"] == 40
else:
error: Final = response.json()["error"]
assert isinstance(error["message"], str) and "Controlled provider failure" in error["message"]
assert str(error["code"]) == str(status)
assert error["type"] == {400: "invalid_request_error", 429: "throttling_error", 500: "internal_server_error"}[status]
assert error["param"] is None
assert "Traceback" not in response.text and "File \"" not in response.text
finally:
gateway.post("/config/update", {"router_settings": {"num_retries": original}})
assert object_value(gateway.get("/router/settings")["current_values"])["num_retries"] == original
@pytest.mark.covers("other.routing.fallback.loaded_configuration_selects_only_permitted_target")
def test_loaded_fallback_selects_expected_deployment_and_keeps_response_identity(tmp_path: Path) -> None:
from litellm import Router
def respond(request: Request) -> Reply:
model: Final = json.loads(request.body)["model"]
assert model in {"primary-wire", "fallback-wire", "unrelated-wire"}
if model == "primary-wire":
return Reply(status=500, body=b'{"error":{"message":"synthetic primary unavailable","type":"api_error","code":"500"}}')
return Reply(body=json.dumps({"id": "response-" + model, "object": "chat.completion", "created": 1, "model": model, "choices": [{"index": 0, "message": {"role": "assistant", "content": "served " + model}, "finish_reason": "stop"}], "usage": {"prompt_tokens": 11, "completion_tokens": 4, "total_tokens": 15}}).encode())
with wire_server(respond) as wire:
path: Final = tmp_path / "fallback.yaml"
path.write_text(yaml.safe_dump({"model_list": [{"model_name": alias, "litellm_params": {"model": "openai/" + upstream, "api_key": "synthetic-routing-key", "api_base": wire.url + "/v1"}} for alias, upstream in (("primary", "primary-wire"), ("fallback", "fallback-wire"), ("unrelated", "unrelated-wire"))], "router_settings": {"num_retries": 0, "disable_cooldowns": True, "fallbacks": [{"primary": ["fallback"]}]}}))
loaded: Final = yaml.safe_load(path.read_text())
router: Final = Router(model_list=loaded["model_list"], **loaded["router_settings"])
try:
result: Final = router.completion(model="primary", messages=[{"role": "user", "content": "fallback control"}])
assert result.id == "response-fallback-wire"
assert result.choices[0].message.content == "served fallback-wire"
assert result.choices[0].finish_reason == "stop"
assert result.usage.prompt_tokens == 11 and result.usage.completion_tokens == 4
assert tuple(json.loads(request.body)["model"] for request in wire.drain()) == ("primary-wire", "fallback-wire")
control: Final = router.completion(model="unrelated", messages=[{"role": "user", "content": "independent route"}])
assert control.id == "response-unrelated-wire"
assert tuple(json.loads(request.body)["model"] for request in wire.drain()) == ("unrelated-wire",)
finally:
router.reset()
@pytest.mark.covers("other.routing.alias_update.persisted_target_changes_only_selected_route")
def test_saved_deployment_target_update_changes_wire_and_preserves_control(gateway: Gateway) -> None:
with httpx.Client(base_url=gateway.upstream_url, timeout=5, trust_env=False) as upstream, gateway.scenario() as scenario:
prefix: Final = "target-" + uuid.uuid4().hex
model: Final = scenario.model(model="openai/" + prefix + "-first", input_cost_per_token=0, output_cost_per_token=0)
other: Final = scenario.model(model="openai/" + prefix + "-control", input_cost_per_token=0, output_cost_per_token=0)
target: Final = next(entry for entry in gateway.get("/model/info")["data"] if entry["model_name"] == model)
for generation, suffix in enumerate(("first", "second")):
if generation:
response: Final = gateway.request("PATCH", f"/model/{target['model_info']['id']}/update", {"litellm_params": {"model": "openai/" + prefix + "-second"}})
assert response.status_code == 200, response.text
upstream.get("/__observations").raise_for_status()
for alias in (model, other):
assert gateway.chat(alias, text=f"{prefix} generation {generation}")["usage"]["total_tokens"] == 40
requests: Final = upstream.get("/__observations").json()["requests"]
assert [request["body"]["model"] for request in requests] == [prefix + "-" + suffix, prefix + "-control"]

View file

@ -0,0 +1,59 @@
import os
import uuid
from pathlib import Path
from typing import Final
from urllib.parse import urlsplit, urlunsplit
import httpx
import psycopg
import pytest
from psycopg import sql
from redis import Redis
from integration._support.client import Gateway, eventually
from integration._support.process import owned_proxy
from integration._support.redis_process import owned_redis
@pytest.mark.covers("other.routing.redis.owned_outage_recovers_serving_and_response_cache")
def test_owned_redis_outage_recovers_requests_and_real_response_cache(gateway: Gateway, tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None:
original: Final = os.environ["DATABASE_URL"]
identity: Final = "integration_recovery_" + uuid.uuid4().hex
parsed: Final = urlsplit(original)
database_url: Final = urlunsplit((parsed.scheme, parsed.netloc, "/" + identity, "", ""))
with psycopg.connect(original, autocommit=True) as admin:
admin.execute(sql.SQL("CREATE DATABASE {}").format(sql.Identifier(identity)))
try:
with owned_redis(tmp_path) as cache, monkeypatch.context() as environment:
environment.setenv("DATABASE_URL", database_url)
with owned_proxy(gateway, tmp_path, {"DATABASE_URL": database_url, "REDIS_HOST": cache.host, "REDIS_PORT": str(cache.port), "REDIS_CIRCUIT_BREAKER_RECOVERY_TIMEOUT": "1"}) as candidate, candidate.scenario() as scenario, httpx.Client(base_url=gateway.upstream_url, timeout=5, trust_env=False) as upstream:
model: Final = scenario.model()
key: Final = scenario.key(models=[model])
for generation in ("before", "after"):
with Redis(host=cache.host, port=cache.port, socket_timeout=1) as client:
eventually(client.ping, bool)
eventually(lambda: client.pubsub_numsub("litellm_proxy.auth_cache_invalidation")[0][1], lambda count: count >= 1, seconds=8)
upstream.get("/__observations").raise_for_status()
first: Final = candidate.chat(model, key=key, text=identity + generation)
second: Final = candidate.chat(model, key=key, text=identity + generation)
assert first["id"] == second["id"]
assert first["choices"] == second["choices"] and first["usage"]["total_tokens"] == 40
assert len(upstream.get("/__observations").json()["requests"]) == 1
with Redis(host=cache.host, port=cache.port, socket_timeout=1) as client:
eventually(
lambda first=first: tuple(client.get(name) for name in client.scan_iter() if client.type(name) == b"string"),
lambda values, first=first: any(str(first["id"]).encode() in value for value in values if value is not None),
seconds=10,
)
if generation == "before":
cache.stop()
upstream.get("/__observations").raise_for_status()
during: Final = candidate.chat(model, key=key, text=identity + "during")
assert during["usage"]["total_tokens"] == 40
assert len(upstream.get("/__observations").json()["requests"]) == 1
cache.start()
with psycopg.connect(database_url) as fresh:
assert fresh.execute('SELECT count(*) FROM "LiteLLM_VerificationToken"').fetchone()[0] >= 1
finally:
admin.execute(sql.SQL("DROP DATABASE {}").format(sql.Identifier(identity)))
assert admin.execute("SELECT datname FROM pg_database WHERE datname=%s", (identity,)).fetchall() == []

View file

@ -0,0 +1,234 @@
import uuid
from contextlib import ExitStack
from hashlib import sha256
from typing import Final
import httpx
import pytest
from hypothesis import strategies as st
from hypothesis.stateful import RuleBasedStateMachine, rule, run_state_machine_as_test
from integration._support.client import Gateway, eventually
from integration._support.database import read_rows
from integration._support.generation import LIFECYCLE_SETTINGS, bounded_http_requests
@pytest.mark.covers("quota_management.response_cache.generated_sequences_preserve_content_and_accounting")
@pytest.mark.timeout(180)
def test_generated_cache_sequences_preserve_content_usage_and_zero_hit_cost(gateway: Gateway) -> None:
class CacheRequests(RuleBasedStateMachine):
def __init__(self) -> None:
super().__init__()
self.resources = ExitStack()
try:
self.scenario = self.resources.enter_context(gateway.scenario())
self.upstream = self.resources.enter_context(
httpx.Client(base_url=gateway.upstream_url, timeout=5, trust_env=False)
)
self.model = self.scenario.model(input_cost_per_token=0.001, output_cost_per_token=0.002)
self.key = self.scenario.key(models=[self.model])
self.prefix = uuid.uuid4().hex
self.seen: frozenset[int] = frozenset()
self.requests = 0
self.paid = 0
self.failed = False
self.identities: dict[int, str] = {}
except BaseException:
with budget.cleanup():
self.resources.close()
raise
@rule(marker=st.integers(min_value=0, max_value=2))
def request(self, marker: int) -> None:
try:
self.perform_request(marker)
except BaseException:
self.failed = True
raise
def perform_request(self, marker: int) -> None:
self.upstream.get("/__observations").raise_for_status()
response: Final = gateway.request(
"POST",
"/v1/chat/completions",
{
"model": self.model,
"messages": [{"role": "user", "content": f"{self.prefix}-{marker}"}],
},
key=self.key,
)
assert response.status_code == 200, response.text
self.requests += 1
body: Final = response.json()
assert (
body["choices"][0]["message"]["content"]
== "Hello! This is a mock response from the fake OpenAI endpoint."
)
assert body["usage"]["total_tokens"] == 40
observed: Final = self.upstream.get("/__observations").json()["requests"]
expected_calls: Final = 0 if marker in self.seen else 1
assert len(observed) == expected_calls, observed
if marker not in self.seen:
assert float(response.headers["x-litellm-response-cost"]) == pytest.approx(0.06)
if marker in self.identities:
assert body["id"] == self.identities[marker]
else:
assert body["id"] not in self.identities.values()
self.identities = {**self.identities, marker: body["id"]}
self.paid += expected_calls
self.seen = self.seen.union((marker,))
def teardown(self) -> None:
try:
if self.requests and not self.failed:
rows: Final = eventually(
lambda: read_rows(
"SELECT request_id, spend, cache_hit, prompt_tokens, "
'completion_tokens FROM "LiteLLM_SpendLogs" WHERE api_key=%s',
(sha256(self.key.encode()).hexdigest(),),
),
lambda values: len(values) == self.requests,
seconds=70,
)
assert len({row["request_id"] for row in rows}) == self.requests
assert sum(float(row["spend"]) for row in rows) == pytest.approx(self.paid * 0.06)
assert sum(row["cache_hit"] == "True" for row in rows) == self.requests - self.paid
for row in rows:
assert row["prompt_tokens"] == 20 and row["completion_tokens"] == 20
if row["cache_hit"] == "True":
assert float(row["spend"]) == 0 and "_cache_hit" in row["request_id"]
assert any(
row["request_id"].startswith(identity + "_cache_hit")
for identity in self.identities.values()
)
else:
assert row["request_id"] in self.identities.values()
assert float(row["spend"]) == pytest.approx(0.06)
finally:
with budget.cleanup():
self.resources.close()
with bounded_http_requests((gateway,), limit=2000) as budget:
run_state_machine_as_test(CacheRequests, settings=LIFECYCLE_SETTINGS)
@pytest.mark.covers("quota_management.response_cache.repeated_hits_preserve_identity_and_single_charge")
def test_repeated_hits_keep_response_identity_and_create_distinct_zero_cost_rows(gateway: Gateway) -> None:
with (
gateway.scenario() as scenario,
httpx.Client(base_url=gateway.upstream_url, timeout=5, trust_env=False) as upstream,
):
model: Final = scenario.model(input_cost_per_token=0.001, output_cost_per_token=0.002)
key: Final = scenario.key(models=[model])
prompt: Final = f"repeated cache {uuid.uuid4().hex}"
upstream.get("/__observations").raise_for_status()
results: Final = tuple(gateway.chat(model, key=key, text=prompt) for _ in range(3))
assert len(upstream.get("/__observations").json()["requests"]) == 1
assert len({result["id"] for result in results}) == 1
for result in results:
assert (
result["choices"][0]["message"]["content"]
== "Hello! This is a mock response from the fake OpenAI endpoint."
)
assert result["usage"]["total_tokens"] == 40
rows: Final = eventually(
lambda: read_rows(
'SELECT request_id, spend, cache_hit FROM "LiteLLM_SpendLogs" WHERE api_key=%s',
(sha256(key.encode()).hexdigest(),),
),
lambda values: len(values) == 3,
seconds=70,
)
assert len({row["request_id"] for row in rows}) == 3
assert sorted(float(row["spend"]) for row in rows) == [0, 0, 0.06]
for row in rows:
if row["cache_hit"] == "True":
assert float(row["spend"]) == 0
assert row["request_id"].startswith(results[0]["id"] + "_cache_hit")
else:
assert row["request_id"] == results[0]["id"] and float(row["spend"]) == pytest.approx(0.06)
@pytest.mark.covers("quota_management.budget.key.boundary_blocks_before_provider_and_reset_restores")
def test_key_budget_at_boundary_blocks_provider_then_explicit_reset_restores(gateway: Gateway) -> None:
with (
gateway.scenario() as scenario,
httpx.Client(base_url=gateway.upstream_url, timeout=5, trust_env=False) as upstream,
):
model: Final = scenario.model(input_cost_per_token=0.001, output_cost_per_token=0.002)
key: Final = scenario.key(models=[model], max_budget=0.06)
control: Final = scenario.key(models=[model])
first: Final = gateway.chat(model, key=key, text=f"budget {uuid.uuid4().hex}")
assert first["usage"]["total_tokens"] == 40
digest: Final = sha256(key.encode()).hexdigest()
spent: Final = eventually(
lambda: read_rows('SELECT spend FROM "LiteLLM_VerificationToken" WHERE token=%s', (digest,)),
lambda values: len(values) == 1 and float(values[0]["spend"]) >= 0.06,
seconds=70,
)
assert float(spent[0]["spend"]) == pytest.approx(0.06)
upstream.get("/__observations").raise_for_status()
denied: Final = gateway.request(
"POST",
"/v1/chat/completions",
{"model": model, "messages": [{"role": "user", "content": f"over budget {uuid.uuid4().hex}"}]},
key=key,
)
assert denied.status_code == 429 and denied.json()["error"]["type"] == "budget_exceeded", denied.text
assert upstream.get("/__observations").json()["requests"] == []
assert gateway.chat(model, key=control, text=f"control {uuid.uuid4().hex}")["usage"]["total_tokens"] == 40
gateway.post("/key/update", {"key": key, "spend": 0})
assert read_rows('SELECT spend, max_budget FROM "LiteLLM_VerificationToken" WHERE token=%s', (digest,)) == [
{"spend": 0.0, "max_budget": 0.06}
]
assert gateway.chat(model, key=key, text=f"reset {uuid.uuid4().hex}")["usage"]["total_tokens"] == 40
eventually(
lambda: read_rows('SELECT spend FROM "LiteLLM_VerificationToken" WHERE token=%s', (digest,)),
lambda values: len(values) == 1 and float(values[0]["spend"]) >= 0.06,
seconds=70,
)
upstream.get("/__observations").raise_for_status()
denied_again: Final = gateway.request(
"POST",
"/v1/chat/completions",
{"model": model, "messages": [{"role": "user", "content": f"boundary again {uuid.uuid4().hex}"}]},
key=key,
)
assert denied_again.status_code == 429 and denied_again.json()["error"]["type"] == "budget_exceeded", (
denied_again.text
)
assert upstream.get("/__observations").json()["requests"] == []
@pytest.mark.covers("quota_management.response_cache.system_messages_partition_cache_identity")
def test_different_system_messages_do_not_share_a_cached_response(gateway: Gateway) -> None:
with (
gateway.scenario() as scenario,
httpx.Client(base_url=gateway.upstream_url, timeout=5, trust_env=False) as upstream,
):
model: Final = scenario.model()
prompt: Final = uuid.uuid4().hex
identities: dict[str, str] = {}
for system, expected_calls in (("first policy", 1), ("second policy", 1), ("first policy", 0)):
upstream.get("/__observations").raise_for_status()
response: Final = gateway.request(
"POST",
"/v1/chat/completions",
{
"model": model,
"messages": [{"role": "system", "content": system}, {"role": "user", "content": prompt}],
},
)
assert response.status_code == 200 and response.json()["usage"]["total_tokens"] == 40, response.text
calls: Final = upstream.get("/__observations").json()["requests"]
assert len(calls) == expected_calls
if system in identities:
assert response.json()["id"] == identities[system]
else:
assert response.json()["id"] not in identities.values()
identities = {**identities, system: response.json()["id"]}
if calls:
assert calls[0]["body"]["messages"] == [
{"role": "system", "content": system},
{"role": "user", "content": prompt},
]

View file

@ -0,0 +1,149 @@
import asyncio
import json
import threading
import uuid
from typing import Final
import pytest
from hypothesis import Phase, example, given, settings, strategies as st
from openai import OpenAI
from integration._support.client import Gateway, eventually
from integration._support.database import read_rows
from integration._support.wire import Reply, wire_server
def frame(identity: str, delta: dict, *, finish: str | None = None) -> bytes:
value: Final = {"id": identity, "object": "chat.completion.chunk", "created": 1, "model": "gpt-4o-mini", "choices": [{"index": 0, "delta": delta, "finish_reason": finish}]}
return b"data: " + json.dumps(value, ensure_ascii=False).encode() + b"\n\n"
def text_stream(identity: str) -> tuple[bytes, ...]:
usage: Final = {"id": identity, "object": "chat.completion.chunk", "created": 1, "model": "gpt-4o-mini", "choices": [], "usage": {"prompt_tokens": 11, "completion_tokens": 4, "total_tokens": 15}}
return (frame(identity, {"role": "assistant", "content": "Hello "}), frame(identity, {"content": "雪 café"}), frame(identity, {}, finish="stop"), b"data: " + json.dumps(usage).encode() + b"\n\n", b"data: [DONE]\n\n")
@pytest.mark.covers("other.streaming.byte_partitions.preserve_text_identity_and_usage")
def test_generated_tcp_partitions_preserve_unicode_text_identity_and_final_usage() -> None:
import litellm
body: Final = b"".join(text_stream("stream-partition-control"))
@settings(max_examples=20, deadline=None, database=None, phases=(Phase.explicit, Phase.generate, Phase.shrink))
@example(cuts=tuple(range(1, len(body))))
@example(cuts=())
@given(cuts=st.lists(st.integers(min_value=1, max_value=len(body) - 1), max_size=35, unique=True).map(tuple))
def check(cuts: tuple[int, ...]) -> None:
boundaries: Final = (0, *sorted(cuts), len(body))
pieces: Final = tuple(body[left:right] for left, right in zip(boundaries, boundaries[1:]))
with wire_server(lambda request: Reply(content_type="text/event-stream", chunks=pieces)) as wire:
stream: Final = litellm.completion(model="openai/gpt-4o-mini", api_base=wire.url + "/v1", api_key="synthetic-stream-key", messages=[{"role": "user", "content": "partition control"}], stream=True, stream_options={"include_usage": True}, timeout=5, num_retries=0)
try:
chunks: Final = tuple(stream)
finally:
asyncio.run(stream.aclose())
assert "".join(choice.delta.content or "" for chunk in chunks for choice in chunk.choices) == "Hello 雪 café"
assert {chunk.id for chunk in chunks} == {"stream-partition-control"}
assert [choice.finish_reason for chunk in chunks for choice in chunk.choices if choice.finish_reason] == ["stop"]
usages: Final = tuple(chunk.usage for chunk in chunks if getattr(chunk, "usage", None) is not None)
assert len(usages) == 1
assert usages[0].prompt_tokens == 11 and usages[0].completion_tokens == 4
assert len(wire.drain()) == 1
check()
@pytest.mark.covers("other.streaming.tools.fragmented_calls_keep_independent_arguments")
def test_fragmented_tool_names_and_arguments_keep_each_call_identity() -> None:
import litellm
identity: Final = "stream-tools-control"
deltas: Final = (
{"role": "assistant", "tool_calls": [{"index": 0, "id": "call-add", "type": "function", "function": {"name": "ad", "arguments": ""}}, {"index": 1, "id": "call-multiply", "type": "function", "function": {"name": "multi", "arguments": ""}}]},
{"tool_calls": [{"index": 1, "function": {"name": "ply", "arguments": '{"x":3,'}}, {"index": 0, "function": {"arguments": '{"x":1,'}}]},
{"tool_calls": [{"index": 0, "function": {"name": "d", "arguments": '"y":2}'}}, {"index": 1, "function": {"arguments": '"y":4}'}}]},
)
frames: Final = (*tuple(frame(identity, delta) for delta in deltas), frame(identity, {}, finish="tool_calls"), b"data: [DONE]\n\n")
with wire_server(lambda request: Reply(content_type="text/event-stream", chunks=frames)) as wire:
stream: Final = litellm.completion(model="openai/gpt-4o-mini", api_base=wire.url + "/v1", api_key="synthetic-stream-key", messages=[{"role": "user", "content": "tool control"}], stream=True, timeout=5, num_retries=0)
try:
chunks: Final = tuple(stream)
finally:
asyncio.run(stream.aclose())
events: Final = tuple((choice.index, tool) for chunk in chunks for choice in chunk.choices for tool in (choice.delta.tool_calls or ()))
for index, name, call_id, arguments in ((0, "add", "call-add", {"x": 1, "y": 2}), (1, "multiply", "call-multiply", {"x": 3, "y": 4})):
selected: Final = tuple(tool for choice, tool in events if (choice, tool.index) == (0, index))
assert "".join(tool.id or "" for tool in selected) == call_id
assert "".join(tool.function.name or "" for tool in selected) == name
assert json.loads("".join(tool.function.arguments or "" for tool in selected)) == arguments
assert {tool.index for _, tool in events} == {0, 1}
assert [choice.finish_reason for chunk in chunks for choice in chunk.choices if choice.finish_reason] == ["tool_calls"]
assert len(wire.drain()) == 1
@pytest.mark.covers("other.streaming.usage.client_visibility_preserves_persisted_accounting")
def test_proxy_stream_usage_visibility_keeps_exact_persisted_charge(gateway: Gateway) -> None:
with gateway.scenario() as scenario:
for include in (None, False, True):
identity: Final = "stream-usage-" + uuid.uuid4().hex
with wire_server(lambda request, identity=identity: Reply(content_type="text/event-stream", chunks=text_stream(identity))) as wire:
model: Final = scenario.model(api_base=wire.url + "/v1", input_cost_per_token=0.001, output_cost_per_token=0.002)
with OpenAI(api_key=gateway.key, base_url=str(gateway.client.base_url), timeout=5, max_retries=0) as client:
stream: Final = client.chat.completions.create(model=model, messages=[{"role": "user", "content": identity}], stream=True, **({} if include is None else {"stream_options": {"include_usage": include}}))
with stream:
chunks: Final = tuple(stream)
assert "".join(choice.delta.content or "" for chunk in chunks for choice in chunk.choices) == "Hello 雪 café"
assert {chunk.id for chunk in chunks} == {identity}
usages: Final = tuple(chunk.usage for chunk in chunks if chunk.usage is not None)
assert len(usages) == (1 if include else 0)
if include:
assert usages[0].prompt_tokens == 11 and usages[0].completion_tokens == 4
requests: Final = wire.drain()
assert len(requests) == 1
assert json.loads(requests[0].body)["stream_options"]["include_usage"] is True
rows: Final = eventually(lambda identity=identity: read_rows('SELECT spend, prompt_tokens, completion_tokens FROM "LiteLLM_SpendLogs" WHERE request_id=%s', (identity,)), lambda values: len(values) == 1, seconds=70)
assert rows[0]["prompt_tokens"] == 11 and rows[0]["completion_tokens"] == 4
assert float(rows[0]["spend"]) == pytest.approx(0.019)
@pytest.mark.covers("other.streaming.failure.truncated_transport_raises_and_control_recovers")
def test_truncated_http_stream_is_an_error_and_next_stream_succeeds() -> None:
import litellm
for truncated in (True, False):
with wire_server(lambda request, truncated=truncated: Reply(content_type="text/event-stream", chunks=text_stream("stream-truncated"), abort_after=1 if truncated else None)) as wire:
stream: Final = litellm.completion(model="openai/gpt-4o-mini", api_base=wire.url + "/v1", api_key="synthetic-stream-key", messages=[{"role": "user", "content": "truncation control"}], stream=True, timeout=5, num_retries=0)
try:
if truncated:
with pytest.raises(litellm.exceptions.MidStreamFallbackError, match="incomplete chunked read") as failure:
tuple(stream)
assert isinstance(failure.value.original_exception, litellm.APIConnectionError)
assert failure.value.generated_content == "Hello "
assert failure.value.is_pre_first_chunk is False
else:
chunks: Final = tuple(stream)
assert "".join(choice.delta.content or "" for chunk in chunks for choice in chunk.choices) == "Hello 雪 café"
assert any(choice.finish_reason == "stop" for chunk in chunks for choice in chunk.choices)
finally:
asyncio.run(stream.aclose())
assert len(wire.drain()) == 1
@pytest.mark.covers("other.streaming.cancellation.closes_actual_provider_connection")
def test_client_cancellation_releases_the_actual_provider_connection() -> None:
import litellm
gate: Final = threading.Event()
frames: Final = (frame("stream-cancel", {"role": "assistant", "content": "first"}), b":" + b"x" * 4_000_000 + b"\n\n", b"data: [DONE]\n\n")
with wire_server(lambda request: Reply(content_type="text/event-stream", chunks=frames, gate_after_first=gate)) as wire:
stream: Final = litellm.completion(model="openai/gpt-4o-mini", api_base=wire.url + "/v1", api_key="synthetic-stream-key", messages=[{"role": "user", "content": "cancellation control"}], stream=True, timeout=5, num_retries=0)
try:
first: Final = next(stream)
assert first.choices[0].delta.content == "first"
finally:
try:
asyncio.run(stream.aclose())
finally:
gate.set()
assert wire.disconnected.get(timeout=5) == "/v1/chat/completions"
assert len(wire.drain()) == 1

View file

@ -4287,3 +4287,36 @@ def test_system_string_after_a_developer_message_stays_in_input_in_client_order(
assert instructions is None
assert [item["role"] for item in input_items] == ["developer", "system", "user"]
assert input_items[1] == _system_input_item("Be brief.")
def test_map_optional_params_verbosity_merges_into_text():
"""Chat verbosity must land on Responses text.verbosity alongside text.format regardless of key order."""
from litellm.completion_extras.litellm_responses_transformation.transformation import (
LiteLLMResponsesTransformationHandler,
)
from litellm.types.llms.openai import ResponsesAPIOptionalRequestParams
handler: Final = LiteLLMResponsesTransformationHandler()
responses_api_request = ResponsesAPIOptionalRequestParams()
handler._map_optional_params_to_responses_api_request(
{"verbosity": "low", "response_format": {"type": "json_object"}},
responses_api_request,
)
assert responses_api_request["text"]["verbosity"] == "low"
assert responses_api_request["text"]["format"]["type"] == "json_object"
reversed_request = ResponsesAPIOptionalRequestParams()
handler._map_optional_params_to_responses_api_request(
{"response_format": {"type": "json_object"}, "verbosity": "low"},
reversed_request,
)
assert reversed_request["text"]["verbosity"] == "low"
assert reversed_request["text"]["format"]["type"] == "json_object"
verbosity_only_request = ResponsesAPIOptionalRequestParams()
handler._map_optional_params_to_responses_api_request(
{"verbosity": "low"},
verbosity_only_request,
)
assert verbosity_only_request["text"] == {"verbosity": "low"}

View file

@ -257,6 +257,18 @@ class TestBedrockMantleConfig:
assert "temperature" in params
assert "stream" in params
assert "max_tokens" in params
assert "verbosity" not in params
def test_verbosity_passes_through_for_gpt_5_models(self):
cfg = BedrockMantleChatConfig()
assert "verbosity" in cfg.get_supported_openai_params("openai.gpt-5.6-sol")
optional_params = litellm.get_optional_params(
model="openai.gpt-5.6-sol",
custom_llm_provider="bedrock_mantle",
verbosity="low",
drop_params=False,
)
assert optional_params["verbosity"] == "low"
class TestBedrockMantleChatAuth:

View file

@ -23,7 +23,7 @@ from litellm.proxy._experimental.mcp_server.outbound_credentials.adapter import
to_subject,
validate_static_credential,
)
from litellm.proxy._experimental.mcp_server.outbound_credentials.result import Ok
from litellm.proxy._experimental.mcp_server.outbound_credentials.result import Error, Ok
from litellm.proxy._experimental.mcp_server.outbound_credentials.types import (
ApiKeyConfig,
AuthorizationCodeConfig,
@ -59,6 +59,22 @@ def test_static_credential_preserves_supported_api_key_and_raw_headers(
assert isinstance(result, Ok)
@pytest.mark.parametrize("auth_type,headers,static_header_names,expected", [
(MCPAuth.api_key, {"apikey": "static-key"}, ("apikey",), Ok),
(MCPAuth.api_key, {"apikey": "static-key", "X-API-Key": ""}, ("apikey",), Ok),
(MCPAuth.api_key, {"apikey": ""}, ("apikey",), Error),
(MCPAuth.api_key, {"apikey": "static-key"}, (), Error),
(MCPAuth.api_key, {"apikey": "static-key"}, ("X-Tenant",), Error),
(MCPAuth.bearer_token, {"apikey": "static-key"}, ("apikey",), Error),
(MCPAuth.token, {"apikey": "static-key"}, ("apikey",), Error),
])
def test_static_credential_counts_api_key_static_headers_only(
auth_type: MCPAuthType, headers: dict[str, str], static_header_names: tuple[str, ...], expected: type,
) -> None:
result: Final = validate_static_credential(auth_type, headers, static_header_names=static_header_names)
assert isinstance(result, expected)
def _server(**kwargs) -> MCPServer:
return MCPServer(server_id="s", name="n", transport=MCPTransport.http, **kwargs)

View file

@ -13652,6 +13652,28 @@ class TestProtectedCredentialPreparation:
assert client._credential_slot == "X-Custom"
assert await client.discovery_auth_fingerprint()
@pytest.mark.asyncio
@pytest.mark.parametrize("static_headers,accepted", [
({"apikey": "static-key"}, True),
({"apikey": ""}, False),
({"X-Tenant": "tenant"}, True),
])
async def test_api_key_carried_by_static_header_passes_fail_closed_check(
self, static_headers: dict[str, str], accepted: bool
) -> None:
server: Final = MCPServer(
server_id="static-slot", name="static-slot", url="https://upstream.example/mcp",
transport=MCPTransport.http, auth_type=MCPAuth.api_key, static_headers=static_headers,
)
if not accepted:
with pytest.raises(HTTPException) as exc:
await MCPServerManager()._create_mcp_client(server, extra_headers=dict(static_headers))
assert exc.value.status_code == 500
return
client: Final = await MCPServerManager()._create_mcp_client(server, extra_headers=dict(static_headers))
request: Final = await client.prepare_request_auth()
assert all(request.headers[name] == value for name, value in static_headers.items())
@pytest.mark.asyncio
@pytest.mark.parametrize("static,forwarded,caller", [
({"X-API-Key": "static"}, {"x-api-key": "forwarded"}, None),

View file

@ -133,6 +133,26 @@ async def test_static_auth_uses_configured_custom_header(
assert destination.call_count == 0
@pytest.mark.asyncio
@pytest.mark.parametrize("credential", ["static-key", ""])
async def test_static_auth_accepts_api_key_carried_by_static_header(
respx_mock: MockRouter, monkeypatch: pytest.MonkeyPatch, credential: str,
) -> None:
monkeypatch.setenv("DISABLE_AIOHTTP_TRANSPORT", "True")
tool: Final = create_tool_function(
"/echo", "get", {}, "https://upstream.example", headers={"apikey": credential}, auth_type=MCPAuth.api_key,
)
destination: Final = respx_mock.get("https://upstream.example/echo").respond(200, text="authenticated")
if credential:
assert await tool() == "authenticated"
assert destination.calls.last.request.headers["apikey"] == credential
assert "x-api-key" not in destination.calls.last.request.headers
else:
with pytest.raises(HTTPException, match="requires a usable upstream credential"):
await tool()
assert destination.call_count == 0
@pytest.mark.asyncio
@pytest.mark.parametrize("auth_type,resolved", [
(MCPAuth.none, None),

View file

@ -972,11 +972,6 @@
"count": 2
}
},
"src/app/(dashboard)/projects/_components/ProjectKeysSection.tsx": {
"react-hooks/set-state-in-effect": {
"count": 1
}
},
"src/app/(dashboard)/projects/_components/ProjectModals/ProjectBaseForm.tsx": {
"react-hooks/set-state-in-effect": {
"count": 2

View file

@ -1,10 +1,23 @@
import { QueryClient, QueryClientProvider } from "@tanstack/react-query";
import { act, render, screen, waitFor } from "@testing-library/react";
import { act, fireEvent, render, screen, waitFor } from "@testing-library/react";
import { NuqsTestingAdapter, type UrlUpdateEvent } from "nuqs/adapters/testing";
import React from "react";
import { beforeEach, describe, expect, it, vi } from "vitest";
import type OrganizationsTableComponent from "./OrganizationsTable";
import type OrganizationInfoViewComponent from "@/components/organization/organization_view";
import type { OrganizationListFilters } from "@/app/(dashboard)/hooks/organizations/useOrganizations";
const useOrganizationsSpy = vi.hoisted(() => vi.fn<(filters?: OrganizationListFilters) => void>());
vi.mock("@/app/(dashboard)/hooks/organizations/useOrganizations", async (importOriginal) => {
const actual = await importOriginal<typeof import("@/app/(dashboard)/hooks/organizations/useOrganizations")>();
return {
...actual,
useOrganizations: (filters?: OrganizationListFilters) => {
useOrganizationsSpy(filters);
return actual.useOrganizations(filters);
},
};
});
vi.mock("@/components/vector_store_management/VectorStoreSelector", () => ({
__esModule: true,
@ -79,10 +92,13 @@ const renderPanel = ({ premiumUser = true, searchParams = "" }: RenderPanelOptio
const expectQueryString = (queryString: string) =>
waitFor(() => expect(onUrlUpdate).toHaveBeenLastCalledWith(expect.objectContaining({ queryString })));
const lastSearchParams = () => onUrlUpdate.mock.calls.at(-1)?.[0].searchParams;
beforeEach(() => {
capturedTableProps = null;
mockOrgInfoView.mockClear();
onUrlUpdate.mockClear();
useOrganizationsSpy.mockClear();
});
describe("OrganizationsPanel", () => {
@ -123,9 +139,7 @@ describe("OrganizationsPanel - org detail deep link (?org=)", () => {
it("opens the org detail directly from a ?org= deep link", () => {
renderPanel({ searchParams: "?org=org-from-url" });
expect(mockOrgInfoView).toHaveBeenLastCalledWith(
expect.objectContaining({ organizationId: "org-from-url", editOrg: false }),
);
expect(mockOrgInfoView).toHaveBeenLastCalledWith(expect.objectContaining({ organizationId: "org-from-url" }));
expect(screen.queryByTestId("organizations-table")).not.toBeInTheDocument();
});
@ -139,23 +153,24 @@ describe("OrganizationsPanel - org detail deep link (?org=)", () => {
expect(screen.getByTestId("organizations-table")).toBeInTheDocument();
});
it("the edit action opens the detail in edit mode with ?org= set", async () => {
it("the edit action pushes ?org= with ?org_tab=settings in one history entry", async () => {
renderPanel();
act(() => capturedTableProps?.onEditClick("org-edit"));
await expectQueryString("?org=org-edit");
expect(mockOrgInfoView).toHaveBeenLastCalledWith(
expect.objectContaining({ organizationId: "org-edit", editOrg: true }),
await expectQueryString("?org=org-edit&org_tab=settings");
expect(onUrlUpdate).toHaveBeenCalledTimes(1);
expect(onUrlUpdate).toHaveBeenLastCalledWith(
expect.objectContaining({ options: expect.objectContaining({ history: "push" }) }),
);
expect(mockOrgInfoView).toHaveBeenLastCalledWith(expect.objectContaining({ organizationId: "org-edit" }));
});
it("a plain row click after leaving an edit view via browser history does not reopen in edit mode", async () => {
it("a plain row click after leaving an edit view via browser history opens the detail without the settings tab", async () => {
const { navigate } = renderPanel();
act(() => capturedTableProps?.onEditClick("org-edit"));
await expectQueryString("?org=org-edit");
expect(mockOrgInfoView).toHaveBeenLastCalledWith(expect.objectContaining({ editOrg: true }));
await expectQueryString("?org=org-edit&org_tab=settings");
navigate("");
expect(screen.getByTestId("organizations-table")).toBeInTheDocument();
@ -163,8 +178,90 @@ describe("OrganizationsPanel - org detail deep link (?org=)", () => {
act(() => capturedTableProps?.onOrganizationClick("org-plain"));
await expectQueryString("?org=org-plain");
expect(mockOrgInfoView).toHaveBeenLastCalledWith(
expect.objectContaining({ organizationId: "org-plain", editOrg: false }),
expect(mockOrgInfoView).toHaveBeenLastCalledWith(expect.objectContaining({ organizationId: "org-plain" }));
});
it("a row click drops a leftover ?org_tab= so the detail opens on its default tab", async () => {
renderPanel({ searchParams: "?org_tab=settings" });
act(() => capturedTableProps?.onOrganizationClick("org-plain"));
await expectQueryString("?org=org-plain");
});
it("closing the org detail keeps the list's search, filter, sort and page in the URL", async () => {
renderPanel({
searchParams:
"?org_search=Acme&filter_org_id=org-7&sort_by=spend&sort_order=asc&page=2&org=org-x&org_tab=members",
});
act(() => mockOrgInfoView.mock.calls.at(-1)?.[0].onClose());
await expectQueryString("?org_search=Acme&filter_org_id=org-7&sort_by=spend&sort_order=asc&page=2");
expect(onUrlUpdate).toHaveBeenCalledTimes(1);
expect(screen.getByPlaceholderText("Search by Organization Name")).toHaveValue("Acme");
expect(useOrganizationsSpy).toHaveBeenLastCalledWith({ org_id: "org-7", org_alias: "Acme" });
});
it("closing the org detail drops ?org_tab= together with ?org=", async () => {
renderPanel({ searchParams: "?org=org-from-url&org_tab=members" });
act(() => mockOrgInfoView.mock.calls.at(-1)?.[0].onClose());
await expectQueryString("");
expect(onUrlUpdate).toHaveBeenCalledTimes(1);
expect(onUrlUpdate).toHaveBeenLastCalledWith(
expect.objectContaining({ options: expect.objectContaining({ history: "push" }) }),
);
});
});
describe("OrganizationsPanel - list filters in the URL", () => {
it("restores the name search and org ID filter from the URL and fetches with both", () => {
renderPanel({ searchParams: "?org_search=Acme&filter_org_id=org-7" });
expect(screen.getByPlaceholderText("Search by Organization Name")).toHaveValue("Acme");
expect(screen.getByPlaceholderText("Search by Organization ID")).toHaveValue("org-7");
expect(useOrganizationsSpy).toHaveBeenLastCalledWith({ org_id: "org-7", org_alias: "Acme" });
expect(capturedTableProps?.searchActive).toBe(true);
});
it("keeps the org ID filter panel collapsed when the URL has no org ID filter", () => {
renderPanel({ searchParams: "?org_search=Acme" });
expect(screen.queryByPlaceholderText("Search by Organization ID")).not.toBeInTheDocument();
expect(useOrganizationsSpy).toHaveBeenLastCalledWith({ org_id: "", org_alias: "Acme" });
});
it("writes the name search to ?org_search= and returns the list to the first page", async () => {
renderPanel({ searchParams: "?page=3" });
fireEvent.change(screen.getByPlaceholderText("Search by Organization Name"), { target: { value: "Acme" } });
await waitFor(() => expect(lastSearchParams()?.get("org_search")).toBe("Acme"));
expect(lastSearchParams()?.has("page")).toBe(false);
expect(useOrganizationsSpy).toHaveBeenLastCalledWith({ org_id: "", org_alias: "Acme" });
});
it("writes the org ID filter to ?filter_org_id= and returns the list to the first page", async () => {
renderPanel({ searchParams: "?page=3" });
fireEvent.click(screen.getByRole("button", { name: "Filters" }));
fireEvent.change(screen.getByPlaceholderText("Search by Organization ID"), { target: { value: "org-9" } });
await waitFor(() => expect(lastSearchParams()?.get("filter_org_id")).toBe("org-9"));
expect(lastSearchParams()?.has("page")).toBe(false);
expect(useOrganizationsSpy).toHaveBeenLastCalledWith({ org_id: "org-9", org_alias: "" });
});
it("clears the search, the org ID filter and the page in one update on reset", async () => {
renderPanel({ searchParams: "?org_search=Acme&filter_org_id=org-7&page=2" });
fireEvent.click(screen.getByRole("button", { name: "Reset Filters" }));
await expectQueryString("");
expect(onUrlUpdate).toHaveBeenCalledTimes(1);
expect(useOrganizationsSpy).toHaveBeenLastCalledWith({ org_id: "", org_alias: "" });
expect(capturedTableProps?.searchActive).toBe(false);
});
});

View file

@ -2,16 +2,18 @@ import { organizationKeys, useOrganizations } from "@/app/(dashboard)/hooks/orga
import { useUserModels } from "@/app/(dashboard)/hooks/models/useModels";
import OrganizationFilters, { FilterState } from "@/app/(dashboard)/organizations/OrganizationFilters";
import { useQueryClient } from "@tanstack/react-query";
import { parseAsString, useQueryState } from "nuqs";
import { parseAsString, parseAsStringLiteral, useQueryStates } from "nuqs";
import React, { useState } from "react";
import DeleteResourceModal from "@/components/common_components/DeleteResourceModal";
import { toast } from "@/lib/toast";
import { organizationDeleteCall } from "@/components/networking";
import { OrgCreateDialog } from "@/components/organization/org-create/OrgCreateDialog";
import OrganizationInfoView from "@/components/organization/organization_view";
import { ORGANIZATION_TAB_URL_KEY, ORGANIZATION_TABS } from "@/components/organization/organizationTabs";
import { Button } from "@/components/ui/button";
import OrganizationsTable from "./OrganizationsTable";
import { organizationIdFilter, useOrganizationsTableState } from "./useOrganizationsTableState";
interface OrganizationsPanelProps {
userRole: string;
@ -19,15 +21,25 @@ interface OrganizationsPanelProps {
premiumUser: boolean;
}
const ORGANIZATION_DETAIL_STATE = {
org: parseAsString,
tab: parseAsStringLiteral(ORGANIZATION_TABS),
};
const ORGANIZATION_DETAIL_URL_KEYS = { tab: ORGANIZATION_TAB_URL_KEY };
const OrganizationsPanel: React.FC<OrganizationsPanelProps> = ({ userRole, accessToken, premiumUser }) => {
const [selectedOrgId, setSelectedOrgId] = useQueryState("org", parseAsString.withOptions({ history: "push" }));
const [editOrg, setEditOrg] = useState(false);
const [{ org: selectedOrgId }, setOrganizationDetail] = useQueryStates(ORGANIZATION_DETAIL_STATE, {
history: "push",
urlKeys: ORGANIZATION_DETAIL_URL_KEYS,
});
const tableState = useOrganizationsTableState();
const { setSearch, onColumnFiltersChange } = tableState;
const filters: FilterState = { org_id: organizationIdFilter(tableState), org_alias: tableState.search };
const [isDeleteModalOpen, setIsDeleteModalOpen] = useState(false);
const [orgToDelete, setOrgToDelete] = useState<string | null>(null);
const [isDeleting, setIsDeleting] = useState(false);
const [isOrgModalVisible, setIsOrgModalVisible] = useState(false);
const [showFilters, setShowFilters] = useState(false);
const [filters, setFilters] = useState<FilterState>({ org_id: "", org_alias: "" });
const [showFilters, setShowFilters] = useState(() => filters.org_id !== "");
const queryClient = useQueryClient();
const { data: organizations = [], isLoading } = useOrganizations({
@ -41,11 +53,16 @@ const OrganizationsPanel: React.FC<OrganizationsPanelProps> = ({ userRole, acces
const refetchOrganizations = () => queryClient.invalidateQueries({ queryKey: organizationKeys.lists() });
const handleFilterChange = (key: keyof FilterState, value: string) => {
setFilters((previousFilters) => ({ ...previousFilters, [key]: value }));
if (key === "org_alias") {
setSearch(value);
return;
}
onColumnFiltersChange(value ? [{ id: "org_id", value }] : []);
};
const handleFilterReset = () => {
setFilters({ org_id: "", org_alias: "" });
setSearch("");
onColumnFiltersChange([]);
};
const handleDelete = (orgId: string | null) => {
@ -108,15 +125,11 @@ const OrganizationsPanel: React.FC<OrganizationsPanelProps> = ({ userRole, acces
{selectedOrgId ? (
<OrganizationInfoView
organizationId={selectedOrgId}
onClose={() => {
void setSelectedOrgId(null);
setEditOrg(false);
}}
onClose={() => void setOrganizationDetail(null)}
accessToken={accessToken}
is_org_admin={true}
is_proxy_admin={userRole === "Admin"}
userModels={userModels}
editOrg={editOrg}
/>
) : (
<>
@ -133,14 +146,8 @@ const OrganizationsPanel: React.FC<OrganizationsPanelProps> = ({ userRole, acces
isLoading={isLoading}
userRole={userRole}
searchActive={searchActive}
onOrganizationClick={(organizationId) => {
setEditOrg(false);
void setSelectedOrgId(organizationId);
}}
onEditClick={(organizationId) => {
void setSelectedOrgId(organizationId);
setEditOrg(true);
}}
onOrganizationClick={(organizationId) => void setOrganizationDetail({ org: organizationId, tab: null })}
onEditClick={(organizationId) => void setOrganizationDetail({ org: organizationId, tab: "settings" })}
onDeleteClick={handleDelete}
/>
</>

View file

@ -1,7 +1,9 @@
import { render, screen, within } from "@testing-library/react";
import userEvent from "@testing-library/user-event";
import type { OnUrlUpdateFunction } from "nuqs/adapters/testing";
import React from "react";
import { describe, expect, it, vi } from "vitest";
import { describe, expect, it, vi, type Mock } from "vitest";
import { renderWithProviders, screen, waitFor, within } from "../../../../../tests/test-utils";
import { Organization } from "@/components/networking";
@ -26,6 +28,34 @@ const makeOrganization = (overrides: Partial<Organization> = {}): Organization =
...overrides,
});
const thirtyOrganizations = Array.from({ length: 30 }, (_, index) =>
makeOrganization({ organization_id: `org-${index}`, organization_alias: `Org ${index}` }),
);
const sortableOrganization = (alias: string, createdAt: string, spend: number): Organization => {
const overrides: Partial<Organization> = {
organization_id: `org-${alias.toLowerCase()}`,
organization_alias: alias,
created_at: createdAt,
spend,
};
return makeOrganization(overrides);
};
const sortableOrganizations = [
sortableOrganization("Mid", "2024-03-01T00:00:00Z", 5),
sortableOrganization("Zed", "2023-01-01T00:00:00Z", 1),
sortableOrganization("Ace", "2025-01-01T00:00:00Z", 3),
];
const bodyRowAliases = () =>
screen
.getAllByRole("row")
.slice(1)
.map((row) => ["Ace", "Mid", "Zed"].find((alias) => within(row).queryByText(alias) !== null));
const lastSearchParams = (onUrlUpdate: Mock<OnUrlUpdateFunction>) => onUrlUpdate.mock.calls.at(-1)?.[0].searchParams;
const baseProps = {
isLoading: false,
userRole: "Admin",
@ -37,7 +67,7 @@ const baseProps = {
describe("OrganizationsTable", () => {
it("renders every column header", () => {
render(<OrganizationsTable {...baseProps} organizations={[]} />);
renderWithProviders(<OrganizationsTable {...baseProps} organizations={[]} />);
for (const header of [
"Organization ID",
"Organization Name",
@ -55,7 +85,7 @@ describe("OrganizationsTable", () => {
it("opens the detail view when the organization ID cell is clicked", async () => {
const user = userEvent.setup();
const onOrganizationClick = vi.fn();
render(
renderWithProviders(
<OrganizationsTable
{...baseProps}
onOrganizationClick={onOrganizationClick}
@ -72,7 +102,7 @@ describe("OrganizationsTable", () => {
const user = userEvent.setup();
const onEditClick = vi.fn();
const onDeleteClick = vi.fn();
render(
renderWithProviders(
<OrganizationsTable
{...baseProps}
userRole="Admin"
@ -92,7 +122,7 @@ describe("OrganizationsTable", () => {
});
it("hides the row actions menu from non-admins", () => {
render(
renderWithProviders(
<OrganizationsTable
{...baseProps}
userRole="Internal User"
@ -104,7 +134,7 @@ describe("OrganizationsTable", () => {
});
it("sorts by created_at descending by default", () => {
render(
renderWithProviders(
<OrganizationsTable
{...baseProps}
organizations={[
@ -129,7 +159,7 @@ describe("OrganizationsTable", () => {
});
it("renders budget, limits, members, and models for a fully-populated organization", () => {
render(
renderWithProviders(
<OrganizationsTable
{...baseProps}
organizations={[
@ -151,7 +181,7 @@ describe("OrganizationsTable", () => {
});
it("shows Unlimited budget and All Proxy Models when unset", () => {
render(
renderWithProviders(
<OrganizationsTable
{...baseProps}
organizations={[makeOrganization({ organization_id: "org-empty", litellm_budget_table: {}, models: [] })]}
@ -166,7 +196,7 @@ describe("OrganizationsTable", () => {
});
it("renders a tpm/rpm limit of 0 as 0, never as Unlimited", () => {
render(
renderWithProviders(
<OrganizationsTable
{...baseProps}
organizations={[makeOrganization({ litellm_budget_table: { max_budget: null, tpm_limit: 0, rpm_limit: 0 } })]}
@ -180,7 +210,7 @@ describe("OrganizationsTable", () => {
});
it("renders loading skeletons instead of rows while loading", () => {
render(
renderWithProviders(
<OrganizationsTable
{...baseProps}
isLoading
@ -194,10 +224,8 @@ describe("OrganizationsTable", () => {
it("pages long lists client-side with the shared size selector and footer", async () => {
const user = userEvent.setup();
const organizations = Array.from({ length: 30 }, (_, index) =>
makeOrganization({ organization_id: `org-${index}`, organization_alias: `Org ${index}` }),
);
render(<OrganizationsTable {...baseProps} organizations={organizations} />);
const onUrlUpdate = vi.fn<OnUrlUpdateFunction>();
renderWithProviders(<OrganizationsTable {...baseProps} organizations={thirtyOrganizations} />, { onUrlUpdate });
expect(screen.getAllByRole("row")).toHaveLength(26);
expect(screen.getByTestId("pagination-range")).toHaveTextContent("Showing 1-25 of 30");
@ -207,13 +235,87 @@ describe("OrganizationsTable", () => {
expect(screen.getAllByRole("row")).toHaveLength(31);
expect(screen.getByTestId("pagination-range")).toHaveTextContent("Showing 1-30 of 30");
await waitFor(() => expect(onUrlUpdate.mock.calls.at(-1)?.[0].searchParams.get("page_size")).toBe("50"));
});
it("uses a search-aware empty state", () => {
const { rerender } = render(<OrganizationsTable {...baseProps} searchActive={false} organizations={[]} />);
const { rerender } = renderWithProviders(
<OrganizationsTable {...baseProps} searchActive={false} organizations={[]} />,
);
expect(screen.getByText("No organizations yet")).toBeInTheDocument();
rerender(<OrganizationsTable {...baseProps} searchActive={true} organizations={[]} />);
expect(screen.getByText("No matching organizations")).toBeInTheDocument();
});
});
describe("OrganizationsTable URL state", () => {
it("restores the sort column and direction from ?sort_by=&sort_order=", () => {
renderWithProviders(<OrganizationsTable {...baseProps} organizations={sortableOrganizations} />, {
searchParams: "?sort_by=spend&sort_order=desc",
});
expect(bodyRowAliases()).toEqual(["Mid", "Ace", "Zed"]);
});
it("falls back to sorting by creation date for a ?sort_by= column that cannot be sorted", () => {
renderWithProviders(<OrganizationsTable {...baseProps} organizations={sortableOrganizations} />, {
searchParams: "?sort_by=members&sort_order=asc",
});
expect(bodyRowAliases()).toEqual(["Zed", "Mid", "Ace"]);
});
it("writes the clicked sort column to the URL and returns to the first page", async () => {
const user = userEvent.setup();
const onUrlUpdate = vi.fn<OnUrlUpdateFunction>();
renderWithProviders(<OrganizationsTable {...baseProps} organizations={thirtyOrganizations} />, {
searchParams: "?page=2",
onUrlUpdate,
});
expect(screen.getByTestId("pagination-range")).toHaveTextContent("Showing 26-30 of 30");
await user.click(screen.getByTestId("sort-header-organization_alias"));
await waitFor(() => expect(lastSearchParams(onUrlUpdate)?.get("sort_by")).toBe("organization_alias"));
expect(onUrlUpdate).toHaveBeenCalledTimes(1);
expect(lastSearchParams(onUrlUpdate)?.get("sort_order")).toBe("asc");
expect(lastSearchParams(onUrlUpdate)?.has("page")).toBe(false);
expect(screen.getByTestId("pagination-range")).toHaveTextContent("Showing 1-25 of 30");
expect(within(screen.getAllByRole("row")[1]).getByText("Org 0")).toBeInTheDocument();
});
it("opens the page named by ?page= and writes page changes back to the URL", async () => {
const user = userEvent.setup();
const onUrlUpdate = vi.fn<OnUrlUpdateFunction>();
renderWithProviders(<OrganizationsTable {...baseProps} organizations={thirtyOrganizations} />, {
searchParams: "?page=2",
onUrlUpdate,
});
expect(screen.getByTestId("pagination-range")).toHaveTextContent("Showing 26-30 of 30");
expect(screen.getByText("org-29")).toBeInTheDocument();
await user.click(screen.getByTestId("pagination-prev"));
await waitFor(() => expect(screen.getByTestId("pagination-range")).toHaveTextContent("Showing 1-25 of 30"));
expect(lastSearchParams(onUrlUpdate)?.has("page")).toBe(false);
await user.click(screen.getByTestId("pagination-next"));
await waitFor(() => expect(lastSearchParams(onUrlUpdate)?.get("page")).toBe("2"));
});
it("keeps a deep-linked ?page= while the organization list is still loading", async () => {
const onUrlUpdate = vi.fn<OnUrlUpdateFunction>();
const { rerender } = renderWithProviders(<OrganizationsTable {...baseProps} isLoading organizations={[]} />, {
searchParams: "?page=2",
onUrlUpdate,
});
rerender(<OrganizationsTable {...baseProps} organizations={thirtyOrganizations} />);
await waitFor(() => expect(screen.getByTestId("pagination-range")).toHaveTextContent("Showing 26-30 of 30"));
expect(onUrlUpdate).not.toHaveBeenCalled();
});
});

View file

@ -1,13 +1,13 @@
"use client";
import { SortingState } from "@tanstack/react-table";
import { Building2, SearchX } from "lucide-react";
import React, { useMemo, useState } from "react";
import React, { useMemo } from "react";
import { DataTable } from "@/components/shared/DataTable";
import { Organization } from "@/components/networking";
import { getOrganizationsTableColumns } from "./OrganizationsTableColumns";
import { useOrganizationsTableState } from "./useOrganizationsTableState";
interface OrganizationsTableProps {
organizations: Organization[];
@ -19,8 +19,6 @@ interface OrganizationsTableProps {
onDeleteClick: (organizationId: string) => void;
}
const DEFAULT_SORTING: SortingState = [{ id: "created_at", desc: true }];
function EmptyState({ searchActive }: { searchActive: boolean }) {
const Icon = searchActive ? SearchX : Building2;
return (
@ -49,7 +47,7 @@ const OrganizationsTable: React.FC<OrganizationsTableProps> = ({
onEditClick,
onDeleteClick,
}) => {
const [sorting, setSorting] = useState<SortingState>(DEFAULT_SORTING);
const { sorting, onSortingChange, pagination, onPaginationChange } = useOrganizationsTableState();
const columns = useMemo(() => {
const deps = { userRole, onOrganizationClick, onEditClick, onDeleteClick };
@ -60,11 +58,13 @@ const OrganizationsTable: React.FC<OrganizationsTableProps> = ({
<DataTable
data={organizations}
paginationMode="client"
pagination={pagination}
onPaginationChange={onPaginationChange}
columns={columns}
getRowId={(organization, index) => organization.organization_id || String(index)}
sortingMode="client"
sorting={sorting}
onSortingChange={setSorting}
onSortingChange={onSortingChange}
isLoading={isLoading}
loadingMessage="Loading organizations…"
noDataMessage={<EmptyState searchActive={searchActive} />}

View file

@ -0,0 +1,19 @@
import { useUrlTableState, type UrlTableState, type UrlTableStateOptions } from "@/components/shared/DataTable";
const FILTER_COLUMNS = ["org_id"] as const;
type FilterColumn = (typeof FILTER_COLUMNS)[number];
const TABLE_STATE_OPTIONS: UrlTableStateOptions<FilterColumn> = {
sortFields: ["organization_id", "organization_alias", "created_at", "spend"],
defaultSort: { id: "created_at", desc: true },
defaultPageSize: 25,
filterColumns: FILTER_COLUMNS,
urlKeys: { search: "org_search" },
};
export const useOrganizationsTableState = (): UrlTableState => useUrlTableState(TABLE_STATE_OPTIONS);
export const organizationIdFilter = ({ columnFilters }: Pick<UrlTableState, "columnFilters">): string => {
const value = columnFilters.find((filter) => filter.id === "org_id")?.value;
return typeof value === "string" ? value : "";
};

View file

@ -1,5 +1,7 @@
import { describe, it, expect, vi } from "vitest";
import { renderWithProviders, screen } from "../../../../../tests/test-utils";
import { describe, it, expect, vi, beforeEach, type Mock } from "vitest";
import userEvent from "@testing-library/user-event";
import type { OnUrlUpdateFunction } from "nuqs/adapters/testing";
import { fireEvent, renderWithProviders, screen, waitFor } from "../../../../../tests/test-utils";
import { ProjectKeysSection } from "./ProjectKeysSection";
const mockUseKeys = vi.fn();
@ -70,3 +72,136 @@ describe("ProjectKeysSection", () => {
);
});
});
describe("ProjectKeysSection URL state (keys_ prefix)", () => {
const fortyTwoKeys = {
data: { keys: [], total_count: 42, current_page: 1, total_pages: 9 },
isLoading: false,
isError: false,
};
const lastSearchParams = (onUrlUpdate: Mock<OnUrlUpdateFunction>) => onUrlUpdate.mock.calls.at(-1)?.[0].searchParams;
beforeEach(() => {
mockUseKeys.mockReset();
});
it("should fetch the page, page size and key name filter named by the keys_ params", () => {
mockUseKeys.mockReturnValue(fortyTwoKeys);
renderWithProviders(<ProjectKeysSection projectId="proj-1" />, {
searchParams: "?page=4&keys_page=2&keys_page_size=10&keys_search=prod",
});
expect(mockUseKeys).toHaveBeenLastCalledWith(
2,
10,
expect.objectContaining({ projectID: "proj-1", selectedKeyAlias: "prod" }),
);
expect(screen.getByPlaceholderText("Filter by key name...")).toHaveValue("prod");
expect(screen.getByTestId("pagination-page")).toHaveTextContent("Page 2 of 5");
});
it("should cap an oversized ?keys_page_size= at the largest offered page size", () => {
mockUseKeys.mockReturnValue(fortyTwoKeys);
renderWithProviders(<ProjectKeysSection projectId="proj-1" />, { searchParams: "?keys_page_size=500" });
expect(mockUseKeys).toHaveBeenLastCalledWith(1, 25, expect.anything());
expect(screen.getByTestId("pagination-page")).toHaveTextContent("Page 1 of 2");
});
it("should fall back to the default page size for a ?keys_page_size= outside the offered options", () => {
mockUseKeys.mockReturnValue(fortyTwoKeys);
renderWithProviders(<ProjectKeysSection projectId="proj-1" />, { searchParams: "?keys_page_size=7" });
expect(mockUseKeys).toHaveBeenLastCalledWith(1, 5, expect.anything());
expect(screen.getByTestId("pagination-page")).toHaveTextContent("Page 1 of 9");
});
it("should drop an unsupported ?keys_page_size= when the user pages forward", async () => {
const user = userEvent.setup();
mockUseKeys.mockReturnValue(fortyTwoKeys);
const onUrlUpdate = vi.fn<OnUrlUpdateFunction>();
renderWithProviders(<ProjectKeysSection projectId="proj-1" />, { searchParams: "?keys_page_size=7", onUrlUpdate });
await user.click(screen.getByTestId("pagination-next"));
await waitFor(() => expect(onUrlUpdate.mock.calls.at(-1)?.[0].queryString).toBe("?keys_page=2"));
expect(mockUseKeys).toHaveBeenLastCalledWith(2, 5, expect.anything());
});
it("should write the key name filter to ?keys_search= and return the keys to their first page", async () => {
mockUseKeys.mockReturnValue(fortyTwoKeys);
const onUrlUpdate = vi.fn<OnUrlUpdateFunction>();
renderWithProviders(<ProjectKeysSection projectId="proj-1" />, {
searchParams: "?page=4&keys_page=3",
onUrlUpdate,
});
fireEvent.change(screen.getByPlaceholderText("Filter by key name..."), { target: { value: "prod" } });
await waitFor(() => expect(lastSearchParams(onUrlUpdate)?.get("keys_search")).toBe("prod"));
expect(lastSearchParams(onUrlUpdate)?.has("keys_page")).toBe(false);
expect(lastSearchParams(onUrlUpdate)?.get("page")).toBe("4");
expect(mockUseKeys).toHaveBeenLastCalledWith(1, 5, expect.objectContaining({ selectedKeyAlias: "prod" }));
});
it("should remove ?keys_search= when the key filter is cleared", async () => {
const user = userEvent.setup();
mockUseKeys.mockReturnValue(fortyTwoKeys);
const onUrlUpdate = vi.fn<OnUrlUpdateFunction>();
renderWithProviders(<ProjectKeysSection projectId="proj-1" />, { searchParams: "?keys_search=prod", onUrlUpdate });
await user.click(screen.getByRole("button", { name: /clear key filter/i }));
await waitFor(() => expect(lastSearchParams(onUrlUpdate)?.has("keys_search")).toBe(false));
expect(mockUseKeys).toHaveBeenLastCalledWith(1, 5, expect.objectContaining({ selectedKeyAlias: null }));
});
it("should write key pages to ?keys_page= without touching the projects list's ?page=", async () => {
const user = userEvent.setup();
mockUseKeys.mockReturnValue(fortyTwoKeys);
const onUrlUpdate = vi.fn<OnUrlUpdateFunction>();
renderWithProviders(<ProjectKeysSection projectId="proj-1" />, { searchParams: "?page=4", onUrlUpdate });
await user.click(screen.getByTestId("pagination-next"));
await waitFor(() => expect(lastSearchParams(onUrlUpdate)?.get("keys_page")).toBe("2"));
expect(lastSearchParams(onUrlUpdate)?.get("page")).toBe("4");
expect(mockUseKeys).toHaveBeenLastCalledWith(2, 5, expect.anything());
});
it("should snap a ?keys_page= past the last page back to the last page once the keys load", async () => {
mockUseKeys.mockReturnValue({ data: undefined, isLoading: true, isError: false });
const onUrlUpdate = vi.fn<OnUrlUpdateFunction>();
const { rerender } = renderWithProviders(<ProjectKeysSection projectId="proj-1" />, {
searchParams: "?keys_page=9",
onUrlUpdate,
});
expect(mockUseKeys).toHaveBeenLastCalledWith(9, 5, expect.anything());
mockUseKeys.mockReturnValue({
data: { keys: [], total_count: 6, current_page: 9, total_pages: 2 },
isLoading: false,
isError: false,
});
rerender(<ProjectKeysSection projectId="proj-1" />);
await waitFor(() => expect(lastSearchParams(onUrlUpdate)?.get("keys_page")).toBe("2"));
expect(mockUseKeys).toHaveBeenLastCalledWith(2, 5, expect.anything());
});
it("should keep a deep-linked ?keys_page= when the key fetch fails", async () => {
mockUseKeys.mockReturnValue({ data: undefined, isLoading: true, isError: false });
const onUrlUpdate = vi.fn<OnUrlUpdateFunction>();
const { rerender } = renderWithProviders(<ProjectKeysSection projectId="proj-1" />, {
searchParams: "?keys_page=3",
onUrlUpdate,
});
mockUseKeys.mockReturnValue({ data: undefined, isLoading: false, isError: true });
rerender(<ProjectKeysSection projectId="proj-1" />);
await new Promise((resolve) => setTimeout(resolve, 100));
expect(onUrlUpdate).not.toHaveBeenCalled();
expect(mockUseKeys).toHaveBeenLastCalledWith(3, 5, expect.anything());
});
});

View file

@ -1,30 +1,27 @@
import { useKeys } from "@/app/(dashboard)/hooks/keys/useKeys";
import { PaginationState } from "@tanstack/react-table";
import { KeyIcon, SearchIcon, X } from "lucide-react";
import { useEffect, useState } from "react";
import { Card, CardContent, CardHeader, CardTitle } from "@/components/ui/card";
import { InputGroup, InputGroupAddon, InputGroupButton, InputGroupInput } from "@/components/ui/input-group";
import { ProjectKeysTable } from "./ProjectKeysTable";
import { useProjectKeysTableState } from "./useProjectsUrlState";
interface ProjectKeysSectionProps {
projectId: string;
}
const PAGE_SIZE = 5;
export function ProjectKeysSection({ projectId }: ProjectKeysSectionProps) {
const [pagination, setPagination] = useState<PaginationState>({ pageIndex: 0, pageSize: PAGE_SIZE });
const [keyAlias, setKeyAlias] = useState<string>("");
const {
search: keyAlias,
setSearch: setKeyAlias,
pagination,
onPaginationChange: setPagination,
} = useProjectKeysTableState();
const { data, isLoading } = useKeys(pagination.pageIndex + 1, pagination.pageSize, {
const { data, isLoading, isError } = useKeys(pagination.pageIndex + 1, pagination.pageSize, {
projectID: projectId,
selectedKeyAlias: keyAlias || null,
});
useEffect(() => {
setPagination((current) => ({ ...current, pageIndex: 0 }));
}, [keyAlias]);
const keys = data?.keys ?? [];
const totalCount = data?.total_count ?? 0;
@ -60,6 +57,7 @@ export function ProjectKeysSection({ projectId }: ProjectKeysSectionProps) {
keys={keys}
totalCount={totalCount}
isLoading={isLoading}
isError={isError}
pagination={pagination}
onPaginationChange={setPagination}
/>

View file

@ -8,17 +8,17 @@ import { KeyResponse } from "@/components/key_team_helpers/key_list";
import { DataTable } from "@/components/shared/DataTable";
import { getProjectKeysTableColumns } from "./ProjectKeysTableColumns";
import { PROJECT_KEYS_PAGE_SIZE_OPTIONS } from "./useProjectsUrlState";
interface ProjectKeysTableProps {
keys: KeyResponse[];
totalCount: number;
isLoading: boolean;
isError?: boolean;
pagination: PaginationState;
onPaginationChange: OnChangeFn<PaginationState>;
}
const PAGE_SIZE_OPTIONS = [5, 10, 25];
function EmptyState() {
return (
<div className="flex flex-col items-center gap-1 py-6">
@ -35,6 +35,7 @@ export function ProjectKeysTable({
keys,
totalCount,
isLoading,
isError = false,
pagination,
onPaginationChange,
}: ProjectKeysTableProps) {
@ -49,8 +50,9 @@ export function ProjectKeysTable({
pagination={pagination}
onPaginationChange={onPaginationChange}
rowCount={totalCount}
pageSizeOptions={PAGE_SIZE_OPTIONS}
pageSizeOptions={PROJECT_KEYS_PAGE_SIZE_OPTIONS}
isLoading={isLoading}
isError={isError}
loadingMessage="Loading keys…"
noDataMessage={<EmptyState />}
size="compact"

View file

@ -190,22 +190,48 @@ describe("ProjectsPage", () => {
it("should reset to the first page when the search text changes", async () => {
const user = userEvent.setup();
const onUrlUpdate = vi.fn<(event: UrlUpdateEvent) => void>();
const manyProjects = Array.from({ length: 12 }, (_, i) => ({
...mockProjects[0],
project_id: `proj-${i + 1}`,
project_alias: `Project ${String(i + 1).padStart(2, "0")}`,
}));
mockUseProjects.mockReturnValue({ data: manyProjects, isLoading: false });
renderWithProviders(<ProjectsPage />);
renderWithProviders(<ProjectsPage />, { onUrlUpdate });
await user.click(screen.getByTestId("pagination-next"));
expect(screen.getByTestId("pagination-page")).toHaveTextContent("Page 2 of 2");
await waitFor(() => expect(onUrlUpdate.mock.calls.at(-1)?.[0].searchParams.get("page")).toBe("2"));
fireEvent.change(screen.getByPlaceholderText(/search projects/i), { target: { value: "Project 01" } });
await waitFor(() => {
expect(screen.getByText("Project 01")).toBeInTheDocument();
expect(screen.getByTestId("pagination-page")).toHaveTextContent("Page 1 of 1");
});
await waitFor(() => expect(onUrlUpdate.mock.calls.at(-1)?.[0].queryString).toBe("?project_search=Project+01"));
expect(onUrlUpdate).toHaveBeenCalledTimes(2);
});
it("should restore the search box and filtered list from a ?project_search= deep link", () => {
mockUseProjects.mockReturnValue({ data: mockProjects, isLoading: false });
renderWithProviders(<ProjectsPage />, { searchParams: "?project_search=Beta" });
expect(screen.getByPlaceholderText(/search projects/i)).toHaveValue("Beta");
expect(screen.getByText("Beta Project")).toBeInTheDocument();
expect(screen.queryByText("Alpha Project")).not.toBeInTheDocument();
});
it("should remove ?project_search= when the search is cleared", async () => {
const user = userEvent.setup();
const onUrlUpdate = vi.fn<(event: UrlUpdateEvent) => void>();
mockUseProjects.mockReturnValue({ data: mockProjects, isLoading: false });
renderWithProviders(<ProjectsPage />, { searchParams: "?project_search=Beta", onUrlUpdate });
await user.click(screen.getByRole("button", { name: /clear search/i }));
await waitFor(() => expect(onUrlUpdate).toHaveBeenLastCalledWith(expect.objectContaining({ queryString: "" })));
expect(screen.getByPlaceholderText(/search projects/i)).toHaveValue("");
expect(screen.getByText("Alpha Project")).toBeInTheDocument();
});
it("should open the detail view directly from a ?project= deep link", () => {
@ -250,6 +276,24 @@ describe("ProjectsPage", () => {
expect(screen.getByText("Alpha Project")).toBeInTheDocument();
});
it("should drop the project's key table state but keep the list's search and page when the detail view is closed", async () => {
const user = userEvent.setup();
const onUrlUpdate = vi.fn<(event: UrlUpdateEvent) => void>();
mockUseProjects.mockReturnValue({ data: mockProjects, isLoading: false });
renderWithProviders(<ProjectsPage />, {
searchParams:
"?page=2&project_search=Project&project=proj-1&keys_page=3&keys_page_size=10&keys_search=prod&keys_sort_by=spend&keys_sort_order=asc",
onUrlUpdate,
});
await user.click(screen.getByRole("button", { name: /back to projects/i }));
await waitFor(() => expect(onUrlUpdate).toHaveBeenCalledTimes(1));
const [update] = onUrlUpdate.mock.calls[0];
expect(update.queryString).toBe("?page=2&project_search=Project");
expect(update.options.history).toBe("replace");
});
it("should resolve team alias from the teams list in the Team column", () => {
mockUseTeams.mockReturnValue({
data: [{ team_id: "team-1", team_alias: "Engineering", models: [] }],

View file

@ -9,6 +9,7 @@ import { InputGroup, InputGroupAddon, InputGroupButton, InputGroupInput } from "
import { CreateProjectModal } from "./ProjectModals/CreateProjectModal";
import { ProjectDetail } from "./ProjectDetailsPage";
import { ProjectsTable } from "./ProjectsTable";
import { useClearProjectKeysTableState, useProjectsTableState } from "./useProjectsUrlState";
export function ProjectsPage() {
const { data: projects, isLoading } = useProjects();
@ -18,8 +19,9 @@ export function ProjectsPage() {
"project",
parseAsString.withOptions({ history: "push" }),
);
const clearProjectKeysTableState = useClearProjectKeysTableState();
const { search: searchText, setSearch: setSearchText } = useProjectsTableState();
const [isCreateModalVisible, setIsCreateModalVisible] = useState(false);
const [searchText, setSearchText] = useState("");
const teamAliasMap = useMemo(() => {
const map = new Map<string, string>();
@ -44,13 +46,13 @@ export function ProjectsPage() {
});
}, [projects, searchText, teamAliasMap]);
const closeProject = () => {
void setSelectedProjectId(null, { history: "replace" });
clearProjectKeysTableState();
};
if (selectedProjectId) {
return (
<ProjectDetail
projectId={selectedProjectId}
onBack={() => void setSelectedProjectId(null, { history: "replace" })}
/>
);
return <ProjectDetail projectId={selectedProjectId} onBack={closeProject} />;
}
return (

View file

@ -83,6 +83,7 @@ describe("ProjectsTable pagination URL state", () => {
await waitFor(() => expect(onUrlUpdate).toHaveBeenCalled());
const [update] = onUrlUpdate.mock.calls[0];
expect(update.searchParams.get("page")).toBe("2");
expect(update.searchParams.has("page_size")).toBe(false);
expect(update.options.history).toBe("push");
expect(firstDataRow().getByText("Project 11")).toBeInTheDocument();
});
@ -147,6 +148,7 @@ describe("ProjectsTable pagination URL state", () => {
const lastUpdate = onUrlUpdate.mock.calls.at(-1)?.[0];
expect(lastUpdate.searchParams.get("page")).toBeNull();
expect(lastUpdate.searchParams.get("page_size")).toBe("25");
expect(lastUpdate.options.history).toBe("push");
});
it("should apply both params from a ?page=2&page_size=25 deep link so the restored view matches", () => {

View file

@ -2,13 +2,13 @@
import { SortingState } from "@tanstack/react-table";
import { FolderKanban } from "lucide-react";
import { parseAsInteger, useQueryStates } from "nuqs";
import { useMemo, useState } from "react";
import { ProjectResponse } from "@/app/(dashboard)/hooks/projects/useProjects";
import { DataTable, DataTablePagination } from "@/components/shared/DataTable";
import { getProjectsTableColumns } from "./ProjectsTableColumns";
import { PROJECTS_DEFAULT_PAGE_SIZE, useProjectsTableState } from "./useProjectsUrlState";
interface ProjectsTableProps {
projects: ProjectResponse[];
@ -19,8 +19,7 @@ interface ProjectsTableProps {
isTeamsLoading: boolean;
}
const DEFAULT_PAGE_SIZE = 10;
const PAGE_SIZE_OPTIONS = [DEFAULT_PAGE_SIZE, 25, 50];
const PAGE_SIZE_OPTIONS = [PROJECTS_DEFAULT_PAGE_SIZE, 25, 50];
function EmptyState({ isFiltered }: { isFiltered: boolean }) {
return (
@ -47,11 +46,8 @@ export function ProjectsTable({
isTeamsLoading,
}: ProjectsTableProps) {
const [sorting, setSorting] = useState<SortingState>([]);
const [{ page, page_size }, setPagination] = useQueryStates(
{ page: parseAsInteger.withDefault(1), page_size: parseAsInteger.withDefault(DEFAULT_PAGE_SIZE) },
{ history: "push" },
);
const pageSize = PAGE_SIZE_OPTIONS.includes(page_size) ? page_size : DEFAULT_PAGE_SIZE;
const { pagination, onPaginationChange } = useProjectsTableState();
const pageSize = PAGE_SIZE_OPTIONS.includes(pagination.pageSize) ? pagination.pageSize : PROJECTS_DEFAULT_PAGE_SIZE;
const columns = useMemo(() => {
const deps = { onProjectClick, teamAliasMap, isTeamsLoading };
@ -59,7 +55,7 @@ export function ProjectsTable({
}, [onProjectClick, teamAliasMap, isTeamsLoading]);
const pageCount = Math.max(Math.ceil(projects.length / pageSize), 1);
const pageIndex = page >= 1 && page <= pageCount ? page - 1 : 0;
const pageIndex = pagination.pageIndex < pageCount ? pagination.pageIndex : 0;
return (
<DataTable
@ -77,8 +73,8 @@ export function ProjectsTable({
page={pageIndex}
pageSize={pageSize}
rowCount={projects.length}
onPageChange={(nextPageIndex) => void setPagination({ page: nextPageIndex + 1 })}
onPageSizeChange={(nextPageSize) => void setPagination({ page_size: nextPageSize, page: null })}
onPageChange={(nextPageIndex) => onPaginationChange({ pageIndex: nextPageIndex, pageSize })}
onPageSizeChange={(nextPageSize) => onPaginationChange({ pageIndex: 0, pageSize: nextPageSize })}
pageSizeOptions={PAGE_SIZE_OPTIONS}
isLoading={isLoading}
/>

View file

@ -0,0 +1,79 @@
import { functionalUpdate, type OnChangeFn, type PaginationState } from "@tanstack/react-table";
import { useUrlTableState, type UrlTableState, type UrlTableStateOptions } from "@/components/shared/DataTable";
import { parseAsInteger, useQueryStates } from "nuqs";
import { useCallback, useMemo } from "react";
export const PROJECTS_DEFAULT_PAGE_SIZE = 10;
export const PROJECT_KEYS_DEFAULT_PAGE_SIZE = 5;
export const PROJECT_KEYS_PAGE_SIZE_OPTIONS = [PROJECT_KEYS_DEFAULT_PAGE_SIZE, 10, 25];
const PROJECTS_TABLE_STATE_OPTIONS: UrlTableStateOptions<never> = {
sortFields: [],
defaultSort: { id: "created_at", desc: true },
defaultPageSize: PROJECTS_DEFAULT_PAGE_SIZE,
filterColumns: [],
urlKeys: { search: "project_search" },
};
const PROJECTS_PAGE_PARAMS = {
page: parseAsInteger.withDefault(1),
page_size: parseAsInteger.withDefault(PROJECTS_DEFAULT_PAGE_SIZE),
};
const PROJECT_KEYS_TABLE_STATE_OPTIONS: UrlTableStateOptions<never> = {
sortFields: [],
defaultSort: { id: "created_at", desc: true },
defaultPageSize: PROJECT_KEYS_DEFAULT_PAGE_SIZE,
maxPageSize: Math.max(...PROJECT_KEYS_PAGE_SIZE_OPTIONS),
filterColumns: [],
keyPrefix: "keys_",
};
export function useProjectsTableState(): UrlTableState {
const tableState = useUrlTableState(PROJECTS_TABLE_STATE_OPTIONS);
const [, setPageParams] = useQueryStates(PROJECTS_PAGE_PARAMS, { history: "push" });
const { pagination } = tableState;
const onPaginationChange = useCallback<OnChangeFn<PaginationState>>(
(updaterOrValue) => {
const next = functionalUpdate(updaterOrValue, pagination);
void setPageParams({ page: next.pageIndex + 1, page_size: next.pageSize });
},
[pagination, setPageParams],
);
return useMemo(() => ({ ...tableState, onPaginationChange }), [tableState, onPaginationChange]);
}
export function useProjectKeysTableState(): UrlTableState {
const tableState = useUrlTableState(PROJECT_KEYS_TABLE_STATE_OPTIONS);
const { pagination: urlPagination, onPaginationChange: writePagination } = tableState;
const pageSize = PROJECT_KEYS_PAGE_SIZE_OPTIONS.includes(urlPagination.pageSize)
? urlPagination.pageSize
: PROJECT_KEYS_DEFAULT_PAGE_SIZE;
const pagination = useMemo<PaginationState>(
() => ({ pageIndex: urlPagination.pageIndex, pageSize }),
[urlPagination.pageIndex, pageSize],
);
const onPaginationChange = useCallback<OnChangeFn<PaginationState>>(
(updaterOrValue) => writePagination(functionalUpdate(updaterOrValue, pagination)),
[pagination, writePagination],
);
return useMemo(
() => ({ ...tableState, pagination, onPaginationChange }),
[tableState, pagination, onPaginationChange],
);
}
export function useClearProjectKeysTableState(): () => void {
const { setSearch, onSortingChange, onColumnFiltersChange, onPaginationChange } = useProjectKeysTableState();
return useCallback(() => {
setSearch("");
onSortingChange([]);
onColumnFiltersChange([]);
onPaginationChange({ pageIndex: 0, pageSize: PROJECT_KEYS_DEFAULT_PAGE_SIZE });
}, [setSearch, onSortingChange, onColumnFiltersChange, onPaginationChange]);
}

View file

@ -0,0 +1,3 @@
export const ORGANIZATION_TABS = ["overview", "members", "settings"] as const;
export type OrganizationTab = (typeof ORGANIZATION_TABS)[number];
export const ORGANIZATION_TAB_URL_KEY = "org_tab";

View file

@ -1,8 +1,10 @@
import React from "react";
import { fireEvent, screen, waitFor, within } from "@testing-library/react";
import { fireEvent, render, screen, waitFor, within } from "@testing-library/react";
import userEvent from "@testing-library/user-event";
import { vi, test, expect, beforeEach } from "vitest";
import { renderWithProviders } from "../../../tests/test-utils";
import { QueryClientProvider } from "@tanstack/react-query";
import { NuqsTestingAdapter, type OnUrlUpdateFunction } from "nuqs/adapters/testing";
import { vi, test, expect, beforeEach, describe, type Mock } from "vitest";
import { renderWithProviders, testQueryClient } from "../../../tests/test-utils";
import OrganizationInfoView from "./organization_view";
import { useOrganization } from "@/app/(dashboard)/hooks/organizations/useOrganizations";
@ -115,7 +117,6 @@ test("renders organization view after loading data", async () => {
is_org_admin={false}
is_proxy_admin={false}
userModels={[]}
editOrg={false}
/>,
);
@ -135,7 +136,6 @@ test("should display empty state when organization has no members", async () =>
is_org_admin={false}
is_proxy_admin={false}
userModels={[]}
editOrg={false}
/>,
);
@ -165,7 +165,6 @@ test("should display team aliases when teams are available", async () => {
is_org_admin={false}
is_proxy_admin={false}
userModels={[]}
editOrg={false}
/>,
);
@ -199,7 +198,6 @@ test("should display team ID as fallback when alias is not found", async () => {
is_org_admin={false}
is_proxy_admin={false}
userModels={[]}
editOrg={false}
/>,
);
@ -223,7 +221,6 @@ test("links each team badge to that team's detail page", async () => {
is_org_admin={false}
is_proxy_admin={false}
userModels={[]}
editOrg={false}
/>,
);
@ -250,7 +247,6 @@ test("model badges stay non-clickable", async () => {
is_org_admin={false}
is_proxy_admin={false}
userModels={[]}
editOrg={false}
/>,
);
@ -272,7 +268,6 @@ test("should keep unsaved settings edits when switching tabs and back", async ()
is_org_admin={false}
is_proxy_admin={true}
userModels={[]}
editOrg={false}
/>,
);
@ -308,7 +303,6 @@ test("renders a tpm/rpm limit of 0 as 0 in the overview and settings tabs, never
is_org_admin={false}
is_proxy_admin={true}
userModels={[]}
editOrg={false}
/>,
);
@ -323,3 +317,99 @@ test("renders a tpm/rpm limit of 0 as 0 in the overview and settings tabs, never
expect(screen.queryByText("TPM: Unlimited")).not.toBeInTheDocument();
expect(screen.queryByText("RPM: Unlimited")).not.toBeInTheDocument();
});
const renderOrgView = (props: { is_proxy_admin?: boolean } = {}) => (
<OrganizationInfoView
organizationId="org_123"
onClose={() => {}}
accessToken="test-token"
is_org_admin={false}
is_proxy_admin={props.is_proxy_admin ?? false}
userModels={[]}
/>
);
const lastSearchParams = (onUrlUpdate: Mock<OnUrlUpdateFunction>) => onUrlUpdate.mock.calls.at(-1)?.[0].searchParams;
describe("organization detail tab in the URL (?org_tab=)", () => {
beforeEach(() => {
mockUseOrganization.mockReturnValue({ data: mockOrg, isLoading: false } as unknown as ReturnType<
typeof useOrganization
>);
});
test("opens on the tab named by ?org_tab=", () => {
renderWithProviders(renderOrgView(), { searchParams: "?org=org_123&org_tab=members" });
expect(screen.getByRole("tab", { name: "Members" })).toHaveAttribute("aria-selected", "true");
expect(screen.getByText("No members found")).toBeInTheDocument();
});
test("the settings deep link used by the list's Edit action opens the Settings tab", () => {
renderWithProviders(renderOrgView({ is_proxy_admin: true }), { searchParams: "?org=org_123&org_tab=settings" });
expect(screen.getByRole("tab", { name: "Settings" })).toHaveAttribute("aria-selected", "true");
expect(screen.getByRole("button", { name: /Edit Settings/i })).toBeInTheDocument();
});
test("opens on Overview when the URL names no tab", () => {
renderWithProviders(renderOrgView(), { searchParams: "?org=org_123" });
expect(screen.getByRole("tab", { name: "Overview" })).toHaveAttribute("aria-selected", "true");
});
test("writes the selected tab to ?org_tab= and drops it again for Overview", async () => {
const user = userEvent.setup();
const onUrlUpdate = vi.fn<OnUrlUpdateFunction>();
renderWithProviders(renderOrgView(), { searchParams: "?org=org_123", onUrlUpdate });
await user.click(screen.getByRole("tab", { name: "Settings" }));
await waitFor(() => expect(lastSearchParams(onUrlUpdate)?.get("org_tab")).toBe("settings"));
expect(lastSearchParams(onUrlUpdate)?.get("org")).toBe("org_123");
expect(screen.getByRole("tab", { name: "Settings" })).toHaveAttribute("aria-selected", "true");
await user.click(screen.getByRole("tab", { name: "Overview" }));
await waitFor(() => expect(lastSearchParams(onUrlUpdate)?.has("org_tab")).toBe(false));
expect(lastSearchParams(onUrlUpdate)?.get("org")).toBe("org_123");
});
test("falls back to Overview for an unknown ?org_tab= and removes it from the URL", async () => {
const onUrlUpdate = vi.fn<OnUrlUpdateFunction>();
render(renderOrgView(), {
wrapper: ({ children }) => (
<NuqsTestingAdapter
searchParams="?org=org_123&org_tab=billing"
onUrlUpdate={onUrlUpdate}
hasMemory
resetUrlUpdateQueueOnMount={false}
>
<QueryClientProvider client={testQueryClient}>{children}</QueryClientProvider>
</NuqsTestingAdapter>
),
});
expect(screen.getByRole("tab", { name: "Overview" })).toHaveAttribute("aria-selected", "true");
await waitFor(() => expect(onUrlUpdate).toHaveBeenCalled());
expect(lastSearchParams(onUrlUpdate)?.has("org_tab")).toBe(false);
expect(lastSearchParams(onUrlUpdate)?.get("org")).toBe("org_123");
});
test("follows back and forward navigation between tabs while the detail view stays open", () => {
const atUrl = (searchParams: string) => (
<NuqsTestingAdapter searchParams={searchParams} hasMemory>
<QueryClientProvider client={testQueryClient}>{renderOrgView()}</QueryClientProvider>
</NuqsTestingAdapter>
);
const { rerender } = render(atUrl("?org=org_123&org_tab=members"));
expect(screen.getByRole("tab", { name: "Members" })).toHaveAttribute("aria-selected", "true");
rerender(atUrl("?org=org_123"));
expect(screen.getByRole("tab", { name: "Overview" })).toHaveAttribute("aria-selected", "true");
rerender(atUrl("?org=org_123&org_tab=members"));
expect(screen.getByRole("tab", { name: "Members" })).toHaveAttribute("aria-selected", "true");
expect(screen.getByText("No members found")).toBeInTheDocument();
});
});

View file

@ -1,6 +1,7 @@
import { useTeams } from "@/app/(dashboard)/hooks/teams/useTeams";
import { organizationKeys, useOrganization } from "@/app/(dashboard)/hooks/organizations/useOrganizations";
import { useQueryClient } from "@tanstack/react-query";
import { useUrlTab } from "@/hooks/useUrlTab";
import { useVisitedTabs } from "@/hooks/useVisitedTabs";
import { MoneyCell } from "@/components/shared/table_cells";
import CopyButton from "@/components/shared/CopyButton";
@ -25,6 +26,7 @@ import {
import ObjectPermissionsView from "../object_permissions_view";
import MemberModal from "../team/EditMembership";
import { OrgSettingsForm } from "./org-settings/OrgSettingsForm";
import { ORGANIZATION_TAB_URL_KEY, ORGANIZATION_TABS, type OrganizationTab } from "./organizationTabs";
interface OrganizationInfoProps {
organizationId: string;
@ -33,7 +35,6 @@ interface OrganizationInfoProps {
is_org_admin: boolean;
is_proxy_admin: boolean;
userModels: string[];
editOrg: boolean;
}
const OrganizationInfoView: React.FC<OrganizationInfoProps> = ({
@ -43,7 +44,6 @@ const OrganizationInfoView: React.FC<OrganizationInfoProps> = ({
is_org_admin,
is_proxy_admin,
userModels,
editOrg,
}) => {
const queryClient = useQueryClient();
const { data: orgData, isLoading: loading } = useOrganization(organizationId);
@ -53,10 +53,16 @@ const OrganizationInfoView: React.FC<OrganizationInfoProps> = ({
const [selectedEditMember, setSelectedEditMember] = useState<Member | null>(null);
const canEditOrg = is_org_admin || is_proxy_admin;
const { data: teams } = useTeams();
const { onTabChange, hasVisited } = useVisitedTabs(editOrg ? "settings" : "overview");
const [tab, setTab] = useUrlTab(ORGANIZATION_TABS, "overview", ORGANIZATION_TAB_URL_KEY);
const { onTabChange, hasVisited } = useVisitedTabs(tab);
const teamAliasMap = useMemo(() => createTeamAliasMap(teams), [teams]);
const handleTabChange = (value: OrganizationTab) => {
setTab(value);
onTabChange(value);
};
const handleMemberAdd = async (values: any) => {
try {
if (accessToken == null) {
@ -158,7 +164,7 @@ const OrganizationInfoView: React.FC<OrganizationInfoProps> = ({
</div>
</div>
<Tabs defaultValue={editOrg ? "settings" : "overview"} onValueChange={onTabChange} className="mb-4">
<Tabs value={tab} onValueChange={handleTabChange} className="mb-4">
<TabsList variant="line" className="h-auto w-full justify-start rounded-none border-b p-0">
<TabsTrigger value="overview" className="flex-none rounded-none px-4 py-2">
Overview