mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-12 23:01:41 +00:00
Merge remote-tracking branch 'origin/litellm_internal_staging' into litellm_registry_audit_2026_09_02
This commit is contained in:
mateo
commit
71823afd78
18 changed files with 778 additions and 65 deletions
|
|
@ -14,9 +14,22 @@ then fails on a Node binary that was never written. Deleting a cache directory
|
|||
that exists without a Node binary is what turns a killed bootstrap back into a
|
||||
recoverable one.
|
||||
|
||||
Both budgets are overridable so an operator can widen them without a release:
|
||||
``LITELLM_PRISMA_BOOTSTRAP_TIMEOUT`` for the toolchain install and
|
||||
``LITELLM_PRISMA_COMMAND_TIMEOUT`` for every individual Prisma command.
|
||||
``prisma migrate deploy`` is the other command whose runtime is not a
|
||||
constant: it grows with the number of pending migrations, so a fresh database
|
||||
that has to replay every migration this package ships overruns a per-command
|
||||
budget sized for the short bookkeeping commands, on a laptop as much as on a
|
||||
slow CI runner. The Python ``prisma`` wrapper spawns Node and the schema engine
|
||||
as separate children, so killing the wrapper on timeout leaves them running:
|
||||
the retry then contends with that orphan for Prisma's advisory lock and cannot
|
||||
finish any sooner. Migrate deploy therefore runs under its own budget.
|
||||
|
||||
All three budgets are overridable so an operator can widen them without a
|
||||
release: ``LITELLM_PRISMA_BOOTSTRAP_TIMEOUT`` for the toolchain install,
|
||||
``LITELLM_PRISMA_MIGRATE_DEPLOY_TIMEOUT`` for ``prisma migrate deploy`` and
|
||||
``LITELLM_PRISMA_COMMAND_TIMEOUT`` for every other Prisma command. The
|
||||
per-command budget used to bound migrate deploy as well, so a deployment that
|
||||
raised it above the deploy default keeps that larger budget for deploy unless
|
||||
the deploy override says otherwise.
|
||||
"""
|
||||
|
||||
import math
|
||||
|
|
@ -36,10 +49,12 @@ except ImportError:
|
|||
|
||||
PRISMA_COMMAND_TIMEOUT_ENV_VAR = "LITELLM_PRISMA_COMMAND_TIMEOUT"
|
||||
PRISMA_BOOTSTRAP_TIMEOUT_ENV_VAR = "LITELLM_PRISMA_BOOTSTRAP_TIMEOUT"
|
||||
PRISMA_MIGRATE_DEPLOY_TIMEOUT_ENV_VAR = "LITELLM_PRISMA_MIGRATE_DEPLOY_TIMEOUT"
|
||||
NODEENV_CACHE_DIR_ENV_VAR = "PRISMA_NODEENV_CACHE_DIR"
|
||||
|
||||
DEFAULT_PRISMA_COMMAND_TIMEOUT = 60.0
|
||||
DEFAULT_PRISMA_BOOTSTRAP_TIMEOUT = 600.0
|
||||
DEFAULT_PRISMA_MIGRATE_DEPLOY_TIMEOUT = 600.0
|
||||
|
||||
BOOTSTRAP_ARG = "--version"
|
||||
|
||||
|
|
@ -88,6 +103,15 @@ def prisma_bootstrap_timeout() -> float:
|
|||
)
|
||||
|
||||
|
||||
def prisma_migrate_deploy_timeout() -> float:
|
||||
"""Seconds one ``prisma migrate deploy`` may run for, however many migrations are pending."""
|
||||
if os.getenv(PRISMA_MIGRATE_DEPLOY_TIMEOUT_ENV_VAR) is not None:
|
||||
return _timeout_from_env(
|
||||
PRISMA_MIGRATE_DEPLOY_TIMEOUT_ENV_VAR, DEFAULT_PRISMA_MIGRATE_DEPLOY_TIMEOUT
|
||||
)
|
||||
return max(DEFAULT_PRISMA_MIGRATE_DEPLOY_TIMEOUT, prisma_command_timeout())
|
||||
|
||||
|
||||
def nodeenv_cache_dir() -> Optional[Path]:
|
||||
"""Where Prisma installs its private Node runtime, or None if unknowable."""
|
||||
override = os.getenv(NODEENV_CACHE_DIR_ENV_VAR)
|
||||
|
|
|
|||
|
|
@ -15,8 +15,11 @@ from litellm_proxy_extras.replica_identity import (
|
|||
apply_replica_identity_full,
|
||||
)
|
||||
from litellm_proxy_extras.prisma_toolchain import (
|
||||
PRISMA_COMMAND_TIMEOUT_ENV_VAR,
|
||||
PRISMA_MIGRATE_DEPLOY_TIMEOUT_ENV_VAR,
|
||||
ensure_prisma_toolchain,
|
||||
prisma_command_timeout,
|
||||
prisma_migrate_deploy_timeout,
|
||||
)
|
||||
|
||||
|
||||
|
|
@ -698,12 +701,13 @@ class ProxyExtrasDBManager:
|
|||
|
||||
original_dir = os.getcwd()
|
||||
os.chdir(migrations_dir)
|
||||
deploy_timeout = prisma_migrate_deploy_timeout()
|
||||
try:
|
||||
for attempt in range(4):
|
||||
try:
|
||||
result = subprocess.run(
|
||||
[_get_prisma_command(), "migrate", "deploy"],
|
||||
timeout=prisma_command_timeout(),
|
||||
timeout=deploy_timeout,
|
||||
check=True,
|
||||
capture_output=True,
|
||||
text=True,
|
||||
|
|
@ -713,8 +717,12 @@ class ProxyExtrasDBManager:
|
|||
return True
|
||||
|
||||
except subprocess.TimeoutExpired:
|
||||
logger.info(
|
||||
f"prisma migrate deploy attempt {attempt + 1} timed out, retrying"
|
||||
logger.warning(
|
||||
"prisma migrate deploy attempt %s timed out after %ss, retrying. "
|
||||
"Raise %s if this database needs longer to apply its pending migrations.",
|
||||
attempt + 1,
|
||||
deploy_timeout,
|
||||
PRISMA_MIGRATE_DEPLOY_TIMEOUT_ENV_VAR,
|
||||
)
|
||||
time.sleep(random.randrange(5, 15))
|
||||
continue
|
||||
|
|
@ -823,7 +831,8 @@ class ProxyExtrasDBManager:
|
|||
"Database migration failed after 4 attempts (retry loop "
|
||||
"exhausted by timeouts or repeated idempotent-recovery "
|
||||
"continues). Check database connectivity, load, and "
|
||||
"_prisma_migrations ledger state."
|
||||
"_prisma_migrations ledger state, and raise "
|
||||
f"{PRISMA_MIGRATE_DEPLOY_TIMEOUT_ENV_VAR} if the attempts timed out."
|
||||
)
|
||||
finally:
|
||||
os.chdir(original_dir)
|
||||
|
|
@ -908,7 +917,7 @@ class ProxyExtrasDBManager:
|
|||
# Set migrations directory for Prisma
|
||||
result = subprocess.run(
|
||||
[_get_prisma_command(), "migrate", "deploy"],
|
||||
timeout=prisma_command_timeout(),
|
||||
timeout=prisma_migrate_deploy_timeout(),
|
||||
check=True,
|
||||
capture_output=True,
|
||||
text=True,
|
||||
|
|
@ -1126,7 +1135,11 @@ class ProxyExtrasDBManager:
|
|||
)
|
||||
return True
|
||||
except subprocess.TimeoutExpired:
|
||||
logger.info(f"Attempt {attempt + 1} timed out")
|
||||
logger.warning(
|
||||
"Attempt %s timed out. Raise %s if this database needs longer to apply its schema.",
|
||||
attempt + 1,
|
||||
PRISMA_MIGRATE_DEPLOY_TIMEOUT_ENV_VAR if use_migrate else PRISMA_COMMAND_TIMEOUT_ENV_VAR,
|
||||
)
|
||||
time.sleep(random.randrange(5, 15))
|
||||
except subprocess.CalledProcessError as e:
|
||||
attempts_left = 3 - attempt
|
||||
|
|
|
|||
|
|
@ -10,8 +10,19 @@ from collections.abc import AsyncIterator, Mapping
|
|||
from typing import Any, Final
|
||||
|
||||
from litellm._logging import verbose_logger
|
||||
from litellm.a2a_protocol.litellm_completion_bridge.handler import (
|
||||
A2A_USER_API_KEY_HASH_PARAM,
|
||||
)
|
||||
from litellm.a2a_protocol.utils import (
|
||||
get_session_id_from_a2a_params,
|
||||
scope_session_to_principal,
|
||||
)
|
||||
from litellm.exceptions import BadRequestError
|
||||
from litellm.llms.bedrock.chat.agentcore.transformation import AmazonAgentCoreConfig
|
||||
|
||||
RUNTIME_SESSION_ID_MIN_LENGTH: Final = 33
|
||||
RUNTIME_SESSION_ID_MAX_LENGTH: Final = 256
|
||||
|
||||
# Reserved outbound header names that must never be sourced from per-request
|
||||
# ``agent_extra_headers`` for AgentCore requests. ``agent_extra_headers`` carries
|
||||
# values rewritten from the client-controlled ``x-a2a-{agent}-*`` convention, so
|
||||
|
|
@ -19,8 +30,9 @@ from litellm.llms.bedrock.chat.agentcore.transformation import AmazonAgentCoreCo
|
|||
# request identity / SigV4 metadata by overwriting headers the proxy sets from
|
||||
# trusted server-side config.
|
||||
#
|
||||
# The runtime headers (session / user id) are derived server-side from
|
||||
# ``runtimeSessionId`` / ``runtimeUserId`` in the agent's ``litellm_params``;
|
||||
# The runtime headers (session / user id) are derived server-side from the A2A
|
||||
# ``message.contextId`` and ``runtimeSessionId`` / ``runtimeUserId`` in the
|
||||
# agent's ``litellm_params``;
|
||||
# ``authorization`` is set by the AgentCore signer (JWT or SigV4); ``host`` and
|
||||
# the ``x-amz-*`` family are owned by SigV4 itself.
|
||||
_RESERVED_EXACT_HEADERS: Final = frozenset(
|
||||
|
|
@ -66,6 +78,31 @@ def _filter_reserved_headers(
|
|||
return filtered or None
|
||||
|
||||
|
||||
def _request_scoped_runtime_session_id(
|
||||
params: Mapping[str, Any],
|
||||
litellm_params: Mapping[str, Any],
|
||||
) -> str | None:
|
||||
context_id: Final = get_session_id_from_a2a_params(params)
|
||||
if not isinstance(context_id, str) or not context_id:
|
||||
return None
|
||||
return scope_session_to_principal(context_id, litellm_params.get(A2A_USER_API_KEY_HASH_PARAM))
|
||||
|
||||
|
||||
def _validate_runtime_session_id(session_id: str, model: str) -> str:
|
||||
if RUNTIME_SESSION_ID_MIN_LENGTH <= len(session_id) <= RUNTIME_SESSION_ID_MAX_LENGTH:
|
||||
return session_id
|
||||
raise BadRequestError(
|
||||
message=(
|
||||
f"Invalid AgentCore runtime session id {session_id!r}: AWS requires "
|
||||
f"{RUNTIME_SESSION_ID_MIN_LENGTH}-{RUNTIME_SESSION_ID_MAX_LENGTH} characters. It is built from the A2A "
|
||||
"message.contextId (prefixed with a 16-hex-char hash of the calling key and '-') when set, "
|
||||
"otherwise from the agent's configured runtimeSessionId."
|
||||
),
|
||||
model=model,
|
||||
llm_provider="bedrock",
|
||||
)
|
||||
|
||||
|
||||
class BedrockAgentCoreA2ATransformation:
|
||||
"""
|
||||
Request/response transformation for Bedrock AgentCore A2A agents.
|
||||
|
|
@ -100,7 +137,9 @@ class BedrockAgentCoreA2ATransformation:
|
|||
here to prevent a caller-controlled ``x-a2a-{agent}-*`` header from
|
||||
spoofing the AgentCore runtime user id or other SigV4 metadata. Use
|
||||
``api_key`` / ``runtimeUserId`` / ``runtimeSessionId`` in litellm_params
|
||||
(not ``agent_extra_headers``) to override those values.
|
||||
(not ``agent_extra_headers``) to override those values. The runtime
|
||||
session id is taken from ``params["message"]["contextId"]`` (scoped to
|
||||
the calling key) when present, then ``runtimeSessionId``, else generated.
|
||||
|
||||
Returns:
|
||||
Tuple of (url, signed_headers, signed_body_bytes)
|
||||
|
|
@ -139,7 +178,11 @@ class BedrockAgentCoreA2ATransformation:
|
|||
# Set required AgentCore session headers (normally set by transform_request,
|
||||
# which we skip because it also builds {"prompt": "..."})
|
||||
headers: Final[dict] = {}
|
||||
session_id: Final = agentcore_config._get_runtime_session_id(optional_params)
|
||||
session_id: Final = _validate_runtime_session_id(
|
||||
_request_scoped_runtime_session_id(params, litellm_params)
|
||||
or agentcore_config._get_runtime_session_id(optional_params),
|
||||
model=model,
|
||||
)
|
||||
headers["X-Amzn-Bedrock-AgentCore-Runtime-Session-Id"] = session_id
|
||||
runtime_user_id: Final = agentcore_config._get_runtime_user_id(optional_params)
|
||||
if runtime_user_id:
|
||||
|
|
|
|||
|
|
@ -2,6 +2,8 @@
|
|||
Utility functions for A2A protocol.
|
||||
"""
|
||||
|
||||
import hashlib
|
||||
from collections.abc import Mapping
|
||||
from typing import TYPE_CHECKING, Any, Final
|
||||
|
||||
import litellm
|
||||
|
|
@ -140,6 +142,29 @@ class A2ARequestUtils:
|
|||
return prompt_tokens, completion_tokens, total_tokens
|
||||
|
||||
|
||||
def get_session_id_from_a2a_params(params: Mapping[str, Any]) -> str | None:
|
||||
message: Final = params.get("message", {})
|
||||
if isinstance(message, dict):
|
||||
return message.get("contextId")
|
||||
return getattr(message, "contextId", None)
|
||||
|
||||
|
||||
def scope_session_to_principal(session_id: str, principal: str | None) -> str:
|
||||
"""
|
||||
Bind a client-supplied A2A contextId to the authenticated principal.
|
||||
|
||||
Without this, two distinct keys authorized for the same agent could set the
|
||||
same contextId and read/append to each other's backend memory. The
|
||||
principal is hashed (it is already a hashed token) so the raw value is never
|
||||
sent to the agent backend, while the original contextId is kept as a suffix
|
||||
for operator-side correlation.
|
||||
"""
|
||||
if not principal:
|
||||
return session_id
|
||||
principal_prefix: Final = hashlib.sha256(principal.encode("utf-8")).hexdigest()[:16]
|
||||
return f"{principal_prefix}-{session_id}"
|
||||
|
||||
|
||||
# Backwards compatibility aliases
|
||||
def extract_text_from_a2a_message(message: Any) -> str:
|
||||
return A2ARequestUtils.extract_text_from_message(message)
|
||||
|
|
|
|||
|
|
@ -1,7 +1,7 @@
|
|||
import asyncio
|
||||
import contextvars
|
||||
import json
|
||||
from collections.abc import Coroutine, Mapping
|
||||
from collections.abc import Callable, Coroutine, Mapping
|
||||
from functools import partial
|
||||
from typing import Final, Literal, overload
|
||||
|
||||
|
|
@ -47,6 +47,13 @@ __all__ = [
|
|||
|
||||
|
||||
##### Container Create #######################
|
||||
async def _encode_created_container_id(
|
||||
pending: Coroutine[object, object, ContainerObject],
|
||||
encode: Callable[[ContainerObject], ContainerObject],
|
||||
) -> ContainerObject:
|
||||
return encode(await pending)
|
||||
|
||||
|
||||
@client
|
||||
async def acreate_container(
|
||||
name: str,
|
||||
|
|
@ -256,16 +263,16 @@ def create_container(
|
|||
_is_async=_is_async,
|
||||
)
|
||||
|
||||
# Encode container_id with provider/model metadata for routing
|
||||
encode: Final = partial(
|
||||
ContainerRequestUtils.encode_container_id_in_response,
|
||||
custom_llm_provider=custom_llm_provider,
|
||||
litellm_metadata=kwargs.get("litellm_metadata"),
|
||||
extra_body=extra_body,
|
||||
)
|
||||
if isinstance(container_obj, ContainerObject):
|
||||
container_obj = ContainerRequestUtils.encode_container_id_in_response(
|
||||
response_obj=container_obj,
|
||||
custom_llm_provider=custom_llm_provider,
|
||||
litellm_metadata=kwargs.get("litellm_metadata"),
|
||||
extra_body=extra_body,
|
||||
)
|
||||
return encode(container_obj)
|
||||
|
||||
return container_obj
|
||||
return _encode_created_container_id(pending=container_obj, encode=encode)
|
||||
|
||||
except Exception as e:
|
||||
raise litellm.exception_type(
|
||||
|
|
|
|||
|
|
@ -1,28 +1,9 @@
|
|||
import hashlib
|
||||
from typing import Any, Final
|
||||
|
||||
|
||||
def get_session_id_from_a2a_params(params: dict[str, Any]) -> str | None:
|
||||
message: Final = params.get("message", {})
|
||||
if isinstance(message, dict):
|
||||
return message.get("contextId")
|
||||
return getattr(message, "contextId", None)
|
||||
|
||||
|
||||
def scope_session_to_principal(session_id: str, principal: str | None) -> str:
|
||||
"""
|
||||
Bind a client-supplied A2A contextId to the authenticated principal.
|
||||
|
||||
Without this, two distinct keys authorized for the same LangFlow agent could
|
||||
set the same contextId and read/append to each other's LangFlow memory. The
|
||||
principal is hashed (it is already a hashed token) so the raw value is never
|
||||
sent to the LangFlow backend, while the original contextId is kept as a
|
||||
suffix for operator-side correlation.
|
||||
"""
|
||||
if not principal:
|
||||
return session_id
|
||||
principal_prefix: Final = hashlib.sha256(principal.encode("utf-8")).hexdigest()[:16]
|
||||
return f"{principal_prefix}-{session_id}"
|
||||
from litellm.a2a_protocol.utils import (
|
||||
get_session_id_from_a2a_params,
|
||||
scope_session_to_principal,
|
||||
)
|
||||
|
||||
|
||||
def merge_a2a_session_into_litellm_params(
|
||||
|
|
|
|||
|
|
@ -1019,4 +1019,6 @@ async def invoke_agent_a2a(
|
|||
)
|
||||
except Exception:
|
||||
pass
|
||||
if isinstance(e, litellm.BadRequestError):
|
||||
return _jsonrpc_error(body.get("id"), -32602, e.message, 400)
|
||||
return _jsonrpc_error(body.get("id"), -32603, f"Internal error: {e}", 500)
|
||||
|
|
|
|||
|
|
@ -208,7 +208,7 @@ async def list_containers(
|
|||
|
||||
# Read query parameters
|
||||
query_params: Final = dict(request.query_params)
|
||||
data: Final[dict[str, Any]] = {"query_params": query_params}
|
||||
data: Final[dict[str, Any]] = {"query_params": query_params, "model": query_params.get("model")}
|
||||
|
||||
# Extract custom_llm_provider using priority chain
|
||||
custom_llm_provider: Final = (
|
||||
|
|
|
|||
|
|
@ -1,6 +1,6 @@
|
|||
import json
|
||||
import time
|
||||
from collections.abc import Iterable, Mapping
|
||||
from collections.abc import Iterable, Mapping, Sequence
|
||||
from dataclasses import dataclass
|
||||
from types import MappingProxyType
|
||||
from typing import Final
|
||||
|
|
@ -199,18 +199,10 @@ async def build_model_max_budget_usage(
|
|||
)
|
||||
for budget_model, budget_config in budgets
|
||||
)
|
||||
batched: Final = await cache.async_batch_get_cache(
|
||||
keys=list(spend_keys) # mutable-ok: async_batch_get_cache annotates keys as list, so one must exist here
|
||||
)
|
||||
# async_batch_get_cache returns None if it fails internally, and its result is
|
||||
# index-aligned with `keys` otherwise. An unusable result reads as a miss,
|
||||
# which is what a never-written counter already reads as.
|
||||
current_spends: Final = (
|
||||
tuple(batched) if isinstance(batched, list) and len(batched) == len(budgets) else (None,) * len(budgets)
|
||||
)
|
||||
current_spends: Final = await _current_window_spends(cache=cache, spend_keys=spend_keys)
|
||||
return {
|
||||
budget_model: {
|
||||
"current_spend": round(_as_spend(current_spend), 4),
|
||||
"current_spend": round(current_spend, 4),
|
||||
"budget_limit": budget_config.max_budget,
|
||||
"time_period": budget_config.budget_duration,
|
||||
}
|
||||
|
|
@ -218,6 +210,22 @@ async def build_model_max_budget_usage(
|
|||
}
|
||||
|
||||
|
||||
async def _current_window_spends(cache: DualCache, spend_keys: Sequence[str]) -> tuple[float, ...]:
|
||||
"""Redis holds the window total across replicas; the in-memory copy is one replica's share."""
|
||||
keys: Final = list(spend_keys) # mutable-ok: both batch readers annotate their key argument as list
|
||||
redis_cache: Final = cache.redis_cache
|
||||
if redis_cache is not None:
|
||||
shared: Final = await redis_cache.async_batch_get_cache(key_list=keys)
|
||||
return tuple(_as_spend(shared.get(key)) for key in keys)
|
||||
# async_batch_get_cache returns None if it fails internally, and its result is
|
||||
# index-aligned with `keys` otherwise. An unusable result reads as a miss,
|
||||
# which is what a never-written counter already reads as.
|
||||
batched: Final = await cache.async_batch_get_cache(keys=keys)
|
||||
if not isinstance(batched, list) or len(batched) != len(keys):
|
||||
return (0.0,) * len(keys)
|
||||
return tuple(_as_spend(current_spend) for current_spend in batched)
|
||||
|
||||
|
||||
def _usable_budget_config(raw_budget_config: object) -> BudgetConfig | None:
|
||||
try:
|
||||
budget_config: Final = BudgetConfig.model_validate(raw_budget_config)
|
||||
|
|
@ -404,7 +412,10 @@ class _PROXY_VirtualKeyModelMaxBudgetLimiter(RouterBudgetLimiting):
|
|||
return current_spend + _as_spend(await self._cached_spend(legacy_spend_key))
|
||||
|
||||
async def _cached_spend(self, spend_key: str) -> float | None:
|
||||
return await self.dual_cache.async_get_cache(key=spend_key)
|
||||
redis_cache: Final = self.dual_cache.redis_cache
|
||||
if redis_cache is None:
|
||||
return await self.dual_cache.async_get_cache(key=spend_key)
|
||||
return await redis_cache.async_get_cache(key=spend_key)
|
||||
|
||||
async def async_filter_deployments(
|
||||
self,
|
||||
|
|
|
|||
|
|
@ -2280,7 +2280,7 @@ user_api_key_cache: UserApiKeyCache = UserApiKeyCache(
|
|||
)
|
||||
spend_counter_cache: Final = DualCache(default_in_memory_ttl=UserAPIKeyCacheTTLEnum.in_memory_cache_ttl.value)
|
||||
cli_sso_session_cache: Final = DualCache(default_in_memory_ttl=CLI_SSO_SESSION_TTL_SECONDS)
|
||||
model_max_budget_limiter: Final = _PROXY_VirtualKeyModelMaxBudgetLimiter(dual_cache=user_api_key_cache)
|
||||
model_max_budget_limiter: Final = _PROXY_VirtualKeyModelMaxBudgetLimiter(dual_cache=spend_counter_cache)
|
||||
litellm.logging_callback_manager.add_litellm_callback(model_max_budget_limiter)
|
||||
redis_usage_cache: RedisCache | None = None # redis cache used for tracking spend, tpm/rpm limits
|
||||
polling_via_cache_enabled: Literal["all"] | list[str] | bool = False
|
||||
|
|
|
|||
|
|
@ -6714,7 +6714,10 @@ class Router:
|
|||
metadata. When present, decode the ID, replace ``container_id`` with the
|
||||
upstream value, and route through ``_ageneric_api_call_with_fallbacks`` so
|
||||
deployment credentials (e.g. regional ``api_base`` for Azure) match
|
||||
:meth:`_init_responses_api_endpoints`. Otherwise call the handler directly.
|
||||
:meth:`_init_responses_api_endpoints`. Create/list calls carry no container ID, so
|
||||
they route through the deployment named by ``model`` when the caller passes one,
|
||||
falling back to the direct call when no deployment matches. Otherwise call the
|
||||
handler directly with global provider credentials.
|
||||
"""
|
||||
if custom_llm_provider and "custom_llm_provider" not in kwargs:
|
||||
kwargs["custom_llm_provider"] = custom_llm_provider
|
||||
|
|
@ -6746,6 +6749,14 @@ class Router:
|
|||
**kwargs,
|
||||
)
|
||||
|
||||
requested_model: Final = kwargs.get("model")
|
||||
if isinstance(requested_model, str) and requested_model.strip():
|
||||
return await self._ageneric_api_call_with_fallbacks(
|
||||
original_function=original_function,
|
||||
passthrough_on_no_deployment=True,
|
||||
**kwargs,
|
||||
)
|
||||
|
||||
return await original_function(**kwargs)
|
||||
|
||||
async def _init_responses_api_endpoints(
|
||||
|
|
|
|||
|
|
@ -7,26 +7,36 @@ attempt fails identically. These tests pin the two behaviours that keep a
|
|||
container recoverable: an incomplete cache is deleted before Prisma is
|
||||
invoked, and the install gets a budget of its own rather than sharing the one
|
||||
that bounds each migration command.
|
||||
|
||||
``prisma migrate deploy`` gets a budget of its own for the same reason: its
|
||||
runtime grows with the number of pending migrations, so a fresh database that
|
||||
replays every migration overran the per-command budget on slow machines and
|
||||
the proxy gave up after four identical timeouts.
|
||||
"""
|
||||
|
||||
import ast
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
import sys
|
||||
import time
|
||||
from collections.abc import Callable
|
||||
from pathlib import Path
|
||||
|
||||
import pytest
|
||||
|
||||
from litellm_proxy_extras.prisma_toolchain import (
|
||||
DEFAULT_PRISMA_COMMAND_TIMEOUT,
|
||||
DEFAULT_PRISMA_MIGRATE_DEPLOY_TIMEOUT,
|
||||
PRISMA_BOOTSTRAP_TIMEOUT_ENV_VAR,
|
||||
PRISMA_COMMAND_TIMEOUT_ENV_VAR,
|
||||
PRISMA_MIGRATE_DEPLOY_TIMEOUT_ENV_VAR,
|
||||
ensure_prisma_toolchain,
|
||||
heal_incomplete_nodeenv_cache,
|
||||
node_binary_path,
|
||||
prisma_bootstrap_timeout,
|
||||
prisma_command_timeout,
|
||||
prisma_migrate_deploy_timeout,
|
||||
)
|
||||
from litellm_proxy_extras.utils import ProxyExtrasDBManager
|
||||
|
||||
|
|
@ -42,14 +52,27 @@ import time
|
|||
|
||||
args = sys.argv[1:]
|
||||
cache_dir = os.environ["PRISMA_NODEENV_CACHE_DIR"]
|
||||
with pathlib.Path(os.environ["FAKE_PRISMA_LOG"]).open("a") as log:
|
||||
log_path = pathlib.Path(os.environ["FAKE_PRISMA_LOG"])
|
||||
earlier_same_command = sum(
|
||||
1
|
||||
for line in (log_path.read_text().splitlines() if log_path.exists() else [])
|
||||
if json.loads(line)["args"][:2] == args[:2]
|
||||
)
|
||||
with log_path.open("a") as log:
|
||||
log.write(
|
||||
json.dumps({{"args": args, "cache_dir_present": os.path.isdir(cache_dir)}})
|
||||
+ "\\n"
|
||||
)
|
||||
time.sleep(float(os.environ.get("FAKE_PRISMA_SLEEP", "0")))
|
||||
if args[:2] == ["migrate", "deploy"]:
|
||||
if earlier_same_command == 0:
|
||||
time.sleep(float(os.environ.get("FAKE_PRISMA_FIRST_DEPLOY_SLEEP", "0")))
|
||||
elif os.environ.get("FAKE_PRISMA_LATER_DEPLOY_STDERR"):
|
||||
print(os.environ["FAKE_PRISMA_LATER_DEPLOY_STDERR"], file=sys.stderr)
|
||||
sys.exit(1)
|
||||
print("No pending migrations to apply")
|
||||
if args[:2] == ["db", "push"] and earlier_same_command == 0:
|
||||
time.sleep(float(os.environ.get("FAKE_PRISMA_FIRST_PUSH_SLEEP", "0")))
|
||||
sys.exit(0)
|
||||
"""
|
||||
|
||||
|
|
@ -80,9 +103,14 @@ def toolchain_env(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> tuple[Path
|
|||
monkeypatch.setenv("PATH", f"{bin_dir}{os.pathsep}{os.environ['PATH']}")
|
||||
monkeypatch.delenv(PRISMA_COMMAND_TIMEOUT_ENV_VAR, raising=False)
|
||||
monkeypatch.delenv(PRISMA_BOOTSTRAP_TIMEOUT_ENV_VAR, raising=False)
|
||||
monkeypatch.delenv(PRISMA_MIGRATE_DEPLOY_TIMEOUT_ENV_VAR, raising=False)
|
||||
return cache_dir, log_path
|
||||
|
||||
|
||||
def _deploy_calls(log_path: Path) -> list[list[str]]:
|
||||
return [call["args"] for call in _fake_prisma_calls(log_path) if call["args"][:2] == ["migrate", "deploy"]]
|
||||
|
||||
|
||||
def _make_incomplete_cache(cache_dir: Path) -> None:
|
||||
(cache_dir / "lib").mkdir(parents=True)
|
||||
(cache_dir / "bin").mkdir()
|
||||
|
|
@ -209,25 +237,111 @@ def test_setup_database_prepares_the_toolchain_before_migrating(
|
|||
assert calls[0]["cache_dir_present"] is False
|
||||
|
||||
|
||||
@pytest.mark.parametrize("use_v2_resolver", [False, True], ids=["v1", "v2"])
|
||||
def test_migrate_deploy_is_not_bounded_by_the_per_command_timeout(
|
||||
toolchain_env: tuple[Path, Path],
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
use_v2_resolver: bool,
|
||||
) -> None:
|
||||
"""A fresh database replays every migration, which takes longer than any bookkeeping command."""
|
||||
_, log_path = toolchain_env
|
||||
monkeypatch.setenv("DATABASE_URL", "postgresql://u:p@localhost:9/x")
|
||||
monkeypatch.setenv(PRISMA_COMMAND_TIMEOUT_ENV_VAR, "1")
|
||||
monkeypatch.setenv("FAKE_PRISMA_FIRST_DEPLOY_SLEEP", "3")
|
||||
|
||||
assert ProxyExtrasDBManager.setup_database(use_migrate=True, use_v2_resolver=use_v2_resolver) is True
|
||||
assert _deploy_calls(log_path) == [["migrate", "deploy"]]
|
||||
|
||||
|
||||
def test_migrate_deploy_stops_at_its_own_timeout(
|
||||
toolchain_env: tuple[Path, Path], monkeypatch: pytest.MonkeyPatch
|
||||
) -> None:
|
||||
"""The deploy budget still bounds a deploy that hangs, so boot cannot wait forever."""
|
||||
_, log_path = toolchain_env
|
||||
monkeypatch.setenv("DATABASE_URL", "postgresql://u:p@localhost:9/x")
|
||||
monkeypatch.setenv(PRISMA_MIGRATE_DEPLOY_TIMEOUT_ENV_VAR, "1")
|
||||
monkeypatch.setenv("FAKE_PRISMA_FIRST_DEPLOY_SLEEP", "60")
|
||||
monkeypatch.setenv("FAKE_PRISMA_LATER_DEPLOY_STDERR", "Error: P3018 permission denied for schema public")
|
||||
|
||||
started = time.monotonic()
|
||||
with pytest.raises(RuntimeError, match="insufficient permissions"):
|
||||
ProxyExtrasDBManager.setup_database(use_migrate=True, use_v2_resolver=True)
|
||||
elapsed = time.monotonic() - started
|
||||
|
||||
assert len(_deploy_calls(log_path)) == 2
|
||||
assert elapsed < 30
|
||||
|
||||
|
||||
def test_db_push_timeout_hint_names_the_per_command_budget(
|
||||
toolchain_env: tuple[Path, Path], monkeypatch: pytest.MonkeyPatch, caplog: pytest.LogCaptureFixture
|
||||
) -> None:
|
||||
"""``db push`` keeps the per-command budget, so its timeout hint has to name that variable."""
|
||||
_, log_path = toolchain_env
|
||||
monkeypatch.setenv("DATABASE_URL", "postgresql://u:p@localhost:9/x")
|
||||
monkeypatch.setenv(PRISMA_COMMAND_TIMEOUT_ENV_VAR, "1")
|
||||
monkeypatch.setenv("FAKE_PRISMA_FIRST_PUSH_SLEEP", "3")
|
||||
|
||||
with caplog.at_level(logging.WARNING, logger="litellm_proxy_extras"):
|
||||
assert ProxyExtrasDBManager.setup_database(use_migrate=False, use_v2_resolver=False) is True
|
||||
|
||||
assert [call["args"][:2] for call in _fake_prisma_calls(log_path)].count(["db", "push"]) == 2
|
||||
assert [record.getMessage() for record in caplog.records if "timed out" in record.getMessage()] == [
|
||||
f"Attempt 1 timed out. Raise {PRISMA_COMMAND_TIMEOUT_ENV_VAR} if this database needs longer to apply its schema."
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("command_timeout", "deploy_timeout", "expected"),
|
||||
[
|
||||
("900", None, 900.0),
|
||||
("12", None, DEFAULT_PRISMA_MIGRATE_DEPLOY_TIMEOUT),
|
||||
("900", "1200", 1200.0),
|
||||
("900", "300", 300.0),
|
||||
],
|
||||
ids=["raised_command_budget_carries_over", "lowered_command_budget_does_not", "override_wins_upward", "override_wins_downward"],
|
||||
)
|
||||
def test_migrate_deploy_budget_keeps_a_raised_command_budget(
|
||||
command_timeout: str, deploy_timeout: str | None, expected: float, monkeypatch: pytest.MonkeyPatch
|
||||
) -> None:
|
||||
"""Deployments that raised the per-command budget to survive a long deploy keep that budget for deploy."""
|
||||
monkeypatch.setenv(PRISMA_COMMAND_TIMEOUT_ENV_VAR, command_timeout)
|
||||
if deploy_timeout is None:
|
||||
monkeypatch.delenv(PRISMA_MIGRATE_DEPLOY_TIMEOUT_ENV_VAR, raising=False)
|
||||
else:
|
||||
monkeypatch.setenv(PRISMA_MIGRATE_DEPLOY_TIMEOUT_ENV_VAR, deploy_timeout)
|
||||
|
||||
assert prisma_migrate_deploy_timeout() == expected
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"raw",
|
||||
["", "0", "-5", "not-a-number", "nan", "inf", "-inf", "1e400"],
|
||||
)
|
||||
@pytest.mark.parametrize(
|
||||
("env_var", "read_timeout", "default"),
|
||||
[
|
||||
(PRISMA_COMMAND_TIMEOUT_ENV_VAR, prisma_command_timeout, DEFAULT_PRISMA_COMMAND_TIMEOUT),
|
||||
(PRISMA_MIGRATE_DEPLOY_TIMEOUT_ENV_VAR, prisma_migrate_deploy_timeout, DEFAULT_PRISMA_MIGRATE_DEPLOY_TIMEOUT),
|
||||
],
|
||||
ids=["command", "migrate_deploy"],
|
||||
)
|
||||
def test_unusable_timeout_override_falls_back_to_the_default(
|
||||
raw: str, monkeypatch: pytest.MonkeyPatch
|
||||
raw: str, env_var: str, read_timeout: Callable[[], float], default: float, monkeypatch: pytest.MonkeyPatch
|
||||
) -> None:
|
||||
"""A non-finite override would silently disable the timeout it configures."""
|
||||
monkeypatch.setenv(PRISMA_COMMAND_TIMEOUT_ENV_VAR, raw)
|
||||
monkeypatch.setenv(env_var, raw)
|
||||
|
||||
assert prisma_command_timeout() == DEFAULT_PRISMA_COMMAND_TIMEOUT
|
||||
assert read_timeout() == default
|
||||
|
||||
|
||||
def test_timeout_overrides_are_independent(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
monkeypatch.setenv(PRISMA_COMMAND_TIMEOUT_ENV_VAR, "12")
|
||||
monkeypatch.setenv(PRISMA_BOOTSTRAP_TIMEOUT_ENV_VAR, "900")
|
||||
monkeypatch.setenv(PRISMA_MIGRATE_DEPLOY_TIMEOUT_ENV_VAR, "1200")
|
||||
|
||||
assert prisma_command_timeout() == 12
|
||||
assert prisma_bootstrap_timeout() == 900
|
||||
assert prisma_migrate_deploy_timeout() == 1200
|
||||
|
||||
|
||||
@pytest.mark.parametrize("module", ["utils.py", "replica_identity.py"])
|
||||
|
|
|
|||
|
|
@ -1,3 +1,5 @@
|
|||
import asyncio
|
||||
from types import MappingProxyType
|
||||
from unittest.mock import AsyncMock, patch
|
||||
|
||||
|
||||
|
|
@ -5,6 +7,7 @@ import pytest
|
|||
|
||||
import litellm
|
||||
from litellm.caching.caching import DualCache
|
||||
from litellm.caching.redis_cache import RedisCache
|
||||
from datetime import datetime, timezone
|
||||
|
||||
from litellm.litellm_core_utils.duration_parser import duration_in_seconds
|
||||
|
|
@ -1332,3 +1335,85 @@ async def test_the_user_scope_has_no_pre_upgrade_counter_to_carry():
|
|||
await limiter.is_user_within_model_budget(
|
||||
user_id="u1", user_model_max_budget=model_max_budget, model="openai/gpt-4"
|
||||
)
|
||||
|
||||
|
||||
class _SharedFakeRedis(RedisCache):
|
||||
"""Stand-in for the one Redis every replica's DualCache is attached to.
|
||||
|
||||
Only the methods the limiter and DualCache call are implemented, and
|
||||
``super().__init__`` is skipped so no connection is opened.
|
||||
"""
|
||||
|
||||
def __init__(self):
|
||||
self._store = MappingProxyType({})
|
||||
|
||||
async def async_set_cache(self, key, value, **kwargs):
|
||||
self._store = MappingProxyType({**self._store, key: value})
|
||||
|
||||
async def async_get_cache(self, key, **kwargs):
|
||||
return self._store.get(key)
|
||||
|
||||
async def async_batch_get_cache(self, key_list, **kwargs):
|
||||
return {key: self._store.get(key) for key in key_list}
|
||||
|
||||
async def async_increment_pipeline(self, increment_list, **kwargs):
|
||||
for op in increment_list:
|
||||
total = self._store.get(op["key"], 0.0) + op["increment_value"]
|
||||
self._store = MappingProxyType({**self._store, op["key"]: total})
|
||||
return [self._store[op["key"]] for op in increment_list]
|
||||
|
||||
|
||||
async def _log_spend(limiter, *, key_hash, model_max_budget, response_cost):
|
||||
await limiter.async_log_success_event(
|
||||
_success_kwargs(
|
||||
model_group="gpt-4",
|
||||
response_cost=response_cost,
|
||||
key_hash=key_hash,
|
||||
key_model_max_budget=model_max_budget,
|
||||
),
|
||||
response_obj=None,
|
||||
start_time=None,
|
||||
end_time=None,
|
||||
)
|
||||
# The Redis push is scheduled as a task rather than awaited inline.
|
||||
await asyncio.gather(*(t for t in asyncio.all_tasks() if t is not asyncio.current_task()))
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_spend_logged_on_one_replica_is_enforced_and_reported_on_another():
|
||||
"""
|
||||
Each replica increments its own in-memory copy of the per-model counter and
|
||||
pushes the increment to the shared Redis, so only Redis holds the window's
|
||||
total. A replica that has served part of the traffic must still enforce and
|
||||
report the total, not its own share.
|
||||
|
||||
Regression: reads went to the in-memory tier first, so a replica whose local
|
||||
copy sat under the cap kept admitting requests and /key/info on it reported
|
||||
that local share, while the shared counter was already over the cap.
|
||||
"""
|
||||
shared_redis = _SharedFakeRedis()
|
||||
replica_a = _PROXY_VirtualKeyModelMaxBudgetLimiter(dual_cache=DualCache(redis_cache=shared_redis))
|
||||
replica_b = _PROXY_VirtualKeyModelMaxBudgetLimiter(dual_cache=DualCache(redis_cache=shared_redis))
|
||||
key_hash = "vk-shared"
|
||||
model_max_budget = {"gpt-4": {"budget_limit": 1.0, "time_period": "30d"}}
|
||||
user_api_key = UserAPIKeyAuth(token=key_hash, model_max_budget=model_max_budget)
|
||||
|
||||
await _log_spend(replica_b, key_hash=key_hash, model_max_budget=model_max_budget, response_cost=0.25)
|
||||
await _log_spend(replica_a, key_hash=key_hash, model_max_budget=model_max_budget, response_cost=0.5)
|
||||
await _log_spend(replica_a, key_hash=key_hash, model_max_budget=model_max_budget, response_cost=0.5)
|
||||
|
||||
with pytest.raises(litellm.BudgetExceededError):
|
||||
await replica_b.is_key_within_model_budget(user_api_key, "gpt-4")
|
||||
|
||||
usage_on_b = await build_model_max_budget_usage(
|
||||
entity_type=Litellm_EntityType.KEY,
|
||||
entity_id=key_hash,
|
||||
model_max_budget=model_max_budget,
|
||||
cache=replica_b.dual_cache,
|
||||
)
|
||||
assert usage_on_b["gpt-4"]["current_spend"] == 1.25
|
||||
|
||||
# Control: a replica that never served this key reads the same total.
|
||||
replica_c = _PROXY_VirtualKeyModelMaxBudgetLimiter(dual_cache=DualCache(redis_cache=shared_redis))
|
||||
with pytest.raises(litellm.BudgetExceededError):
|
||||
await replica_c.is_key_within_model_budget(user_api_key, "gpt-4")
|
||||
|
|
|
|||
|
|
@ -1324,6 +1324,98 @@ async def test_init_containers_api_endpoints_managed_id_without_model_id_applies
|
|||
assert call_kw["custom_llm_provider"] == "azure"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_init_containers_api_endpoints_create_with_model_uses_deployment_credentials(monkeypatch):
|
||||
"""
|
||||
``POST /v1/containers`` carries no container ID, so a ``model`` in the request
|
||||
body is the only way to pick a deployment. The upstream call must receive that
|
||||
deployment's ``api_key``/``api_base`` instead of falling back to the global
|
||||
``OPENAI_API_KEY`` (which may be unset on the proxy).
|
||||
"""
|
||||
monkeypatch.delenv("OPENAI_API_KEY", raising=False)
|
||||
router = Router(
|
||||
model_list=[
|
||||
{
|
||||
"model_name": "gpt-5.4",
|
||||
"litellm_params": {
|
||||
"model": "openai/gpt-5.4",
|
||||
"api_key": "sk-model-list-key",
|
||||
"api_base": "https://custom.openai.example/v1",
|
||||
},
|
||||
}
|
||||
]
|
||||
)
|
||||
mock_original_function = AsyncMock(return_value={"id": "cntr_test", "name": "Test Container"})
|
||||
|
||||
await router._init_containers_api_endpoints(
|
||||
original_function=mock_original_function,
|
||||
custom_llm_provider="openai",
|
||||
name="Test Container",
|
||||
model="gpt-5.4",
|
||||
)
|
||||
|
||||
mock_original_function.assert_called_once()
|
||||
call_kw = mock_original_function.call_args.kwargs
|
||||
assert call_kw["api_key"] == "sk-model-list-key"
|
||||
assert call_kw["api_base"] == "https://custom.openai.example/v1"
|
||||
assert call_kw["model"] == "openai/gpt-5.4"
|
||||
assert call_kw["name"] == "Test Container"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_init_containers_api_endpoints_create_without_model_calls_directly():
|
||||
"""
|
||||
Without ``model`` (or with ``model=None`` as the proxy forwards it), create/list
|
||||
must keep calling the handler directly with global provider credentials.
|
||||
"""
|
||||
router = Router(model_list=[])
|
||||
router._ageneric_api_call_with_fallbacks = AsyncMock()
|
||||
mock_original_function = AsyncMock(return_value={"id": "cntr_test"})
|
||||
|
||||
await router._init_containers_api_endpoints(
|
||||
original_function=mock_original_function,
|
||||
custom_llm_provider="openai",
|
||||
name="Test Container",
|
||||
model=None,
|
||||
)
|
||||
|
||||
router._ageneric_api_call_with_fallbacks.assert_not_called()
|
||||
mock_original_function.assert_called_once_with(custom_llm_provider="openai", name="Test Container", model=None)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_init_containers_api_endpoints_create_with_unknown_model_passes_through(monkeypatch):
|
||||
"""
|
||||
A ``model`` that names no configured deployment must not turn into a 400. The call
|
||||
falls through to the handler with the caller's model and no injected deployment
|
||||
credentials, matching the behaviour before model-based routing existed.
|
||||
"""
|
||||
monkeypatch.delenv("OPENAI_API_KEY", raising=False)
|
||||
router = Router(
|
||||
model_list=[
|
||||
{
|
||||
"model_name": "gpt-5.4",
|
||||
"litellm_params": {"model": "openai/gpt-5.4", "api_key": "sk-model-list-key"},
|
||||
}
|
||||
]
|
||||
)
|
||||
mock_original_function = AsyncMock(return_value={"id": "cntr_test"})
|
||||
|
||||
await router._init_containers_api_endpoints(
|
||||
original_function=mock_original_function,
|
||||
custom_llm_provider="openai",
|
||||
name="Test Container",
|
||||
model="does-not-exist",
|
||||
)
|
||||
|
||||
mock_original_function.assert_called_once()
|
||||
call_kw = mock_original_function.call_args.kwargs
|
||||
assert call_kw["model"] == "does-not-exist"
|
||||
assert call_kw["name"] == "Test Container"
|
||||
assert "api_key" not in call_kw
|
||||
assert "api_base" not in call_kw
|
||||
|
||||
|
||||
def test_router_model_group_encrypted_content_affinity_callback_registration():
|
||||
from litellm.router_utils.pre_call_checks.deployment_affinity_check import (
|
||||
DeploymentAffinityCheck,
|
||||
|
|
|
|||
|
|
@ -11,7 +11,9 @@ Verifies that:
|
|||
|
||||
import json
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
import respx
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
|
||||
|
|
@ -295,6 +297,195 @@ class TestTransformation:
|
|||
assert headers["Authorization"].startswith("AWS4-HMAC-SHA256")
|
||||
|
||||
|
||||
SESSION_HEADER = "X-Amzn-Bedrock-AgentCore-Runtime-Session-Id"
|
||||
CONTEXT_ID = "conversation-alpha-0001-0000000000000000"
|
||||
KEY_HASH = "hashed-key-of-caller-one"
|
||||
|
||||
|
||||
def _params_with_context(context_id: object) -> dict:
|
||||
return {"message": {**SAMPLE_PARAMS["message"], "contextId": context_id}}
|
||||
|
||||
|
||||
def _scoped(context_id: str, key_hash: str) -> str:
|
||||
import hashlib
|
||||
|
||||
return f"{hashlib.sha256(key_hash.encode()).hexdigest()[:16]}-{context_id}"
|
||||
|
||||
|
||||
def _session_header(params: dict, litellm_params: dict) -> str:
|
||||
from litellm.a2a_protocol.providers.bedrock_agentcore.transformation import (
|
||||
BedrockAgentCoreA2ATransformation,
|
||||
)
|
||||
|
||||
_, headers, _ = BedrockAgentCoreA2ATransformation.get_url_and_signed_request(
|
||||
request_id="req-001",
|
||||
params=params,
|
||||
litellm_params=litellm_params,
|
||||
)
|
||||
return headers[SESSION_HEADER]
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def httpx_transport(monkeypatch):
|
||||
import litellm
|
||||
|
||||
monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
|
||||
monkeypatch.setenv("DISABLE_AIOHTTP_TRANSPORT", "True")
|
||||
litellm.in_memory_llm_clients_cache.flush_cache()
|
||||
yield
|
||||
litellm.in_memory_llm_clients_cache.flush_cache()
|
||||
|
||||
|
||||
class TestRequestScopedRuntimeSession:
|
||||
"""message.contextId selects the AgentCore runtime session, scoped to the calling key."""
|
||||
|
||||
def test_context_id_scoped_to_calling_key(self):
|
||||
from litellm.a2a_protocol.litellm_completion_bridge.handler import (
|
||||
A2A_USER_API_KEY_HASH_PARAM,
|
||||
)
|
||||
|
||||
litellm_params = {**SAMPLE_LITELLM_PARAMS, A2A_USER_API_KEY_HASH_PARAM: KEY_HASH}
|
||||
assert _session_header(_params_with_context(CONTEXT_ID), litellm_params) == _scoped(CONTEXT_ID, KEY_HASH)
|
||||
|
||||
def test_context_id_used_verbatim_without_principal(self):
|
||||
assert _session_header(_params_with_context(CONTEXT_ID), SAMPLE_LITELLM_PARAMS) == CONTEXT_ID
|
||||
|
||||
def test_same_context_id_reuses_session_and_other_context_isolated(self):
|
||||
first = _session_header(_params_with_context(CONTEXT_ID), SAMPLE_LITELLM_PARAMS)
|
||||
second = _session_header(_params_with_context(CONTEXT_ID), SAMPLE_LITELLM_PARAMS)
|
||||
other = _session_header(
|
||||
_params_with_context("conversation-beta-00002-0000000000000000"),
|
||||
SAMPLE_LITELLM_PARAMS,
|
||||
)
|
||||
assert first == second
|
||||
assert other != first
|
||||
|
||||
def test_same_context_id_from_different_keys_is_isolated(self):
|
||||
from litellm.a2a_protocol.litellm_completion_bridge.handler import (
|
||||
A2A_USER_API_KEY_HASH_PARAM,
|
||||
)
|
||||
|
||||
params = _params_with_context(CONTEXT_ID)
|
||||
caller_one = _session_header(params, {**SAMPLE_LITELLM_PARAMS, A2A_USER_API_KEY_HASH_PARAM: KEY_HASH})
|
||||
caller_two = _session_header(
|
||||
params, {**SAMPLE_LITELLM_PARAMS, A2A_USER_API_KEY_HASH_PARAM: "hashed-key-of-caller-two"}
|
||||
)
|
||||
assert caller_one != caller_two
|
||||
assert caller_one.endswith(f"-{CONTEXT_ID}")
|
||||
assert caller_two.endswith(f"-{CONTEXT_ID}")
|
||||
|
||||
def test_context_id_takes_precedence_over_configured_session(self):
|
||||
litellm_params = {**SAMPLE_LITELLM_PARAMS, "runtimeSessionId": "a" * 40}
|
||||
assert _session_header(_params_with_context(CONTEXT_ID), litellm_params) == CONTEXT_ID
|
||||
|
||||
def test_configured_session_is_fallback_without_context_id(self):
|
||||
litellm_params = {**SAMPLE_LITELLM_PARAMS, "runtimeSessionId": "a" * 40}
|
||||
assert _session_header(SAMPLE_PARAMS, litellm_params) == "a" * 40
|
||||
assert _session_header(_params_with_context(""), litellm_params) == "a" * 40
|
||||
|
||||
def test_no_context_id_and_no_config_generates_new_session_per_request(self):
|
||||
first = _session_header(SAMPLE_PARAMS, SAMPLE_LITELLM_PARAMS)
|
||||
second = _session_header(SAMPLE_PARAMS, SAMPLE_LITELLM_PARAMS)
|
||||
assert first != second
|
||||
assert 33 <= len(first) <= 256
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"context_id",
|
||||
[
|
||||
"short-context-id",
|
||||
"x" * 257,
|
||||
],
|
||||
)
|
||||
def test_invalid_context_id_rejected_with_clear_error(self, context_id):
|
||||
import litellm
|
||||
|
||||
with pytest.raises(litellm.BadRequestError, match="Invalid AgentCore runtime session id") as exc_info:
|
||||
_session_header(_params_with_context(context_id), SAMPLE_LITELLM_PARAMS)
|
||||
assert exc_info.value.status_code == 400
|
||||
assert "33-256" in str(exc_info.value)
|
||||
|
||||
def test_scoped_context_id_shorter_than_33_rejected(self):
|
||||
import litellm
|
||||
from litellm.a2a_protocol.litellm_completion_bridge.handler import (
|
||||
A2A_USER_API_KEY_HASH_PARAM,
|
||||
)
|
||||
|
||||
litellm_params = {**SAMPLE_LITELLM_PARAMS, A2A_USER_API_KEY_HASH_PARAM: KEY_HASH}
|
||||
with pytest.raises(litellm.BadRequestError, match=_scoped("c" * 15, KEY_HASH)):
|
||||
_session_header(_params_with_context("c" * 15), litellm_params)
|
||||
assert _session_header(_params_with_context("c" * 16), litellm_params) == _scoped("c" * 16, KEY_HASH)
|
||||
|
||||
def test_invalid_configured_session_rejected(self):
|
||||
import litellm
|
||||
|
||||
litellm_params = {**SAMPLE_LITELLM_PARAMS, "runtimeSessionId": "too-short"}
|
||||
with pytest.raises(litellm.BadRequestError, match="Invalid AgentCore runtime session id"):
|
||||
_session_header(SAMPLE_PARAMS, litellm_params)
|
||||
|
||||
def test_non_string_context_id_falls_back(self):
|
||||
litellm_params = {**SAMPLE_LITELLM_PARAMS, "runtimeSessionId": "a" * 40}
|
||||
assert _session_header(_params_with_context(12345), litellm_params) == "a" * 40
|
||||
|
||||
def test_spoofed_session_header_does_not_override_context_id(self):
|
||||
from litellm.a2a_protocol.providers.bedrock_agentcore.transformation import (
|
||||
BedrockAgentCoreA2ATransformation,
|
||||
)
|
||||
|
||||
_, headers, _ = BedrockAgentCoreA2ATransformation.get_url_and_signed_request(
|
||||
request_id="req-001",
|
||||
params=_params_with_context(CONTEXT_ID),
|
||||
litellm_params=SAMPLE_LITELLM_PARAMS,
|
||||
agent_extra_headers={SESSION_HEADER: "s" * 40},
|
||||
)
|
||||
assert headers[SESSION_HEADER] == CONTEXT_ID
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_context_id_session_header_on_outbound_non_streaming_post(self, httpx_transport):
|
||||
from litellm.a2a_protocol.litellm_completion_bridge.handler import (
|
||||
A2A_USER_API_KEY_HASH_PARAM,
|
||||
)
|
||||
from litellm.a2a_protocol.providers.bedrock_agentcore.config import (
|
||||
BedrockAgentCoreA2AConfig,
|
||||
)
|
||||
|
||||
with respx.mock(assert_all_called=True) as router:
|
||||
route = router.post(url__regex=r".*/invocations.*").mock(
|
||||
return_value=httpx.Response(200, json={"jsonrpc": "2.0", "id": "req-001", "result": {}})
|
||||
)
|
||||
await BedrockAgentCoreA2AConfig().handle_non_streaming(
|
||||
request_id="req-001",
|
||||
params=_params_with_context(CONTEXT_ID),
|
||||
litellm_params={**SAMPLE_LITELLM_PARAMS, A2A_USER_API_KEY_HASH_PARAM: KEY_HASH},
|
||||
)
|
||||
|
||||
assert route.calls.last.request.headers[SESSION_HEADER] == _scoped(CONTEXT_ID, KEY_HASH)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_context_id_session_header_on_outbound_streaming_post(self, httpx_transport):
|
||||
from litellm.a2a_protocol.litellm_completion_bridge.handler import (
|
||||
A2A_USER_API_KEY_HASH_PARAM,
|
||||
)
|
||||
from litellm.a2a_protocol.providers.bedrock_agentcore.config import (
|
||||
BedrockAgentCoreA2AConfig,
|
||||
)
|
||||
|
||||
with respx.mock(assert_all_called=True) as router:
|
||||
route = router.post(url__regex=r".*/invocations.*").mock(
|
||||
return_value=httpx.Response(200, json={"jsonrpc": "2.0", "id": "req-001", "result": {}})
|
||||
)
|
||||
events = [
|
||||
event
|
||||
async for event in BedrockAgentCoreA2AConfig().handle_streaming(
|
||||
request_id="req-001",
|
||||
params=_params_with_context(CONTEXT_ID),
|
||||
litellm_params={**SAMPLE_LITELLM_PARAMS, A2A_USER_API_KEY_HASH_PARAM: KEY_HASH},
|
||||
)
|
||||
]
|
||||
|
||||
assert events == [{"jsonrpc": "2.0", "id": "req-001", "result": {}}]
|
||||
assert route.calls.last.request.headers[SESSION_HEADER] == _scoped(CONTEXT_ID, KEY_HASH)
|
||||
|
||||
|
||||
class TestNonStreaming:
|
||||
"""Test end-to-end non-streaming flow."""
|
||||
|
||||
|
|
|
|||
|
|
@ -152,6 +152,42 @@ class TestContainerAPI:
|
|||
assert response.id == "cntr_async_123"
|
||||
assert response.name == "Async Test Container"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_acreate_container_encodes_router_model_id(self):
|
||||
"""
|
||||
The async handler returns a coroutine, so the managed-ID encoding must run
|
||||
after it resolves. Otherwise follow-up calls (retrieve/delete/files) lose the
|
||||
deployment and fall back to global provider credentials.
|
||||
"""
|
||||
upstream_response = ContainerObject(
|
||||
id="cntr_upstream_123",
|
||||
object="container",
|
||||
created_at=1747857508,
|
||||
status="running",
|
||||
expires_after={"anchor": "last_active_at", "minutes": 20},
|
||||
last_active_at=1747857508,
|
||||
name="Routed Container",
|
||||
)
|
||||
|
||||
async def _resolve_upstream():
|
||||
return upstream_response
|
||||
|
||||
with patch.object( # test-quality-ok: create_container exposes no client seam, only the handler
|
||||
base_llm_http_handler,
|
||||
"container_create_handler",
|
||||
side_effect=lambda **kwargs: _resolve_upstream() if kwargs["_is_async"] else upstream_response,
|
||||
):
|
||||
response = await acreate_container(
|
||||
name="Routed Container",
|
||||
custom_llm_provider="openai",
|
||||
litellm_metadata={"model_info": {"id": "deployment-abc"}},
|
||||
)
|
||||
|
||||
decoded = ResponsesAPIRequestUtils._decode_container_id(response.id)
|
||||
assert decoded["model_id"] == "deployment-abc"
|
||||
assert decoded["custom_llm_provider"] == "openai"
|
||||
assert decoded["response_id"] == "cntr_upstream_123"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_alist_containers_basic(self):
|
||||
"""Test basic async container listing functionality."""
|
||||
|
|
|
|||
|
|
@ -918,6 +918,62 @@ async def test_task_method_failure_hook_uses_enriched_request_data():
|
|||
assert failure_data.get("agent_id") == "test-agent"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_agentcore_invalid_context_id_returns_jsonrpc_invalid_params_400():
|
||||
from litellm.proxy._types import UserAPIKeyAuth
|
||||
|
||||
agent = _make_agent_mock()
|
||||
agent.litellm_params = {
|
||||
"custom_llm_provider": "bedrock",
|
||||
"model": "bedrock/agentcore/arn:aws:bedrock-agentcore:us-west-2:123456789012:runtime/demo",
|
||||
"api_key": "test-jwt-token",
|
||||
}
|
||||
mock_request = _make_request_mock(
|
||||
"message/send",
|
||||
{
|
||||
"message": {
|
||||
"role": "user",
|
||||
"parts": [{"kind": "text", "text": "Hello"}],
|
||||
"messageId": "msg-1",
|
||||
"contextId": "too-short",
|
||||
}
|
||||
},
|
||||
)
|
||||
user_api_key_dict = UserAPIKeyAuth(api_key="sk-test", user_id="u1", team_id="t1")
|
||||
|
||||
mock_proxy_logging = MagicMock()
|
||||
mock_proxy_logging.pre_call_hook = AsyncMock(
|
||||
side_effect=lambda user_api_key_dict, data, call_type: data
|
||||
)
|
||||
mock_proxy_logging.post_call_failure_hook = AsyncMock(return_value=None)
|
||||
|
||||
with ExitStack() as stack:
|
||||
for p in _base_patches(agent):
|
||||
stack.enter_context(p)
|
||||
stack.enter_context(
|
||||
patch( # test-quality-ok: same proxy_logging_obj injection the sibling failure-hook test uses; no HTTP call is made because the request is rejected before signing
|
||||
"litellm.proxy.proxy_server.proxy_logging_obj", mock_proxy_logging
|
||||
)
|
||||
)
|
||||
|
||||
from litellm.proxy.agent_endpoints.a2a_endpoints import invoke_agent_a2a
|
||||
|
||||
response = await invoke_agent_a2a(
|
||||
agent_id="test-agent",
|
||||
request=mock_request,
|
||||
fastapi_response=MagicMock(),
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
)
|
||||
|
||||
body = json.loads(response.body.decode())
|
||||
assert response.status_code == 400
|
||||
assert body["id"] == "req-1"
|
||||
assert body["error"]["code"] == -32602
|
||||
assert "Invalid AgentCore runtime session id" in body["error"]["message"]
|
||||
assert "Internal error" not in body["error"]["message"]
|
||||
mock_proxy_logging.post_call_failure_hook.assert_awaited_once()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_extended_agent_card_rewrites_url():
|
||||
from litellm.proxy._types import UserAPIKeyAuth
|
||||
|
|
|
|||
|
|
@ -5,7 +5,7 @@ Verifies that _init_cache attaches Redis to user_api_key_cache only when
|
|||
the flag is explicitly set to True, and leaves it in-memory-only otherwise.
|
||||
"""
|
||||
|
||||
from contextlib import contextmanager
|
||||
from contextlib import ExitStack, contextmanager
|
||||
import json
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
|
|
@ -167,3 +167,25 @@ class TestRedisAuthCacheFlag:
|
|||
f"cli_sso_session_cache must always get Redis "
|
||||
f"(enable_redis_auth_cache={flag_value!r})"
|
||||
)
|
||||
|
||||
def test_flag_absent_still_shares_the_model_budget_counters_over_redis(self):
|
||||
"""
|
||||
Per-model budget counters are spend counters: the limiter must be able to
|
||||
push and read them through Redis without the auth-cache opt-in, or every
|
||||
worker enforces and reports its own share of a key's spend
|
||||
"""
|
||||
fake_redis = _FakeRedisCache()
|
||||
limiter_cache = ps.model_max_budget_limiter.dual_cache
|
||||
touched_caches = (
|
||||
limiter_cache,
|
||||
ps.spend_counter_cache,
|
||||
ps.cli_sso_session_cache,
|
||||
ps.user_api_key_cache,
|
||||
ps.litellm_config_cache,
|
||||
)
|
||||
with ExitStack() as detached:
|
||||
for cache in touched_caches:
|
||||
detached.enter_context(patch.object(cache, "redis_cache", None))
|
||||
ps._attach_redis_usage_cache(fake_redis, enable_redis_auth_cache=False)
|
||||
assert limiter_cache.redis_cache is fake_redis
|
||||
assert ps.user_api_key_cache.redis_cache is None
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue