mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-21 00:21:49 +00:00
Merge remote-tracking branch 'origin/main' into litellm_transcribe_passthrough
This commit is contained in:
commit
291a43c409
95 changed files with 6671 additions and 573 deletions
|
|
@ -2983,7 +2983,7 @@ workflows:
|
|||
name: integration-<< matrix.suite >>
|
||||
matrix:
|
||||
parameters:
|
||||
suite: [management, accounting, providers]
|
||||
suite: [management, accounting, database, providers]
|
||||
filters:
|
||||
branches:
|
||||
only:
|
||||
|
|
|
|||
|
|
@ -244,6 +244,7 @@ telemetry = True
|
|||
max_tokens: int = DEFAULT_MAX_TOKENS # OpenAI Defaults
|
||||
drop_params = drop_params_env_flag(os.environ, verbose_logger)
|
||||
modify_params = bool(os.getenv("LITELLM_MODIFY_PARAMS", False))
|
||||
bedrock_neutralize_orphaned_tool_blocks: bool = True
|
||||
use_chat_completions_url_for_anthropic_messages: bool = bool(
|
||||
os.getenv("LITELLM_USE_CHAT_COMPLETIONS_URL_FOR_ANTHROPIC_MESSAGES", False)
|
||||
) # When True, routes OpenAI /v1/messages requests to chat/completions instead of the Responses API
|
||||
|
|
@ -1818,6 +1819,9 @@ if TYPE_CHECKING:
|
|||
from .llms.azure.responses.o_series_transformation import (
|
||||
AzureOpenAIOSeriesResponsesAPIConfig as AzureOpenAIOSeriesResponsesAPIConfig,
|
||||
)
|
||||
from .llms.azure_ai.responses.transformation import (
|
||||
AzureAIResponsesAPIConfig as AzureAIResponsesAPIConfig,
|
||||
)
|
||||
from .llms.xai.responses.transformation import (
|
||||
XAIResponsesAPIConfig as XAIResponsesAPIConfig,
|
||||
)
|
||||
|
|
|
|||
|
|
@ -234,6 +234,7 @@ LLM_CONFIG_NAMES: Final = (
|
|||
"OpenAIResponsesAPIConfig",
|
||||
"AzureOpenAIResponsesAPIConfig",
|
||||
"AzureOpenAIOSeriesResponsesAPIConfig",
|
||||
"AzureAIResponsesAPIConfig",
|
||||
"XAIResponsesAPIConfig",
|
||||
"LiteLLMProxyResponsesAPIConfig",
|
||||
"HostedVLLMResponsesAPIConfig",
|
||||
|
|
@ -946,6 +947,10 @@ _LLM_CONFIGS_IMPORT_MAP: Final = {
|
|||
".llms.azure.responses.o_series_transformation",
|
||||
"AzureOpenAIOSeriesResponsesAPIConfig",
|
||||
),
|
||||
"AzureAIResponsesAPIConfig": (
|
||||
".llms.azure_ai.responses.transformation",
|
||||
"AzureAIResponsesAPIConfig",
|
||||
),
|
||||
"XAIResponsesAPIConfig": (
|
||||
".llms.xai.responses.transformation",
|
||||
"XAIResponsesAPIConfig",
|
||||
|
|
|
|||
|
|
@ -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())
|
||||
|
|
|
|||
|
|
@ -148,6 +148,7 @@ class _ToolCallChunk(TypedDict):
|
|||
class _UsageBearingChunk(TypedDict, total=False):
|
||||
usage: Usage | None
|
||||
_hidden_params: Mapping[str, str]
|
||||
choices: ReadOnly[Sequence[StreamingChoices | Mapping[str, object]]]
|
||||
|
||||
|
||||
class _UsageSummary(TypedDict):
|
||||
|
|
@ -921,21 +922,22 @@ class ChunkProcessor:
|
|||
|
||||
prompt_tokens_details = attach_cache_creation_token_details(prompt_tokens_details, cache_creation_token_details)
|
||||
|
||||
completion_tokens = self._reset_anthropic_cursor_completion_tokens(
|
||||
recovered_completion_tokens: Final = self._reset_anthropic_cursor_completion_tokens(
|
||||
chunks=chunks,
|
||||
completion_tokens=completion_tokens,
|
||||
completion_usage_updates=completion_usage_updates,
|
||||
)
|
||||
cursor_was_reset: Final = recovered_completion_tokens != completion_tokens
|
||||
|
||||
return UsagePerChunk(
|
||||
prompt_tokens=prompt_tokens,
|
||||
completion_tokens=completion_tokens,
|
||||
completion_tokens=recovered_completion_tokens,
|
||||
cache_creation_input_tokens=cache_creation_input_tokens,
|
||||
cache_read_input_tokens=cache_read_input_tokens,
|
||||
server_tool_use=server_tool_use,
|
||||
web_search_requests=web_search_requests,
|
||||
google_maps_grounding_requests=google_maps_grounding_requests,
|
||||
completion_tokens_details=completion_tokens_details,
|
||||
completion_tokens_details=None if cursor_was_reset else completion_tokens_details,
|
||||
prompt_tokens_details=prompt_tokens_details,
|
||||
cost=cost,
|
||||
inference_geo=self._last_provider_pricing_field(chunks, "inference_geo"),
|
||||
|
|
@ -960,6 +962,30 @@ class ChunkProcessor:
|
|||
]
|
||||
return values[-1] if values else None
|
||||
|
||||
@staticmethod
|
||||
def _finish_reason_of_choice(choice: object) -> str | None:
|
||||
match choice:
|
||||
case StreamingChoices(finish_reason=reason) | Choices(finish_reason=reason):
|
||||
return reason
|
||||
case {"finish_reason": str() as reason}:
|
||||
return reason
|
||||
case _:
|
||||
return None
|
||||
|
||||
@staticmethod
|
||||
def _chunk_choices(chunk: "_UsageBearingChunk | ModelResponse | ModelResponseStream") -> Sequence[object]:
|
||||
if isinstance(chunk, dict):
|
||||
return chunk.get("choices", ())
|
||||
return getattr(chunk, "choices", ())
|
||||
|
||||
@staticmethod
|
||||
def _saw_finish_reason(chunks: Sequence["_UsageBearingChunk | ModelResponse"]) -> bool:
|
||||
return any(
|
||||
ChunkProcessor._finish_reason_of_choice(choice) is not None
|
||||
for chunk in chunks
|
||||
for choice in ChunkProcessor._chunk_choices(chunk)
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _reset_anthropic_cursor_completion_tokens(
|
||||
chunks: Sequence["_UsageBearingChunk | ModelResponse"],
|
||||
|
|
@ -970,18 +996,18 @@ class ChunkProcessor:
|
|||
|
||||
See the ``completion_usage_updates`` comment in
|
||||
``_calculate_usage_per_chunk``. The accumulated value is NOT a stale
|
||||
cursor when either it is > 1 (definitely not a placeholder) or we saw
|
||||
>= 2 completion-bearing usage events (positive evidence ``message_delta``
|
||||
arrived). Otherwise — the only completion update we ever saw was the
|
||||
Anthropic ``message_start`` cursor (=1) — reset to 0 so
|
||||
``calculate_usage()``'s ``or token_counter(text=...)`` fallback estimates
|
||||
from the actually-received completion text instead of trusting the
|
||||
placeholder. Gated on ``custom_llm_provider == "anthropic"`` so the
|
||||
heuristic (which encodes Anthropic's specific message_start SSE shape)
|
||||
does not silently affect other providers that may legitimately report
|
||||
``completion_tokens=1`` from a single usage event.
|
||||
cursor when we saw >= 2 completion-bearing usage events or any chunk
|
||||
carried a ``finish_reason`` (positive evidence ``message_delta``
|
||||
arrived). Otherwise the only completion update we ever saw was the
|
||||
Anthropic ``message_start`` cursor, a small placeholder whose magnitude
|
||||
varies per request (1 and 8 both observed live), so reset to 0 and let
|
||||
``calculate_usage()``'s ``or token_counter(...)`` fallback estimate from
|
||||
the actually-received text and reasoning instead. Gated on
|
||||
``custom_llm_provider == "anthropic"`` so the heuristic (which encodes
|
||||
Anthropic's specific message_start SSE shape) does not silently affect
|
||||
other providers that legitimately report usage from a single event.
|
||||
"""
|
||||
saw_non_cursor_completion: Final = completion_tokens > 1 or completion_usage_updates >= 2
|
||||
saw_non_cursor_completion: Final = completion_usage_updates >= 2 or ChunkProcessor._saw_finish_reason(chunks)
|
||||
if saw_non_cursor_completion:
|
||||
return completion_tokens
|
||||
|
||||
|
|
@ -995,7 +1021,7 @@ class ChunkProcessor:
|
|||
if isinstance(hp, dict):
|
||||
custom_llm_provider = hp.get("custom_llm_provider")
|
||||
|
||||
if custom_llm_provider == "anthropic" and completion_tokens == 1:
|
||||
if custom_llm_provider == "anthropic":
|
||||
return 0
|
||||
return completion_tokens
|
||||
|
||||
|
|
@ -1039,10 +1065,13 @@ class ChunkProcessor:
|
|||
returned_usage.prompt_tokens = 0
|
||||
returned_usage.completion_tokens = (
|
||||
completion_tokens
|
||||
or token_counter(
|
||||
model=model,
|
||||
text=completion_output,
|
||||
count_response_tokens=True, # count_response_tokens is a Flag to tell token counter this is a response, No need to add extra tokens we do for input messages
|
||||
or (
|
||||
token_counter(
|
||||
model=model,
|
||||
text=completion_output,
|
||||
count_response_tokens=True, # count_response_tokens is a Flag to tell token counter this is a response, No need to add extra tokens we do for input messages
|
||||
)
|
||||
+ (reasoning_tokens or 0)
|
||||
)
|
||||
)
|
||||
returned_usage.total_tokens = returned_usage.prompt_tokens + returned_usage.completion_tokens
|
||||
|
|
@ -1066,15 +1095,16 @@ class ChunkProcessor:
|
|||
returned_usage.completion_tokens_details = completion_tokens_details
|
||||
|
||||
if reasoning_tokens is not None:
|
||||
capped_reasoning_tokens: Final = min(max(0, reasoning_tokens), returned_usage.completion_tokens)
|
||||
if returned_usage.completion_tokens_details is None:
|
||||
returned_usage.completion_tokens_details = CompletionTokensDetailsWrapper(
|
||||
reasoning_tokens=reasoning_tokens
|
||||
reasoning_tokens=capped_reasoning_tokens,
|
||||
text_tokens=returned_usage.completion_tokens - capped_reasoning_tokens,
|
||||
)
|
||||
elif (
|
||||
returned_usage.completion_tokens_details is not None
|
||||
and returned_usage.completion_tokens_details.reasoning_tokens is None
|
||||
):
|
||||
capped_reasoning_tokens: Final = min(max(0, reasoning_tokens), returned_usage.completion_tokens)
|
||||
returned_usage.completion_tokens_details.reasoning_tokens = capped_reasoning_tokens
|
||||
if returned_usage.completion_tokens_details.text_tokens is None:
|
||||
returned_usage.completion_tokens_details.text_tokens = (
|
||||
|
|
|
|||
|
|
@ -49,12 +49,7 @@ class AzureOpenAIResponsesAPIConfig(OpenAIResponsesAPIConfig):
|
|||
return BaseAzureLLM._base_validate_azure_environment(headers=headers, litellm_params=litellm_params)
|
||||
|
||||
def get_stripped_model_name(self, model: str) -> str:
|
||||
# if "responses/" is in the model name, remove it
|
||||
if "responses/" in model:
|
||||
model = model.replace("responses/", "")
|
||||
if "o_series" in model:
|
||||
model = model.replace("o_series/", "")
|
||||
return model
|
||||
return model.replace("responses/", "").replace("o_series/", "").replace("azure_ai/", "")
|
||||
|
||||
def _handle_reasoning_item(self, item: dict[str, Any]) -> dict[str, Any]:
|
||||
"""
|
||||
|
|
|
|||
|
|
@ -9,6 +9,7 @@ from litellm.types.llms.openai import AllMessageValues
|
|||
from litellm.types.router import GenericLiteLLMParams
|
||||
|
||||
AzureAIApiKeyHeader = Literal["Authorization", "api-key", "Api-Key", "Ocp-Apim-Subscription-Key"]
|
||||
AZURE_OPENAI_V1_HOST_SUFFIXES: Final = (".services.ai.azure.com", ".openai.azure.com")
|
||||
|
||||
|
||||
def is_foundry_model_inference_base(api_base: str) -> bool:
|
||||
|
|
@ -19,11 +20,13 @@ def is_foundry_model_inference_base(api_base: str) -> bool:
|
|||
return "/openai/deployments" not in parsed.path
|
||||
|
||||
|
||||
def api_key_header_for_base(api_base: str | None) -> AzureAIApiKeyHeader:
|
||||
def is_azure_openai_v1_host(api_base: str | None) -> bool:
|
||||
host: Final = urlparse(api_base).hostname if api_base else None
|
||||
if host and (host.endswith(".services.ai.azure.com") or host.endswith(".openai.azure.com")):
|
||||
return "api-key"
|
||||
return "Authorization"
|
||||
return host is not None and host.endswith(AZURE_OPENAI_V1_HOST_SUFFIXES)
|
||||
|
||||
|
||||
def api_key_header_for_base(api_base: str | None) -> AzureAIApiKeyHeader:
|
||||
return "api-key" if is_azure_openai_v1_host(api_base) else "Authorization"
|
||||
|
||||
|
||||
def get_azure_ai_entra_token(litellm_params: Mapping[str, object] | None = None) -> str | None:
|
||||
|
|
@ -70,6 +73,17 @@ def get_azure_ai_auth_headers(
|
|||
AZURE_MODEL_ROUTER_SELECTED_MODEL_KEY: Final = "azure_model_router_selected_model"
|
||||
|
||||
|
||||
def azure_ai_supports_native_responses(model: str | None, api_base: str | None) -> bool:
|
||||
resolved_base: Final = AzureFoundryModelInfo.get_api_base(api_base)
|
||||
if resolved_base is not None and not is_azure_openai_v1_host(resolved_base):
|
||||
return False
|
||||
if model is None:
|
||||
return True
|
||||
if "claude" in model.lower():
|
||||
return False
|
||||
return AzureFoundryModelInfo.get_azure_ai_route(model) == "default"
|
||||
|
||||
|
||||
class AzureFoundryModelInfo(BaseLLMModelInfo):
|
||||
"""Model info for Azure AI / Azure Foundry models."""
|
||||
|
||||
|
|
|
|||
0
litellm/llms/azure_ai/responses/__init__.py
Normal file
0
litellm/llms/azure_ai/responses/__init__.py
Normal file
53
litellm/llms/azure_ai/responses/transformation.py
Normal file
53
litellm/llms/azure_ai/responses/transformation.py
Normal file
|
|
@ -0,0 +1,53 @@
|
|||
from typing import Final
|
||||
|
||||
import httpx
|
||||
|
||||
from litellm.llms.azure.responses.transformation import AzureOpenAIResponsesAPIConfig
|
||||
from litellm.llms.azure_ai.common_utils import (
|
||||
AzureFoundryModelInfo,
|
||||
api_key_header_for_base,
|
||||
get_azure_ai_auth_headers,
|
||||
)
|
||||
from litellm.types.router import GenericLiteLLMParams
|
||||
from litellm.types.utils import LlmProviders
|
||||
|
||||
_PROJECT_PATH_PREFIX: Final = ("api", "projects")
|
||||
_RESPONSES_PATH: Final = ("openai", "v1", "responses")
|
||||
|
||||
|
||||
def _responses_url(api_base: str) -> str:
|
||||
base_url: Final = httpx.URL(api_base)
|
||||
segments: Final = tuple(segment for segment in base_url.path.split("/") if segment)
|
||||
project_root: Final = segments[:3] if segments[:2] == _PROJECT_PATH_PREFIX else ()
|
||||
return str(base_url.copy_with(path="/" + "/".join((*project_root, *_RESPONSES_PATH)), query=None))
|
||||
|
||||
|
||||
class AzureAIResponsesAPIConfig(AzureOpenAIResponsesAPIConfig):
|
||||
@property
|
||||
def custom_llm_provider(self) -> LlmProviders:
|
||||
return LlmProviders.AZURE_AI
|
||||
|
||||
def validate_environment(self, headers: dict, model: str, litellm_params: GenericLiteLLMParams | None) -> dict:
|
||||
params: Final = litellm_params or GenericLiteLLMParams()
|
||||
auth_headers: Final = get_azure_ai_auth_headers(
|
||||
api_key=AzureFoundryModelInfo.get_api_key(params.api_key),
|
||||
litellm_params=params.model_dump(),
|
||||
api_key_header=api_key_header_for_base(AzureFoundryModelInfo.get_api_base(params.api_base)),
|
||||
)
|
||||
return { # mutable-ok: the handler updates the returned headers in place per the dict contract
|
||||
**headers,
|
||||
**auth_headers,
|
||||
"Content-Type": "application/json",
|
||||
}
|
||||
|
||||
def supports_native_websocket(self) -> bool:
|
||||
return False
|
||||
|
||||
def get_complete_url(self, api_base: str | None, litellm_params: dict) -> str:
|
||||
resolved_base: Final = AzureFoundryModelInfo.get_api_base(api_base)
|
||||
if resolved_base is None:
|
||||
raise ValueError(
|
||||
"api_base is required for the Azure AI Foundry Responses API. "
|
||||
"Set the api_base parameter or the AZURE_AI_API_BASE environment variable."
|
||||
)
|
||||
return _responses_url(resolved_base)
|
||||
|
|
@ -53,6 +53,7 @@ from litellm.types.llms.openai import (
|
|||
AllMessageValues,
|
||||
ChatCompletionAnnotation,
|
||||
ChatCompletionAssistantMessage,
|
||||
ChatCompletionAssistantToolCall,
|
||||
ChatCompletionRedactedThinkingBlock,
|
||||
ChatCompletionResponseMessage,
|
||||
ChatCompletionSystemMessage,
|
||||
|
|
@ -205,6 +206,84 @@ class AmazonConverseConfig(BaseConfig):
|
|||
|
||||
return messages_copy
|
||||
|
||||
@staticmethod
|
||||
def _has_orphaned_tool_blocks(messages: list[AllMessageValues]) -> bool:
|
||||
return any(
|
||||
(m.get("role") == "assistant" and m.get("tool_calls")) or m.get("role") in ("tool", "function")
|
||||
for m in messages
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _neutralize_orphaned_tool_blocks(
|
||||
messages: list[AllMessageValues], optional_params: dict
|
||||
) -> list[AllMessageValues]:
|
||||
if optional_params.get("tools") or not AmazonConverseConfig._has_orphaned_tool_blocks(messages):
|
||||
return messages
|
||||
|
||||
from litellm.litellm_core_utils.prompt_templates.common_utils import (
|
||||
convert_content_list_to_str,
|
||||
)
|
||||
|
||||
def _tool_call_text(tool_call: ChatCompletionAssistantToolCall) -> str:
|
||||
function = tool_call.get("function") or {}
|
||||
name = function.get("name") or "unknown_tool"
|
||||
arguments = function.get("arguments") or ""
|
||||
call_id = tool_call.get("id")
|
||||
label = f"tool call {call_id}" if call_id else "tool call"
|
||||
return f"[{label}: {name}({arguments})]"
|
||||
|
||||
def _result_text(message: AllMessageValues) -> str:
|
||||
rendered = convert_content_list_to_str(message).strip()
|
||||
return rendered or "<non-text tool result omitted>"
|
||||
|
||||
guardrail_active: Final = "guardrailConfig" in optional_params
|
||||
|
||||
def _rewrite(message: AllMessageValues) -> AllMessageValues:
|
||||
role = message.get("role")
|
||||
tool_calls = message.get("tool_calls")
|
||||
if role == "assistant" and tool_calls:
|
||||
base_text: Final = convert_content_list_to_str(message)
|
||||
call_texts: Final = tuple(_tool_call_text(call) for call in tool_calls)
|
||||
text: Final = "\n".join(part for part in (base_text, *call_texts) if part)
|
||||
return ChatCompletionAssistantMessage(role="assistant", content=text)
|
||||
if role in ("tool", "function"):
|
||||
tool_call_id = message.get("tool_call_id")
|
||||
name = message.get("name")
|
||||
label = f"tool result for {tool_call_id or name or 'unknown'}"
|
||||
result_text: Final = f"[{label}: {_result_text(message)}]"
|
||||
# Tool results are externally controlled, so guard them wherever they
|
||||
# land in history; _convert_consecutive_user_messages_to_guarded_text
|
||||
# only covers the trailing user turn.
|
||||
content: Final = [{"type": "guarded_text", "text": result_text}] if guardrail_active else result_text
|
||||
return ChatCompletionUserMessage(role="user", content=content)
|
||||
return message
|
||||
|
||||
verbose_logger.warning(
|
||||
"litellm.bedrock: request has tool blocks in message history but no "
|
||||
"`tools=` param; neutralizing orphaned tool blocks to text so Bedrock "
|
||||
"accepts the request without a toolConfig. Non-text tool-result "
|
||||
"payloads are dropped. Pass `tools=` to preserve structured tool calling."
|
||||
)
|
||||
return [_rewrite(message) for message in messages]
|
||||
|
||||
@staticmethod
|
||||
def _handle_orphaned_tool_blocks(messages: list[AllMessageValues], optional_params: dict) -> list[AllMessageValues]:
|
||||
if litellm.bedrock_neutralize_orphaned_tool_blocks:
|
||||
return AmazonConverseConfig._neutralize_orphaned_tool_blocks(messages, optional_params)
|
||||
|
||||
if "tools" in optional_params or not has_tool_call_blocks(messages):
|
||||
return messages
|
||||
|
||||
if litellm.modify_params:
|
||||
optional_params["tools"] = add_dummy_tool(custom_llm_provider="bedrock_converse")
|
||||
return messages
|
||||
|
||||
raise litellm.utils.UnsupportedParamsError(
|
||||
message="Bedrock doesn't support tool calling without `tools=` param specified. Pass `tools=` param OR set `litellm.modify_params = True` // `litellm_settings::modify_params: True` to add dummy tool to the request.",
|
||||
model="",
|
||||
llm_provider="bedrock",
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def get_config(cls):
|
||||
return {
|
||||
|
|
@ -1609,20 +1688,6 @@ class AmazonConverseConfig(BaseConfig):
|
|||
drop_params: bool = False,
|
||||
litellm_params: Mapping[str, object] | None = None,
|
||||
) -> CommonRequestObject:
|
||||
## VALIDATE REQUEST
|
||||
"""
|
||||
Bedrock doesn't support tool calling without `tools=` param specified.
|
||||
"""
|
||||
if "tools" not in optional_params and messages is not None and has_tool_call_blocks(messages):
|
||||
if litellm.modify_params:
|
||||
optional_params["tools"] = add_dummy_tool(custom_llm_provider="bedrock_converse")
|
||||
else:
|
||||
raise litellm.UnsupportedParamsError(
|
||||
message="Bedrock doesn't support tool calling without `tools=` param specified. Pass `tools=` param OR set `litellm.modify_params = True` // `litellm_settings::modify_params: True` to add dummy tool to the request.",
|
||||
model="",
|
||||
llm_provider="bedrock",
|
||||
)
|
||||
|
||||
# Drop thinking param if thinking is enabled but thinking_blocks are missing
|
||||
# This prevents the error: "Expected thinking or redacted_thinking, but found tool_use"
|
||||
#
|
||||
|
|
@ -1735,7 +1800,9 @@ class AmazonConverseConfig(BaseConfig):
|
|||
messages, system_content_blocks = self._transform_system_message(messages, model=model)
|
||||
|
||||
# Convert last user message to guarded_text if guardrailConfig is present
|
||||
messages = self._convert_consecutive_user_messages_to_guarded_text(messages, optional_params)
|
||||
messages = self._convert_consecutive_user_messages_to_guarded_text(
|
||||
self._handle_orphaned_tool_blocks(messages, optional_params), optional_params
|
||||
)
|
||||
## TRANSFORMATION ##
|
||||
|
||||
_data: Final[CommonRequestObject] = self._transform_request_helper(
|
||||
|
|
@ -1796,7 +1863,9 @@ class AmazonConverseConfig(BaseConfig):
|
|||
messages, system_content_blocks = self._transform_system_message(messages, model=model)
|
||||
|
||||
# Convert last user message to guarded_text if guardrailConfig is present
|
||||
messages = self._convert_consecutive_user_messages_to_guarded_text(messages, optional_params)
|
||||
messages = self._convert_consecutive_user_messages_to_guarded_text(
|
||||
self._handle_orphaned_tool_blocks(messages, optional_params), optional_params
|
||||
)
|
||||
|
||||
_data: Final[CommonRequestObject] = self._transform_request_helper(
|
||||
model=model,
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -4,9 +4,9 @@ from typing import Final
|
|||
|
||||
import litellm
|
||||
from litellm.utils import (
|
||||
_is_explicitly_disabled_factory,
|
||||
_supports_factory,
|
||||
declared_value_factory,
|
||||
is_explicitly_disabled_factory,
|
||||
)
|
||||
|
||||
from .gpt_transformation import OpenAIGPTConfig
|
||||
|
|
@ -192,7 +192,7 @@ class OpenAIGPT5Config(OpenAIGPTConfig):
|
|||
|
||||
Use this for opt-out checks where unknown models should be allowed through.
|
||||
"""
|
||||
return _is_explicitly_disabled_factory(
|
||||
return is_explicitly_disabled_factory(
|
||||
model=cls._model_map_lookup_name(model),
|
||||
custom_llm_provider=None,
|
||||
key=f"supports_{level}_reasoning_effort",
|
||||
|
|
|
|||
|
|
@ -79,6 +79,7 @@ from litellm.utils import (
|
|||
CustomStreamWrapper,
|
||||
ModelResponse,
|
||||
is_base64_encoded,
|
||||
is_explicitly_disabled_factory,
|
||||
supports_reasoning,
|
||||
)
|
||||
|
||||
|
|
@ -866,6 +867,14 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig):
|
|||
else:
|
||||
raise _unsupported_reasoning_effort(reasoning_effort)
|
||||
|
||||
@staticmethod
|
||||
def _supports_minimal_thinking_level(model: str) -> bool:
|
||||
lowered: Final = model.lower()
|
||||
is_gemini3flash: Final = "gemini-3" in lowered and "flash" in lowered
|
||||
return is_gemini3flash and not is_explicitly_disabled_factory(
|
||||
model=model, custom_llm_provider=None, key="supports_minimal_reasoning_effort"
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _map_reasoning_effort_to_thinking_level(
|
||||
reasoning_effort: str,
|
||||
|
|
@ -880,13 +889,11 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig):
|
|||
Returns:
|
||||
GeminiThinkingConfig with thinkingLevel and includeThoughts
|
||||
"""
|
||||
# Check if this is gemini-3-flash which supports MINIMAL thinking level
|
||||
# Covers gemini-3-flash, gemini-3-flash-preview, gemini-3.1-flash, gemini-3.1-flash-lite-preview,
|
||||
# gemini-3.5-flash, and any future 3.x-flash variants.
|
||||
is_gemini3flash: Final = model and ("flash" in model.lower() and "gemini-3" in model.lower())
|
||||
supports_minimal: Final = bool(model) and VertexGeminiConfig._supports_minimal_thinking_level(model)
|
||||
is_gemini31pro: Final = model and ("gemini-3.1-pro-preview" in model.lower())
|
||||
if reasoning_effort == "minimal":
|
||||
if is_gemini3flash:
|
||||
if supports_minimal:
|
||||
return {"thinkingLevel": "minimal", "includeThoughts": True}
|
||||
else:
|
||||
return {"thinkingLevel": "low", "includeThoughts": True}
|
||||
|
|
@ -899,18 +906,11 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig):
|
|||
return {"thinkingLevel": "high", "includeThoughts": True}
|
||||
elif reasoning_effort == "high":
|
||||
return {"thinkingLevel": "high", "includeThoughts": True}
|
||||
elif reasoning_effort == "disable":
|
||||
# Gemini 3 cannot fully disable thinking, so we use "minimal" for gemini-3-flash-preview, "low" for others
|
||||
if is_gemini3flash:
|
||||
return {"thinkingLevel": "minimal", "includeThoughts": False}
|
||||
else:
|
||||
return {"thinkingLevel": "low", "includeThoughts": False}
|
||||
elif reasoning_effort == "none":
|
||||
# For gemini-3-flash-preview, use "minimal" instead of "low"
|
||||
if is_gemini3flash:
|
||||
return {"thinkingLevel": "minimal", "includeThoughts": False}
|
||||
else:
|
||||
return {"thinkingLevel": "low", "includeThoughts": False}
|
||||
elif reasoning_effort in ("disable", "none"):
|
||||
return {
|
||||
"thinkingLevel": "minimal" if supports_minimal else "low",
|
||||
"includeThoughts": False,
|
||||
}
|
||||
else:
|
||||
raise _unsupported_reasoning_effort(reasoning_effort)
|
||||
|
||||
|
|
@ -977,8 +977,9 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig):
|
|||
params["includeThoughts"] = True
|
||||
# Follow provider defaults unless explicitly opted into legacy behavior.
|
||||
if litellm.enable_gemini_default_thinking_level_low is True:
|
||||
is_gemini3flash: Final = "gemini-3" in model.lower() and "flash" in model.lower()
|
||||
params["thinkingLevel"] = "minimal" if is_gemini3flash else "low"
|
||||
params["thinkingLevel"] = (
|
||||
"minimal" if VertexGeminiConfig._supports_minimal_thinking_level(model) else "low"
|
||||
)
|
||||
else:
|
||||
# Thinking disabled
|
||||
params["includeThoughts"] = False
|
||||
|
|
|
|||
|
|
@ -26105,6 +26105,7 @@
|
|||
"supports_pdf_input": true,
|
||||
"supports_prompt_caching": true,
|
||||
"supports_reasoning": true,
|
||||
"supports_minimal_reasoning_effort": false,
|
||||
"supports_response_schema": true,
|
||||
"supports_system_messages": true,
|
||||
"supports_tool_choice": true,
|
||||
|
|
@ -26162,6 +26163,7 @@
|
|||
"supports_pdf_input": true,
|
||||
"supports_prompt_caching": true,
|
||||
"supports_reasoning": true,
|
||||
"supports_minimal_reasoning_effort": false,
|
||||
"supports_response_schema": true,
|
||||
"supports_system_messages": true,
|
||||
"supports_tool_choice": true,
|
||||
|
|
@ -28111,6 +28113,7 @@
|
|||
"supports_pdf_input": true,
|
||||
"supports_prompt_caching": true,
|
||||
"supports_reasoning": true,
|
||||
"supports_minimal_reasoning_effort": false,
|
||||
"supports_response_schema": true,
|
||||
"supports_system_messages": true,
|
||||
"supports_tool_choice": true,
|
||||
|
|
@ -28170,6 +28173,7 @@
|
|||
"supports_pdf_input": true,
|
||||
"supports_prompt_caching": true,
|
||||
"supports_reasoning": true,
|
||||
"supports_minimal_reasoning_effort": false,
|
||||
"supports_response_schema": true,
|
||||
"supports_system_messages": true,
|
||||
"supports_tool_choice": true,
|
||||
|
|
@ -28592,6 +28596,7 @@
|
|||
"supports_pdf_input": true,
|
||||
"supports_prompt_caching": true,
|
||||
"supports_reasoning": true,
|
||||
"supports_minimal_reasoning_effort": false,
|
||||
"supports_response_schema": true,
|
||||
"supports_system_messages": true,
|
||||
"supports_tool_choice": true,
|
||||
|
|
@ -28649,6 +28654,7 @@
|
|||
"supports_pdf_input": true,
|
||||
"supports_prompt_caching": true,
|
||||
"supports_reasoning": true,
|
||||
"supports_minimal_reasoning_effort": false,
|
||||
"supports_response_schema": true,
|
||||
"supports_system_messages": true,
|
||||
"supports_tool_choice": true,
|
||||
|
|
|
|||
|
|
@ -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():
|
||||
|
|
|
|||
|
|
@ -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():
|
||||
|
|
|
|||
|
|
@ -845,6 +845,9 @@ class LiteLLMRoutes(enum.Enum):
|
|||
)
|
||||
|
||||
self_managed_routes = [
|
||||
# update_team resolves proxy/org/team admin itself and filters team admins
|
||||
# through the team_admin_editable_team_fields setting
|
||||
"/team/update",
|
||||
"/team/member_add",
|
||||
"/team/member_delete",
|
||||
"/management/v1/teams/{team_id}/members/bulk_delete",
|
||||
|
|
@ -4468,6 +4471,29 @@ class TeamInfoMember(Member):
|
|||
user_alias: str | None = None
|
||||
|
||||
|
||||
class TeamEditUnrestricted(BaseModel):
|
||||
kind: Literal["unrestricted"] = "unrestricted"
|
||||
|
||||
|
||||
class TeamEditAsTeamAdmin(BaseModel):
|
||||
kind: Literal["team_admin"] = "team_admin"
|
||||
editable_fields: tuple[str, ...]
|
||||
|
||||
|
||||
class TeamEditAsTeamAdminDisabled(BaseModel):
|
||||
kind: Literal["team_admin_disabled"] = "team_admin_disabled"
|
||||
|
||||
|
||||
class TeamEditNone(BaseModel):
|
||||
kind: Literal["none"] = "none"
|
||||
|
||||
|
||||
TeamEditAccess = Annotated[
|
||||
TeamEditUnrestricted | TeamEditAsTeamAdmin | TeamEditAsTeamAdminDisabled | TeamEditNone,
|
||||
Field(discriminator="kind"),
|
||||
]
|
||||
|
||||
|
||||
class TeamInfoResponseObjectTeamTable(LiteLLM_TeamTable):
|
||||
members_with_roles: tuple[TeamInfoMember, ...] = ()
|
||||
team_member_budget_table: LiteLLM_BudgetTableFull | None = None
|
||||
|
|
@ -4480,6 +4506,7 @@ class TeamInfoResponseObjectTeamTable(LiteLLM_TeamTable):
|
|||
# None = no org or not a manager; [] or ["all-proxy-models"] = no ceiling.
|
||||
organization_models: list[str] | None = None
|
||||
model_max_budget_usage: Mapping[str, Mapping[str, object]] | None = None
|
||||
caller_edit_access: TeamEditAccess = Field(default_factory=TeamEditNone)
|
||||
|
||||
|
||||
class TeamInfoResponseObject(TypedDict):
|
||||
|
|
|
|||
|
|
@ -0,0 +1,191 @@
|
|||
"""Proxy-wide allow-list of team-settings fields a team admin may change on /team/update."""
|
||||
|
||||
from collections.abc import Mapping
|
||||
from dataclasses import dataclass
|
||||
from types import MappingProxyType
|
||||
from typing import Final, Literal, TypeAlias
|
||||
|
||||
from fastapi import HTTPException
|
||||
from pydantic import TypeAdapter, ValidationError
|
||||
from typing_extensions import assert_never
|
||||
|
||||
from litellm._logging import verbose_proxy_logger
|
||||
from litellm.models.team import LiteLLM_TeamTable
|
||||
from litellm.proxy._types import (
|
||||
LiteLLM_ManagementEndpoint_MetadataFields,
|
||||
LiteLLM_ManagementEndpoint_MetadataFields_Premium,
|
||||
UpdateTeamRequest,
|
||||
)
|
||||
|
||||
TEAM_ADMIN_EDITABLE_TEAM_FIELDS_SETTING: Final = "team_admin_editable_team_fields"
|
||||
|
||||
# TODO(LIT-5722): add the remaining team settings one per PR, each with its value-diff tests and dashboard field
|
||||
SUPPORTED_TEAM_ADMIN_EDITABLE_TEAM_FIELDS: Final[frozenset[str]] = frozenset({"tpm_limit"})
|
||||
|
||||
_FIELD_LIST: Final = TypeAdapter(list[str])
|
||||
_JSON_OBJECT: Final = TypeAdapter(dict[str, object])
|
||||
_EMPTY: Final[Mapping[str, object]] = MappingProxyType({})
|
||||
_METADATA_FOLDED_FIELDS: Final[frozenset[str]] = frozenset(
|
||||
(*LiteLLM_ManagementEndpoint_MetadataFields, *LiteLLM_ManagementEndpoint_MetadataFields_Premium)
|
||||
)
|
||||
_SYSTEM_MANAGED_METADATA_KEYS: Final[frozenset[str]] = frozenset({"team_member_budget_id"})
|
||||
_NOT_COLUMNS: Final[frozenset[str]] = frozenset({"team_id", "metadata"})
|
||||
_SETTINGS_LOCATION: Final = "Settings > UI > Team admin editable fields"
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class TeamAdminEditAllowed:
|
||||
request: UpdateTeamRequest
|
||||
kind: Literal["allowed"] = "allowed"
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class TeamAdminEditingDisabled:
|
||||
kind: Literal["disabled"] = "disabled"
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class TeamAdminFieldNotPermitted:
|
||||
field: str
|
||||
kind: Literal["field_not_permitted"] = "field_not_permitted"
|
||||
|
||||
|
||||
TeamAdminEditVerdict: TypeAlias = TeamAdminEditAllowed | TeamAdminEditingDisabled | TeamAdminFieldNotPermitted
|
||||
|
||||
|
||||
def resolve_team_admin_editable_fields(
|
||||
general_settings: Mapping[str, object],
|
||||
supported: frozenset[str],
|
||||
) -> frozenset[str]:
|
||||
raw: Final = general_settings.get(TEAM_ADMIN_EDITABLE_TEAM_FIELDS_SETTING)
|
||||
if raw is None:
|
||||
return frozenset()
|
||||
try:
|
||||
configured: Final = frozenset(_FIELD_LIST.validate_python(raw))
|
||||
except ValidationError:
|
||||
verbose_proxy_logger.warning(
|
||||
"%s must be a list of field names; ignoring %r", TEAM_ADMIN_EDITABLE_TEAM_FIELDS_SETTING, raw
|
||||
)
|
||||
return frozenset()
|
||||
unsupported: Final = configured - supported
|
||||
if unsupported:
|
||||
verbose_proxy_logger.warning(
|
||||
"%s ignores unsupported field(s) %s; supported: %s",
|
||||
TEAM_ADMIN_EDITABLE_TEAM_FIELDS_SETTING,
|
||||
sorted(unsupported),
|
||||
sorted(supported),
|
||||
)
|
||||
return configured & supported
|
||||
|
||||
|
||||
def _as_object(value: object) -> Mapping[str, object]:
|
||||
try:
|
||||
return _JSON_OBJECT.validate_json(value) if isinstance(value, str) else _JSON_OBJECT.validate_python(value)
|
||||
except ValidationError:
|
||||
return _EMPTY
|
||||
|
||||
|
||||
def _stored_metadata(existing: Mapping[str, object]) -> Mapping[str, object]:
|
||||
return _as_object(existing.get("metadata"))
|
||||
|
||||
|
||||
def _submitted_metadata(
|
||||
data: UpdateTeamRequest, submitted: Mapping[str, object], existing: Mapping[str, object]
|
||||
) -> Mapping[str, object]:
|
||||
"""Metadata as it would be stored: the caller's dict (or the stored one) with top-level folded fields laid over."""
|
||||
base: Final = (
|
||||
_as_object(submitted.get("metadata")) if "metadata" in data.model_fields_set else _stored_metadata(existing)
|
||||
)
|
||||
folded: Final = data.model_fields_set & _METADATA_FOLDED_FIELDS
|
||||
return MappingProxyType({key: submitted[key] if key in folded else base[key] for key in base.keys() | folded})
|
||||
|
||||
|
||||
def _metadata_changes(
|
||||
data: UpdateTeamRequest, submitted: Mapping[str, object], existing: Mapping[str, object]
|
||||
) -> frozenset[str]:
|
||||
merged: Final = _submitted_metadata(data, submitted, existing)
|
||||
stored: Final = _stored_metadata(existing)
|
||||
return frozenset(
|
||||
key if key in _METADATA_FOLDED_FIELDS else "metadata"
|
||||
for key in (merged.keys() | stored.keys()) - _SYSTEM_MANAGED_METADATA_KEYS
|
||||
if merged.get(key) != stored.get(key)
|
||||
)
|
||||
|
||||
|
||||
def _stored_model_aliases(existing_row: LiteLLM_TeamTable) -> Mapping[str, object]:
|
||||
table: Final = existing_row.litellm_model_table
|
||||
return _as_object(_JSON_OBJECT.validate_json(table.model_dump_json()).get("model_aliases")) if table else _EMPTY
|
||||
|
||||
|
||||
def _column_changed(
|
||||
field: str, submitted: Mapping[str, object], existing: Mapping[str, object], existing_row: LiteLLM_TeamTable
|
||||
) -> bool:
|
||||
if field == "model_aliases":
|
||||
return _as_object(submitted.get(field)) != _stored_model_aliases(existing_row)
|
||||
if field in LiteLLM_TeamTable.model_fields:
|
||||
return submitted.get(field) != existing.get(field)
|
||||
return True
|
||||
|
||||
|
||||
def changed_team_fields(data: UpdateTeamRequest, existing_row: LiteLLM_TeamTable) -> frozenset[str]:
|
||||
"""Logical field names whose stored value the request would change.
|
||||
|
||||
Request and stored row are compared as JSON values so both sides share one representation. Fields the
|
||||
server folds into metadata are attributed to their own name whether they arrive top-level or inside
|
||||
``metadata``; anything else in ``metadata`` is attributed to ``metadata``. Fields with no stored
|
||||
counterpart on the team row count as changed whenever they are sent.
|
||||
"""
|
||||
submitted: Final = _JSON_OBJECT.validate_json(data.model_dump_json(exclude_unset=True))
|
||||
existing: Final = _JSON_OBJECT.validate_json(existing_row.model_dump_json())
|
||||
column_fields: Final = frozenset(data.model_fields_set) - _NOT_COLUMNS - _METADATA_FOLDED_FIELDS
|
||||
column_changes: Final = frozenset(
|
||||
field for field in column_fields if _column_changed(field, submitted, existing, existing_row)
|
||||
)
|
||||
return column_changes | _metadata_changes(data, submitted, existing)
|
||||
|
||||
|
||||
def _only_changes(data: UpdateTeamRequest, changed: frozenset[str]) -> UpdateTeamRequest:
|
||||
"""The request without the values it resends unchanged, which would otherwise still trigger derived writes
|
||||
such as a resent budget_duration pushing budget_reset_at back."""
|
||||
sent: Final = frozenset(data.model_fields_set)
|
||||
via_metadata: Final = frozenset({"metadata"}) if changed - sent else frozenset()
|
||||
kept: Final = frozenset({"team_id"}) | (changed & sent) | via_metadata
|
||||
return UpdateTeamRequest.model_validate(data.model_dump(include=MappingProxyType({field: True for field in kept})))
|
||||
|
||||
|
||||
def team_admin_edit_verdict(
|
||||
data: UpdateTeamRequest,
|
||||
existing: LiteLLM_TeamTable,
|
||||
permitted: frozenset[str],
|
||||
) -> TeamAdminEditVerdict:
|
||||
if not permitted:
|
||||
return TeamAdminEditingDisabled()
|
||||
changed: Final = changed_team_fields(data, existing)
|
||||
blocked: Final = sorted(changed - permitted)
|
||||
if blocked:
|
||||
return TeamAdminFieldNotPermitted(field=blocked[0])
|
||||
return TeamAdminEditAllowed(request=_only_changes(data, changed))
|
||||
|
||||
|
||||
def team_admin_request_or_raise(verdict: TeamAdminEditVerdict) -> UpdateTeamRequest:
|
||||
match verdict:
|
||||
case TeamAdminEditAllowed(request=request):
|
||||
return request
|
||||
case TeamAdminEditingDisabled():
|
||||
raise HTTPException(
|
||||
status_code=403,
|
||||
detail=(
|
||||
"Team admins on this proxy cannot edit team settings. "
|
||||
f"Ask a proxy admin to enable fields under {_SETTINGS_LOCATION}."
|
||||
),
|
||||
)
|
||||
case TeamAdminFieldNotPermitted(field=field):
|
||||
raise HTTPException(
|
||||
status_code=403,
|
||||
detail=(
|
||||
f"Team admins on this proxy do not have permission to update '{field}'. "
|
||||
f"Ask a proxy admin to add it under {_SETTINGS_LOCATION}."
|
||||
),
|
||||
)
|
||||
case _:
|
||||
assert_never(verdict)
|
||||
|
|
@ -18,12 +18,23 @@ from collections.abc import Iterable, Mapping, Sequence
|
|||
from collections.abc import Set as AbstractSet
|
||||
from datetime import datetime, timezone
|
||||
from types import MappingProxyType
|
||||
from typing import TYPE_CHECKING, Annotated, Final, NamedTuple, NoReturn, Protocol, TypeVar, cast
|
||||
from typing import (
|
||||
TYPE_CHECKING,
|
||||
Annotated,
|
||||
Final,
|
||||
Literal,
|
||||
NamedTuple,
|
||||
NoReturn,
|
||||
Protocol,
|
||||
TypeAlias,
|
||||
TypeVar,
|
||||
cast,
|
||||
)
|
||||
|
||||
import fastapi
|
||||
from fastapi import APIRouter, Depends, Header, HTTPException, Request, status
|
||||
from pydantic import BaseModel, JsonValue, ValidationError
|
||||
from typing_extensions import ReadOnly, TypedDict
|
||||
from pydantic import BaseModel, JsonValue, TypeAdapter, ValidationError
|
||||
from typing_extensions import ReadOnly, TypedDict, assert_never
|
||||
|
||||
import litellm
|
||||
from litellm._logging import verbose_proxy_logger
|
||||
|
|
@ -62,6 +73,11 @@ from litellm.proxy._types import (
|
|||
SpecialProxyStrings,
|
||||
TeamAccessGroupModelGrant,
|
||||
TeamAddMemberResponse,
|
||||
TeamEditAccess,
|
||||
TeamEditAsTeamAdmin,
|
||||
TeamEditAsTeamAdminDisabled,
|
||||
TeamEditNone,
|
||||
TeamEditUnrestricted,
|
||||
TeamInfoMember,
|
||||
TeamInfoResponseObject,
|
||||
TeamInfoResponseObjectTeamTable,
|
||||
|
|
@ -122,6 +138,12 @@ from litellm.proxy.management_endpoints.router_weights import validate_router_se
|
|||
from litellm.proxy.management_endpoints.tag_management_endpoints import (
|
||||
get_daily_activity,
|
||||
)
|
||||
from litellm.proxy.management_endpoints.team_admin_field_permissions import (
|
||||
SUPPORTED_TEAM_ADMIN_EDITABLE_TEAM_FIELDS,
|
||||
resolve_team_admin_editable_fields,
|
||||
team_admin_edit_verdict,
|
||||
team_admin_request_or_raise,
|
||||
)
|
||||
from litellm.proxy.management_helpers.access_group_team_sync import (
|
||||
TEAM_ADVISORY_LOCK_SQL,
|
||||
AccessGroupSyncTx,
|
||||
|
|
@ -439,32 +461,70 @@ async def _refresh_cached_team(
|
|||
)
|
||||
|
||||
|
||||
async def _can_manage_team(
|
||||
TeamAccessRole: TypeAlias = Literal["proxy_admin", "org_admin", "team_admin"]
|
||||
|
||||
|
||||
def _raise_team_access_denied() -> NoReturn:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_403_FORBIDDEN,
|
||||
detail="You do not have access to this team",
|
||||
)
|
||||
|
||||
|
||||
async def _resolve_team_access(
|
||||
team_obj: LiteLLM_TeamTable,
|
||||
user_api_key_dict: UserAPIKeyAuth,
|
||||
) -> bool:
|
||||
"""True for a proxy admin, an admin of this team, or an org admin for the team's organization."""
|
||||
) -> TeamAccessRole | None:
|
||||
"""Strongest role the caller holds over ``team_obj``, or None when they hold none.
|
||||
|
||||
Org admin outranks team admin so a caller holding both keeps unrestricted edits.
|
||||
"""
|
||||
if user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN:
|
||||
return True
|
||||
return "proxy_admin"
|
||||
|
||||
if await _is_user_org_admin_for_team(user_api_key_dict=user_api_key_dict, team_obj=team_obj):
|
||||
return "org_admin"
|
||||
|
||||
if _is_user_team_admin(user_api_key_dict=user_api_key_dict, team_obj=team_obj):
|
||||
return True
|
||||
return "team_admin"
|
||||
|
||||
return await _is_user_org_admin_for_team(user_api_key_dict=user_api_key_dict, team_obj=team_obj)
|
||||
return None
|
||||
|
||||
|
||||
async def _verify_team_access(
|
||||
team_obj: LiteLLM_TeamTable,
|
||||
user_api_key_dict: UserAPIKeyAuth,
|
||||
) -> None:
|
||||
"""Raise HTTPException(403) unless the caller can manage the given team."""
|
||||
if await _can_manage_team(team_obj=team_obj, user_api_key_dict=user_api_key_dict):
|
||||
return
|
||||
"""Raise 403 unless the caller is a proxy admin, an org admin for the team's org, or a team admin."""
|
||||
if await _resolve_team_access(team_obj=team_obj, user_api_key_dict=user_api_key_dict) is None:
|
||||
_raise_team_access_denied()
|
||||
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_403_FORBIDDEN,
|
||||
detail="You do not have access to this team",
|
||||
)
|
||||
|
||||
_GENERAL_SETTINGS: Final = TypeAdapter(dict[str, object])
|
||||
|
||||
|
||||
def _general_settings() -> Mapping[str, object]:
|
||||
from litellm.proxy.proxy_server import general_settings
|
||||
|
||||
return _GENERAL_SETTINGS.validate_python(general_settings)
|
||||
|
||||
|
||||
def _caller_edit_access(role: TeamAccessRole | None, general_settings: Mapping[str, object]) -> TeamEditAccess:
|
||||
"""What the caller may change on /team/update, reported on /team/info so the dashboard never re-derives it."""
|
||||
match role:
|
||||
case "proxy_admin" | "org_admin":
|
||||
return TeamEditUnrestricted()
|
||||
case "team_admin":
|
||||
permitted: Final = resolve_team_admin_editable_fields(
|
||||
general_settings, SUPPORTED_TEAM_ADMIN_EDITABLE_TEAM_FIELDS
|
||||
)
|
||||
if not permitted:
|
||||
return TeamEditAsTeamAdminDisabled()
|
||||
return TeamEditAsTeamAdmin(editable_fields=tuple(sorted(permitted)))
|
||||
case None:
|
||||
return TeamEditNone()
|
||||
case _:
|
||||
assert_never(role)
|
||||
|
||||
|
||||
class TeamMemberBudgetHandler:
|
||||
|
|
@ -2144,16 +2204,29 @@ async def update_team(
|
|||
)
|
||||
|
||||
if existing_team_row is None:
|
||||
# Non-proxy-admins get the same 403 as an access denial so /team/update
|
||||
# cannot be used to probe which team ids exist
|
||||
if user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN:
|
||||
_raise_team_access_denied()
|
||||
raise HTTPException(
|
||||
status_code=404,
|
||||
detail={"error": f"Team not found, passed team_id={data.team_id}"},
|
||||
)
|
||||
|
||||
# Verify caller has access to manage this team
|
||||
await _verify_team_access(
|
||||
team_obj=LiteLLM_TeamTable.model_validate(existing_team_row.model_dump()),
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
)
|
||||
existing_team: Final = LiteLLM_TeamTable.model_validate(existing_team_row.model_dump())
|
||||
access_role: Final = await _resolve_team_access(team_obj=existing_team, user_api_key_dict=user_api_key_dict)
|
||||
if access_role is None:
|
||||
_raise_team_access_denied()
|
||||
if access_role == "team_admin":
|
||||
data = team_admin_request_or_raise( # rebind-ok: resent values must not reach the derived writes below
|
||||
team_admin_edit_verdict(
|
||||
data=data,
|
||||
existing=existing_team,
|
||||
permitted=resolve_team_admin_editable_fields(
|
||||
_general_settings(), SUPPORTED_TEAM_ADMIN_EDITABLE_TEAM_FIELDS
|
||||
),
|
||||
)
|
||||
)
|
||||
|
||||
await validate_router_settings_weights(
|
||||
data.router_settings,
|
||||
|
|
@ -2257,6 +2330,7 @@ async def update_team(
|
|||
org_id=org_id_to_check,
|
||||
user_api_key_cache=user_api_key_cache,
|
||||
prisma_client=prisma_client,
|
||||
include_budget_table=True,
|
||||
)
|
||||
if org_table is not None:
|
||||
await _check_org_team_limits(
|
||||
|
|
@ -4583,10 +4657,9 @@ async def team_info(
|
|||
)
|
||||
team_table: Final = LiteLLM_TeamTable.model_validate(team_info.model_dump())
|
||||
await validate_membership(user_api_key_dict=user_api_key_dict, team_table=team_table)
|
||||
access_role: Final = await _resolve_team_access(team_obj=team_table, user_api_key_dict=user_api_key_dict)
|
||||
organization_models: Final[list[str] | None] = (
|
||||
_parent_organization_models(team_info)
|
||||
if await _can_manage_team(team_obj=team_table, user_api_key_dict=user_api_key_dict)
|
||||
else None
|
||||
_parent_organization_models(team_info) if access_role is not None else None
|
||||
)
|
||||
|
||||
## GET ALL KEYS ##
|
||||
|
|
@ -4655,6 +4728,7 @@ async def team_info(
|
|||
model_max_budget=resolved_team_info.model_max_budget,
|
||||
cache=model_max_budget_limiter.dual_cache,
|
||||
),
|
||||
"caller_edit_access": _caller_edit_access(access_role, _general_settings()),
|
||||
}
|
||||
)
|
||||
|
||||
|
|
|
|||
|
|
@ -1180,9 +1180,8 @@ async def bedrock_proxy_route(
|
|||
endpoint_func: Final = create_pass_through_route(
|
||||
endpoint=endpoint,
|
||||
target=str(prepped.url),
|
||||
custom_headers=prepped.headers,
|
||||
custom_headers=_upstream_headers_for_bedrock_agent_runtime_route(request, user_api_key_dict, prepped.headers),
|
||||
is_streaming_request=is_streaming_request,
|
||||
_forward_headers=True,
|
||||
) # dynamically construct pass-through endpoint based on incoming path
|
||||
setattr(request.state, LITELLM_PASS_THROUGH_CUSTOM_BODY_STATE_KEY, data)
|
||||
# SigV4 signs an exact payload; pass-through must send prepped.body, not json.dumps
|
||||
|
|
@ -2116,6 +2115,9 @@ _ANTHROPIC_UPSTREAM_CREDENTIAL_HEADERS: Final = frozenset({"authorization", "x-a
|
|||
_HEADERS_NEVER_FORWARDED_TO_ANTHROPIC: Final = frozenset({"content-length", "host", "accept-encoding"}) | (
|
||||
SpecialHeaders.litellm_credential_header_names() - _ANTHROPIC_UPSTREAM_CREDENTIAL_HEADERS
|
||||
)
|
||||
_HEADERS_NEVER_FORWARDED_TO_BEDROCK: Final = (
|
||||
frozenset({"content-length", "host", "accept-encoding"}) | SpecialHeaders.litellm_credential_header_names()
|
||||
)
|
||||
|
||||
|
||||
_MAPPED_ROUTE_CALLER_KEY_HEADER: Final = "litellm_user_api_key"
|
||||
|
|
@ -2214,6 +2216,17 @@ def _upstream_headers_for_anthropic_route(
|
|||
return MappingProxyType({**caller_headers, **(proxy_auth_header or {})})
|
||||
|
||||
|
||||
def _upstream_headers_for_bedrock_agent_runtime_route(
|
||||
request: Request, user_api_key_dict: UserAPIKeyAuth, signed_headers: Mapping[str, object]
|
||||
) -> Mapping[str, object]:
|
||||
caller_headers: Final = _caller_headers_without_litellm_secrets(
|
||||
request,
|
||||
user_api_key_dict,
|
||||
_HEADERS_NEVER_FORWARDED_TO_BEDROCK | frozenset(name.lower() for name in signed_headers),
|
||||
)
|
||||
return MappingProxyType({**caller_headers, **signed_headers})
|
||||
|
||||
|
||||
async def _prepare_vertex_auth_headers(
|
||||
request: Request,
|
||||
vertex_credentials: VertexPassThroughCredentials | None,
|
||||
|
|
|
|||
|
|
@ -689,6 +689,9 @@ from litellm.proxy.types_utils.utils import get_instance_fn
|
|||
from litellm.proxy.ui_crud_endpoints.proxy_setting_endpoints import (
|
||||
router as ui_crud_endpoints_router,
|
||||
)
|
||||
from litellm.proxy.ui_crud_endpoints.proxy_setting_endpoints import (
|
||||
sync_ui_settings_to_general_settings,
|
||||
)
|
||||
from litellm.proxy.ui_crud_endpoints.user_banner_endpoints import (
|
||||
router as user_banner_endpoints_router,
|
||||
)
|
||||
|
|
@ -1753,10 +1756,6 @@ class _SSOConfigRow(Protocol):
|
|||
sso_settings: MutableMapping[str, object]
|
||||
|
||||
|
||||
class _UISettingsRow(Protocol):
|
||||
ui_settings: Mapping[str, object] | str | None
|
||||
|
||||
|
||||
class _InvitationLinkRow(Protocol):
|
||||
user_id: str
|
||||
expires_at: datetime
|
||||
|
|
@ -7412,7 +7411,12 @@ class ProxyConfig:
|
|||
Returns what the reconcile saw, captured before the lock is released so a
|
||||
caller's verdict cannot be corrupted by the next reconcile's own in-flight
|
||||
window. See ReconcileOutcome.
|
||||
|
||||
Also re-reads the UI settings that back runtime flags. That runs before the lock, so a
|
||||
setting written through one pod reaches the others without waiting on a model reconcile.
|
||||
"""
|
||||
await sync_ui_settings_to_general_settings(prisma_client)
|
||||
|
||||
async with MODEL_RECONCILE_LOCK:
|
||||
return await self._add_deployment_locked(prisma_client=prisma_client, proxy_logging_obj=proxy_logging_obj)
|
||||
|
||||
|
|
@ -9648,35 +9652,12 @@ class ProxyStartupEvent:
|
|||
|
||||
@classmethod
|
||||
async def _sync_ui_settings_to_general_settings(cls):
|
||||
"""
|
||||
Load persisted UI settings from the database and sync runtime flags
|
||||
into general_settings so they take effect immediately after startup.
|
||||
"""
|
||||
try:
|
||||
import json
|
||||
|
||||
from litellm.proxy.ui_crud_endpoints.proxy_setting_endpoints import (
|
||||
_RUNTIME_GENERAL_SETTINGS_FLAGS,
|
||||
)
|
||||
|
||||
if prisma_client is None:
|
||||
return
|
||||
db_record: Final[_UISettingsRow | None] = cast( # cast-ok: prisma Json stub is `str`, runtime is a dict
|
||||
"_UISettingsRow | None",
|
||||
await UISettingsRepository(prisma_client).table.find_unique(where={"id": "ui_settings"}),
|
||||
)
|
||||
if db_record and db_record.ui_settings:
|
||||
raw: Final = db_record.ui_settings
|
||||
ui_settings: Final = json.loads(raw) if isinstance(raw, str) else dict(raw)
|
||||
flags_to_sync: Final = {k: ui_settings[k] for k in _RUNTIME_GENERAL_SETTINGS_FLAGS if k in ui_settings}
|
||||
if flags_to_sync:
|
||||
general_settings.update(flags_to_sync)
|
||||
verbose_proxy_logger.info(
|
||||
"Synced UI settings to general_settings on startup: %s",
|
||||
list(flags_to_sync.keys()),
|
||||
)
|
||||
except Exception as e:
|
||||
verbose_proxy_logger.debug("UI settings sync on startup skipped or failed: %s", e)
|
||||
"""Apply the persisted UI settings to general_settings before this pod serves traffic."""
|
||||
if prisma_client is None:
|
||||
return
|
||||
applied: Final = await sync_ui_settings_to_general_settings(prisma_client)
|
||||
if applied:
|
||||
verbose_proxy_logger.info("Synced UI settings to general_settings on startup: %s", list(applied))
|
||||
|
||||
@classmethod
|
||||
async def _load_heuristic_v1_tuning_baselines(
|
||||
|
|
@ -12879,7 +12860,6 @@ from litellm.repositories.table_repositories import (
|
|||
InvitationLinkRepository,
|
||||
PromptRepository,
|
||||
SSOConfigRepository,
|
||||
UISettingsRepository,
|
||||
)
|
||||
from litellm.repositories.team_repository import TeamRepository
|
||||
from litellm.repositories.user_repository import UserRepository
|
||||
|
|
@ -15557,18 +15537,34 @@ async def model_group_info(
|
|||
from litellm.proxy.utils import get_available_models_for_user
|
||||
|
||||
# Get available models for the user
|
||||
all_models_str: Final = await get_available_models_for_user(
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
llm_router=llm_router,
|
||||
general_settings=general_settings,
|
||||
user_model=user_model,
|
||||
prisma_client=prisma_client,
|
||||
proxy_logging_obj=proxy_logging_obj,
|
||||
team_id=None,
|
||||
include_model_access_groups=False,
|
||||
only_model_access_groups=False,
|
||||
return_wildcard_routes=False,
|
||||
user_api_key_cache=user_api_key_cache,
|
||||
is_proxy_admin: Final = user_api_key_dict.user_role in (
|
||||
LitellmUserRoles.PROXY_ADMIN,
|
||||
LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY,
|
||||
)
|
||||
all_models_str: Final = (
|
||||
get_complete_model_list(
|
||||
key_models=(),
|
||||
team_models=(),
|
||||
proxy_model_list=llm_router.get_model_names(),
|
||||
user_model=user_model,
|
||||
infer_model_from_keys=general_settings.get("infer_model_from_keys", False),
|
||||
return_wildcard_routes=False,
|
||||
llm_router=llm_router,
|
||||
)
|
||||
if is_proxy_admin
|
||||
else await get_available_models_for_user(
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
llm_router=llm_router,
|
||||
general_settings=general_settings,
|
||||
user_model=user_model,
|
||||
prisma_client=prisma_client,
|
||||
proxy_logging_obj=proxy_logging_obj,
|
||||
team_id=None,
|
||||
include_model_access_groups=False,
|
||||
only_model_access_groups=False,
|
||||
return_wildcard_routes=False,
|
||||
user_api_key_cache=user_api_key_cache,
|
||||
)
|
||||
)
|
||||
model_groups: list[ModelGroupInfoProxy] = _get_model_group_info(
|
||||
llm_router=llm_router, all_models_str=all_models_str, model_group=model_group
|
||||
|
|
|
|||
|
|
@ -14,7 +14,7 @@ from typing import (
|
|||
from urllib.parse import urlparse
|
||||
|
||||
from fastapi import APIRouter, Body, Depends, File, HTTPException, UploadFile
|
||||
from pydantic import ConfigDict, JsonValue, ValidationError, create_model
|
||||
from pydantic import ConfigDict, JsonValue, TypeAdapter, ValidationError, create_model
|
||||
from pydantic.fields import FieldInfo
|
||||
from typing_extensions import NotRequired, ReadOnly, TypedDict
|
||||
|
||||
|
|
@ -29,6 +29,10 @@ from litellm.proxy.config_resolvers.sso import (
|
|||
SSO_SECRET_FIELDS,
|
||||
resolve_sso_config,
|
||||
)
|
||||
from litellm.proxy.management_endpoints.team_admin_field_permissions import (
|
||||
SUPPORTED_TEAM_ADMIN_EDITABLE_TEAM_FIELDS,
|
||||
TEAM_ADMIN_EDITABLE_TEAM_FIELDS_SETTING,
|
||||
)
|
||||
from litellm.proxy.spend_tracking.ptu_feature_flag import is_ptu_cost_attribution_enabled
|
||||
from litellm.proxy.utils import invalidate_config_param
|
||||
from litellm.repositories.config_repository import ConfigRepository
|
||||
|
|
@ -212,6 +216,9 @@ class UIThemeSettingsResponse(SettingsResponse):
|
|||
"""Response model for UI theme settings"""
|
||||
|
||||
|
||||
_TEAM_ADMIN_FIELD_ENUM: Final = tuple(sorted(SUPPORTED_TEAM_ADMIN_EDITABLE_TEAM_FIELDS))
|
||||
|
||||
|
||||
class UISettings(BaseModel):
|
||||
"""Configuration for UI-specific flags"""
|
||||
|
||||
|
|
@ -304,6 +311,18 @@ class UISettings(BaseModel):
|
|||
description="If true, shows the Chat page in the UI sidebar, letting users chat with an LLM and connect their own MCP server credentials via OAuth.",
|
||||
)
|
||||
|
||||
team_admin_editable_team_fields: Sequence[str] = Field(
|
||||
default=(),
|
||||
description=(
|
||||
"Team settings fields a team admin may change on the teams they administer. "
|
||||
"Empty means team admins cannot edit team settings at all. "
|
||||
"Proxy admins and org admins are not affected."
|
||||
),
|
||||
json_schema_extra={ # mutable-ok: pydantic only merges json_schema_extra when it is a plain dict
|
||||
"items": {"type": "string", "enum": [*_TEAM_ADMIN_FIELD_ENUM]}, # mutable-ok: nested in the dict above
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
class UISettingsResponse(SettingsResponse):
|
||||
"""Response model for UI settings"""
|
||||
|
|
@ -326,6 +345,7 @@ ALLOWED_UI_SETTINGS_FIELDS: Final = {
|
|||
"disable_custom_api_keys",
|
||||
"disable_key_generate_for_org_admin",
|
||||
"enable_chat_ui",
|
||||
TEAM_ADMIN_EDITABLE_TEAM_FIELDS_SETTING,
|
||||
}
|
||||
|
||||
ENABLE_PTU_COST_ATTRIBUTION_UI_SETTING: Final = "enable_ptu_cost_attribution"
|
||||
|
|
@ -360,6 +380,7 @@ _RUNTIME_GENERAL_SETTINGS_FLAGS: Final = [
|
|||
"disable_vector_stores_for_internal_users",
|
||||
"allow_vector_stores_for_team_admins",
|
||||
"disable_key_generate_for_org_admin",
|
||||
TEAM_ADMIN_EDITABLE_TEAM_FIELDS_SETTING,
|
||||
]
|
||||
|
||||
# Extension point: packages outside OSS (e.g. litellm_enterprise) can
|
||||
|
|
@ -1457,6 +1478,42 @@ async def get_ui_settings_cached() -> dict[str, JsonValue]:
|
|||
return ui_settings
|
||||
|
||||
|
||||
_UI_SETTINGS_OBJECT: Final = TypeAdapter(dict[str, JsonValue])
|
||||
|
||||
|
||||
def apply_runtime_general_settings_flags(ui_settings: Mapping[str, JsonValue]) -> Mapping[str, JsonValue]:
|
||||
"""Copy the UI settings that gate runtime behavior into ``general_settings``. Returns what was applied."""
|
||||
from litellm.proxy.proxy_server import general_settings
|
||||
|
||||
flags: Final = {k: ui_settings[k] for k in _RUNTIME_GENERAL_SETTINGS_FLAGS if k in ui_settings}
|
||||
if flags:
|
||||
general_settings.update(flags)
|
||||
return MappingProxyType(flags)
|
||||
|
||||
|
||||
async def sync_ui_settings_to_general_settings(prisma_client: object) -> Mapping[str, JsonValue]:
|
||||
"""Re-read the persisted UI settings and apply the runtime flags to ``general_settings``.
|
||||
|
||||
Runs on startup and on every periodic config reload: the PATCH handler only updates the pod
|
||||
that served it, so every other pod needs its own read to pick up a change without a restart.
|
||||
Never raises. A read that fails leaves this pod on the flags it already had.
|
||||
"""
|
||||
try:
|
||||
db_record: Final = await _ui_settings_db(UISettingsRepository(prisma_client)).find_unique(
|
||||
where={"id": "ui_settings"}
|
||||
)
|
||||
stored: Final = (db_record.ui_settings if db_record else None) or "{}"
|
||||
parsed: Final = (
|
||||
_UI_SETTINGS_OBJECT.validate_json(stored)
|
||||
if isinstance(stored, str)
|
||||
else _UI_SETTINGS_OBJECT.validate_python(stored)
|
||||
)
|
||||
except Exception as e:
|
||||
verbose_proxy_logger.warning("Could not refresh UI settings from the database: %s", e)
|
||||
return MappingProxyType({})
|
||||
return apply_runtime_general_settings_flags(parsed)
|
||||
|
||||
|
||||
@router.get(
|
||||
"/get/ui_settings",
|
||||
tags=["UI Settings"],
|
||||
|
|
@ -1485,13 +1542,7 @@ async def get_ui_settings():
|
|||
# Sanitize any unexpected keys from persisted config before returning
|
||||
ui_settings: Final = {k: v for k, v in parsed.items() if k in ALLOWED_UI_SETTINGS_FIELDS}
|
||||
|
||||
# Sync runtime flags into general_settings so the proxy picks them up
|
||||
# at runtime (covers server restart scenarios).
|
||||
_flags_to_sync: Final = {k: ui_settings[k] for k in _RUNTIME_GENERAL_SETTINGS_FLAGS if k in ui_settings}
|
||||
if _flags_to_sync:
|
||||
from litellm.proxy.proxy_server import general_settings
|
||||
|
||||
general_settings.update(_flags_to_sync)
|
||||
apply_runtime_general_settings_flags(ui_settings)
|
||||
|
||||
# Refresh DualCache so other code paths (e.g. /user/filter/ui) see fresh values
|
||||
from litellm.proxy.proxy_server import user_api_key_cache
|
||||
|
|
@ -1571,6 +1622,20 @@ async def update_ui_settings(
|
|||
except ValidationError as e:
|
||||
raise HTTPException(status_code=422, detail=e.errors())
|
||||
|
||||
unsupported_team_fields: Final = sorted(
|
||||
frozenset(settings.team_admin_editable_team_fields) - SUPPORTED_TEAM_ADMIN_EDITABLE_TEAM_FIELDS
|
||||
)
|
||||
if unsupported_team_fields:
|
||||
raise HTTPException(
|
||||
status_code=400,
|
||||
detail={ # mutable-ok: HTTPException detail must be a plain dict for FastAPI JSON serialization
|
||||
"error": (
|
||||
f"{TEAM_ADMIN_EDITABLE_TEAM_FIELDS_SETTING} does not support {unsupported_team_fields}. "
|
||||
f"Supported fields: {sorted(SUPPORTED_TEAM_ADMIN_EDITABLE_TEAM_FIELDS)}."
|
||||
)
|
||||
},
|
||||
)
|
||||
|
||||
# Only include fields the caller actually sent (not Pydantic defaults).
|
||||
settings_dict: Final[Mapping[str, JsonValue]] = settings.model_dump(exclude_unset=True)
|
||||
|
||||
|
|
@ -1616,13 +1681,7 @@ async def update_ui_settings(
|
|||
},
|
||||
)
|
||||
|
||||
# Sync runtime flags to general_settings so the proxy picks them up
|
||||
# at runtime (general_settings is checked in pre-call utils).
|
||||
_flags_to_sync: Final = {k: ui_settings[k] for k in _RUNTIME_GENERAL_SETTINGS_FLAGS if k in ui_settings}
|
||||
if _flags_to_sync:
|
||||
from litellm.proxy.proxy_server import general_settings
|
||||
|
||||
general_settings.update(_flags_to_sync)
|
||||
apply_runtime_general_settings_flags(ui_settings)
|
||||
|
||||
# Invalidate + set DualCache so subsequent reads see the new values immediately
|
||||
from litellm.proxy.proxy_server import user_api_key_cache
|
||||
|
|
|
|||
|
|
@ -482,18 +482,27 @@ class _AsyncPromptManagementOutcome:
|
|||
|
||||
|
||||
def _resolve_responses_api_provider_config(
|
||||
model: str, custom_llm_provider: str, model_info: object
|
||||
model: str, custom_llm_provider: str, model_info: object, api_base: str | None
|
||||
) -> BaseResponsesAPIConfig | None:
|
||||
provider_config: Final = ProviderConfigManager.get_provider_responses_api_config(
|
||||
model=model, provider=custom_llm_provider
|
||||
model=model, provider=custom_llm_provider, api_base=api_base
|
||||
)
|
||||
if provider_config is not None or not _deployment_passes_through_responses(model_info):
|
||||
return provider_config
|
||||
return OpenAILikeResponsesConfig()
|
||||
|
||||
|
||||
def _api_base_kwarg(kwargs: Mapping[str, object]) -> str | None:
|
||||
api_base: Final = kwargs.get("api_base")
|
||||
return api_base if isinstance(api_base, str) else None
|
||||
|
||||
|
||||
def _will_bridge_to_chat_completions(
|
||||
model: str, custom_llm_provider: str | None, use_chat_completions_api: bool, model_info: object
|
||||
model: str,
|
||||
custom_llm_provider: str | None,
|
||||
use_chat_completions_api: bool,
|
||||
model_info: object,
|
||||
api_base: str | None,
|
||||
) -> bool:
|
||||
"""``_bridges_to_chat_completions`` for callers running before the provider config is resolved.
|
||||
|
||||
|
|
@ -507,7 +516,7 @@ def _will_bridge_to_chat_completions(
|
|||
if custom_llm_provider is None:
|
||||
return True
|
||||
return _bridges_to_chat_completions(
|
||||
_resolve_responses_api_provider_config(normalized_model[0], custom_llm_provider, model_info),
|
||||
_resolve_responses_api_provider_config(normalized_model[0], custom_llm_provider, model_info, api_base),
|
||||
use_chat_completions_api or normalized_model[1],
|
||||
)
|
||||
|
||||
|
|
@ -618,6 +627,7 @@ async def aresponses(
|
|||
custom_llm_provider,
|
||||
bool(kwargs.get("use_chat_completions_api")),
|
||||
kwargs.get("model_info"),
|
||||
_api_base_kwarg(kwargs),
|
||||
),
|
||||
):
|
||||
(
|
||||
|
|
@ -783,7 +793,11 @@ def _apply_prompt_management_to_responses_call(
|
|||
with _prompt_management_sees_a_provisional_message_list(
|
||||
kwargs,
|
||||
bridged=_will_bridge_to_chat_completions(
|
||||
model, custom_llm_provider, use_chat_completions_api, kwargs.get("model_info")
|
||||
model,
|
||||
custom_llm_provider,
|
||||
use_chat_completions_api,
|
||||
kwargs.get("model_info"),
|
||||
_api_base_kwarg(kwargs),
|
||||
),
|
||||
):
|
||||
(
|
||||
|
|
@ -1237,7 +1251,7 @@ def responses(
|
|||
responses_api_provider_config = None
|
||||
else:
|
||||
responses_api_provider_config = _resolve_responses_api_provider_config(
|
||||
model, custom_llm_provider, deployment_model_info
|
||||
model, custom_llm_provider, deployment_model_info, litellm_params.api_base
|
||||
)
|
||||
|
||||
if (
|
||||
|
|
@ -1496,6 +1510,7 @@ def delete_responses(
|
|||
ProviderConfigManager.get_provider_responses_api_config(
|
||||
model=None,
|
||||
provider=custom_llm_provider,
|
||||
api_base=litellm_params.api_base,
|
||||
)
|
||||
)
|
||||
|
||||
|
|
@ -1667,6 +1682,7 @@ def get_responses(
|
|||
ProviderConfigManager.get_provider_responses_api_config(
|
||||
model=None,
|
||||
provider=custom_llm_provider,
|
||||
api_base=litellm_params.api_base,
|
||||
)
|
||||
)
|
||||
|
||||
|
|
@ -1811,6 +1827,7 @@ def list_input_items(
|
|||
ProviderConfigManager.get_provider_responses_api_config(
|
||||
model=None,
|
||||
provider=custom_llm_provider,
|
||||
api_base=litellm_params.api_base,
|
||||
)
|
||||
)
|
||||
|
||||
|
|
@ -1960,6 +1977,7 @@ def cancel_responses(
|
|||
ProviderConfigManager.get_provider_responses_api_config(
|
||||
model=None,
|
||||
provider=custom_llm_provider,
|
||||
api_base=litellm_params.api_base,
|
||||
)
|
||||
)
|
||||
|
||||
|
|
@ -2132,6 +2150,7 @@ def compact_responses(
|
|||
ProviderConfigManager.get_provider_responses_api_config(
|
||||
model=model,
|
||||
provider=custom_llm_provider,
|
||||
api_base=litellm_params.api_base,
|
||||
)
|
||||
)
|
||||
|
||||
|
|
@ -2270,14 +2289,15 @@ async def _aresponses_websocket(
|
|||
custom_llm_provider=_custom_llm_provider,
|
||||
)
|
||||
|
||||
resolved_api_base: Final = dynamic_api_base or litellm_params.api_base or litellm.api_base or None
|
||||
responses_api_provider_config: BaseResponsesAPIConfig | None = None
|
||||
if _custom_llm_provider is not None:
|
||||
responses_api_provider_config = ProviderConfigManager.get_provider_responses_api_config(
|
||||
model=resolved_model,
|
||||
provider=litellm.LlmProviders(_custom_llm_provider),
|
||||
api_base=resolved_api_base,
|
||||
)
|
||||
|
||||
resolved_api_base: Final = dynamic_api_base or litellm_params.api_base or litellm.api_base or None
|
||||
resolved_api_key: Final = (
|
||||
dynamic_api_key
|
||||
or litellm_params.api_key
|
||||
|
|
|
|||
|
|
@ -2689,7 +2689,7 @@ def declared_value_factory(model: str, custom_llm_provider: str | None, key: str
|
|||
"""Return a string value the model map declares for *key*, or ``None`` when it says nothing.
|
||||
|
||||
The string-valued sibling of :func:`_supports_factory` and
|
||||
:func:`_is_explicitly_disabled_factory`, public where those two are not because it is read
|
||||
:func:`is_explicitly_disabled_factory`, public like the latter because both are read
|
||||
from the provider configs rather than from this module, sharing their
|
||||
``get_llm_provider`` -> ``_get_model_info_helper`` chain and their unprefixed-twin
|
||||
fallback (#20885), so a provider-prefixed entry that omits the key still answers
|
||||
|
|
@ -2725,7 +2725,7 @@ def declared_value_factory(model: str, custom_llm_provider: str | None, key: str
|
|||
return None
|
||||
|
||||
|
||||
def _is_explicitly_disabled_factory(model: str, custom_llm_provider: str | None, key: str) -> bool:
|
||||
def is_explicitly_disabled_factory(model: str, custom_llm_provider: str | None, key: str) -> bool:
|
||||
"""Return True only when the model map explicitly sets *key* to ``False``.
|
||||
|
||||
This is the opt-out mirror of :func:`_supports_factory`. Where
|
||||
|
|
@ -2844,7 +2844,7 @@ def is_vision_explicitly_disabled(model: str, custom_llm_provider: str | None =
|
|||
The opt-out mirror of :func:`supports_vision`: a missing declaration reads as not
|
||||
disabled, so unknown or newly added models stay eligible for image routing.
|
||||
"""
|
||||
return _is_explicitly_disabled_factory(model, custom_llm_provider, "supports_vision")
|
||||
return is_explicitly_disabled_factory(model, custom_llm_provider, "supports_vision")
|
||||
|
||||
|
||||
def supports_vision(model: str, custom_llm_provider: str | None = None) -> bool:
|
||||
|
|
@ -8746,6 +8746,7 @@ class ProviderConfigManager:
|
|||
def get_provider_responses_api_config(
|
||||
provider: LlmProviders | str,
|
||||
model: str | None = None,
|
||||
api_base: str | None = None,
|
||||
) -> BaseResponsesAPIConfig | None:
|
||||
from litellm.llms.openai_like.dynamic_config import (
|
||||
create_responses_config_class,
|
||||
|
|
@ -8767,7 +8768,7 @@ class ProviderConfigManager:
|
|||
pass
|
||||
|
||||
# Check Python classes first (custom overrides take priority)
|
||||
result: Final = ProviderConfigManager._get_python_responses_api_config(provider_enum, model)
|
||||
result: Final = ProviderConfigManager._get_python_responses_api_config(provider_enum, model, api_base)
|
||||
if result is not None:
|
||||
return result
|
||||
|
||||
|
|
@ -8783,6 +8784,7 @@ class ProviderConfigManager:
|
|||
def _get_python_responses_api_config(
|
||||
provider: LlmProviders | None,
|
||||
model: str | None = None,
|
||||
api_base: str | None = None,
|
||||
) -> BaseResponsesAPIConfig | None:
|
||||
"""Check for Python-class-based responses API configs (custom overrides)."""
|
||||
if provider is None:
|
||||
|
|
@ -8801,6 +8803,14 @@ class ProviderConfigManager:
|
|||
return litellm.AzureOpenAIOSeriesResponsesAPIConfig()
|
||||
else:
|
||||
return litellm.AzureOpenAIResponsesAPIConfig()
|
||||
elif litellm.LlmProviders.AZURE_AI == provider:
|
||||
from litellm.llms.azure_ai.common_utils import (
|
||||
azure_ai_supports_native_responses,
|
||||
)
|
||||
|
||||
if azure_ai_supports_native_responses(model, api_base):
|
||||
return litellm.AzureAIResponsesAPIConfig()
|
||||
return None
|
||||
elif litellm.LlmProviders.XAI == provider:
|
||||
return litellm.XAIResponsesAPIConfig()
|
||||
elif litellm.LlmProviders.GITHUB_COPILOT == provider:
|
||||
|
|
|
|||
|
|
@ -26105,6 +26105,7 @@
|
|||
"supports_pdf_input": true,
|
||||
"supports_prompt_caching": true,
|
||||
"supports_reasoning": true,
|
||||
"supports_minimal_reasoning_effort": false,
|
||||
"supports_response_schema": true,
|
||||
"supports_system_messages": true,
|
||||
"supports_tool_choice": true,
|
||||
|
|
@ -26162,6 +26163,7 @@
|
|||
"supports_pdf_input": true,
|
||||
"supports_prompt_caching": true,
|
||||
"supports_reasoning": true,
|
||||
"supports_minimal_reasoning_effort": false,
|
||||
"supports_response_schema": true,
|
||||
"supports_system_messages": true,
|
||||
"supports_tool_choice": true,
|
||||
|
|
@ -28111,6 +28113,7 @@
|
|||
"supports_pdf_input": true,
|
||||
"supports_prompt_caching": true,
|
||||
"supports_reasoning": true,
|
||||
"supports_minimal_reasoning_effort": false,
|
||||
"supports_response_schema": true,
|
||||
"supports_system_messages": true,
|
||||
"supports_tool_choice": true,
|
||||
|
|
@ -28170,6 +28173,7 @@
|
|||
"supports_pdf_input": true,
|
||||
"supports_prompt_caching": true,
|
||||
"supports_reasoning": true,
|
||||
"supports_minimal_reasoning_effort": false,
|
||||
"supports_response_schema": true,
|
||||
"supports_system_messages": true,
|
||||
"supports_tool_choice": true,
|
||||
|
|
@ -28592,6 +28596,7 @@
|
|||
"supports_pdf_input": true,
|
||||
"supports_prompt_caching": true,
|
||||
"supports_reasoning": true,
|
||||
"supports_minimal_reasoning_effort": false,
|
||||
"supports_response_schema": true,
|
||||
"supports_system_messages": true,
|
||||
"supports_tool_choice": true,
|
||||
|
|
@ -28649,6 +28654,7 @@
|
|||
"supports_pdf_input": true,
|
||||
"supports_prompt_caching": true,
|
||||
"supports_reasoning": true,
|
||||
"supports_minimal_reasoning_effort": false,
|
||||
"supports_response_schema": true,
|
||||
"supports_system_messages": true,
|
||||
"supports_tool_choice": true,
|
||||
|
|
|
|||
|
|
@ -30,6 +30,9 @@
|
|||
- {id: mgmt.key.health.happy_path, module: mgmt, tier: P1, surface: api, assertions: [happy_path], source: "key_management_endpoints.py:4292", rationale: "Key health endpoint"}
|
||||
- {id: mgmt.key.bulk_update.happy_path, module: mgmt, tier: P1, surface: api, assertions: [happy_path], source: "key_management_endpoints.py:2677", rationale: "Batch key updates"}
|
||||
- {id: mgmt.team.update.persists, module: mgmt, tier: P1, surface: api, assertions: [persists], source: "team_endpoints.py:1582", rationale: "Metadata/budget updates persist"}
|
||||
- {id: mgmt.team.update.team_admin_forbidden_until_enabled, module: mgmt, tier: P0, surface: api, assertions: [team_admin_forbidden_until_enabled], source: "team_admin_field_permissions.py:156", rationale: "With no team admin editable fields enabled, a team admin's /team/update is 403 and /team/info reports editing disabled"}
|
||||
- {id: mgmt.team.update.team_admin_limited_to_enabled_fields, module: mgmt, tier: P0, surface: api, assertions: [team_admin_limited_to_enabled_fields], source: "team_admin_field_permissions.py:156", rationale: "A team admin may change only the enabled fields; a request that also changes any other field is 403 and writes nothing"}
|
||||
- {id: mgmt.team.update.team_admin_resend_keeps_budget_reset, module: mgmt, tier: P1, surface: api, assertions: [team_admin_resend_keeps_budget_reset], source: "team_admin_field_permissions.py:147", fail_before_fix: proven, rationale: "A team admin resending unchanged budget settings with an enabled field must not push the team's budget reset times back"}
|
||||
- {id: mgmt.team.delete.persists, module: mgmt, tier: P1, surface: api, assertions: [persists], source: "team_endpoints.py:1750", rationale: "Deletion prevents key access"}
|
||||
- {id: mgmt.team.block.persists, module: mgmt, tier: P1, surface: api, assertions: [persists], source: "team_endpoints.py", rationale: "Block suspends all members"}
|
||||
- {id: mgmt.team.info.happy_path, module: mgmt, tier: P1, surface: api, assertions: [happy_path], source: "team_endpoints.py:2244", rationale: "Metadata+members+budgets"}
|
||||
|
|
|
|||
|
|
@ -1,5 +1,6 @@
|
|||
"""Live e2e: the /team/* management routes' block, membership, and admin-only
|
||||
contract.
|
||||
contract, plus the team settings a team admin may change on /team/update once a
|
||||
proxy admin enables them under Settings > UI > Team admin editable fields.
|
||||
|
||||
Each test creates its team/user/key resources under unique names (deleted on
|
||||
teardown) and asserts both halves of the contract: the recorded state (the info
|
||||
|
|
@ -8,21 +9,25 @@ Team writes reach the read path once their db/cache entry propagates, so the
|
|||
read-backs poll to a deadline instead of asserting once.
|
||||
|
||||
Everything the shared harness does not already model lives here: the local
|
||||
request/response models for /team/block, /team/member_update, and the
|
||||
/team/info fields (blocked flag and per-member budget) these tests assert on.
|
||||
request/response models for /team/block, /team/member_update, the partial
|
||||
/team/update, the UI settings allow-list, and the /team/info fields (blocked
|
||||
flag, limits, budgets, per-member budget, the caller's edit access) these tests
|
||||
assert on.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import time
|
||||
from collections.abc import Callable
|
||||
from typing import Literal
|
||||
from collections.abc import Callable, Generator
|
||||
from contextlib import contextmanager
|
||||
from datetime import UTC, datetime, timedelta
|
||||
from typing import Final, Literal
|
||||
|
||||
import pytest
|
||||
from pydantic import BaseModel
|
||||
|
||||
from e2e_config import unique_marker
|
||||
from e2e_http import NoBody, StreamingResponse, unwrap
|
||||
from e2e_config import settle_propagation, unique_marker
|
||||
from e2e_http import NoBody, PartialBody, StreamingResponse, unwrap
|
||||
from lifecycle import ResourceManager
|
||||
from management_client import ManagementClient
|
||||
from models import (
|
||||
|
|
@ -39,6 +44,8 @@ pytestmark = pytest.mark.e2e
|
|||
|
||||
TeamRole = Literal["admin", "user"]
|
||||
|
||||
_TEAM_TPM_LIMIT: Final = 1000
|
||||
|
||||
|
||||
class TeamBlockBody(BaseModel):
|
||||
team_id: str
|
||||
|
|
@ -66,11 +73,37 @@ class TeamMembership(BaseModel):
|
|||
litellm_budget_table: MemberBudgetTable | None = None
|
||||
|
||||
|
||||
class TeamInfoData(BaseModel):
|
||||
class CallerEditAccess(BaseModel):
|
||||
kind: Literal["unrestricted", "team_admin", "team_admin_disabled", "none"]
|
||||
editable_fields: list[str] = []
|
||||
|
||||
|
||||
class BudgetWindow(BaseModel):
|
||||
budget_duration: str
|
||||
max_budget: float
|
||||
reset_at: str | None = None
|
||||
|
||||
|
||||
class TeamCustomMetadata(BaseModel):
|
||||
cost_center: str | None = None
|
||||
|
||||
|
||||
class TeamSettings(BaseModel):
|
||||
team_alias: str | None = None
|
||||
models: list[str] = []
|
||||
tpm_limit: int | None = None
|
||||
rpm_limit: int | None = None
|
||||
max_budget: float | None = None
|
||||
budget_duration: str | None = None
|
||||
budget_limits: list[BudgetWindow] | None = None
|
||||
metadata: TeamCustomMetadata | None = None
|
||||
|
||||
|
||||
class TeamInfoData(TeamSettings):
|
||||
blocked: bool | None = None
|
||||
members_with_roles: list[MemberRoleEntry] = []
|
||||
budget_reset_at: datetime | None = None
|
||||
caller_edit_access: CallerEditAccess | None = None
|
||||
|
||||
|
||||
class TeamInfoRead(BaseModel):
|
||||
|
|
@ -79,6 +112,27 @@ class TeamInfoRead(BaseModel):
|
|||
team_memberships: list[TeamMembership] = []
|
||||
|
||||
|
||||
class TeamWithAdminNewBody(TeamNewBody):
|
||||
tpm_limit: int
|
||||
members_with_roles: list[TeamMemberEntry]
|
||||
|
||||
|
||||
class TeamSettingsChange(PartialBody, TeamSettings):
|
||||
pass
|
||||
|
||||
|
||||
class TeamSettingsUpdate(TeamSettingsChange):
|
||||
team_id: str
|
||||
|
||||
|
||||
class TeamAdminEditableFields(BaseModel):
|
||||
team_admin_editable_team_fields: list[str] = []
|
||||
|
||||
|
||||
class UiSettingsRead(BaseModel):
|
||||
values: TeamAdminEditableFields
|
||||
|
||||
|
||||
def _poll[T](client: ManagementClient, attempt: Callable[[], T | None], failure: str) -> T:
|
||||
deadline = time.monotonic() + client.proxy.poll_timeout
|
||||
while time.monotonic() < deadline:
|
||||
|
|
@ -107,17 +161,27 @@ def _generate_key(client: ManagementClient, resources: ResourceManager, body: Ke
|
|||
return key
|
||||
|
||||
|
||||
def _read_team(client: ManagementClient, team_id: str) -> TeamInfoRead:
|
||||
def _read_team(client: ManagementClient, team_id: str, caller_key: str | None = None) -> TeamInfoRead:
|
||||
return unwrap(
|
||||
client.proxy.transport.get(
|
||||
"/team/info",
|
||||
headers=client.proxy.transport.master,
|
||||
headers=client.proxy.transport.master if caller_key is None else client.proxy.transport.bearer(caller_key),
|
||||
params=TeamInfoParams(team_id=team_id),
|
||||
response_type=TeamInfoRead,
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
def _poll_team(
|
||||
client: ManagementClient, team_id: str, ready: Callable[[TeamInfoData], bool], failure: str
|
||||
) -> TeamInfoData:
|
||||
def read() -> TeamInfoData | None:
|
||||
info = _read_team(client, team_id).team_info
|
||||
return info if ready(info) else None
|
||||
|
||||
return _poll(client, read, failure)
|
||||
|
||||
|
||||
def _set_blocked(client: ManagementClient, team_id: str, *, blocked: bool) -> None:
|
||||
_ = unwrap(
|
||||
client.proxy.transport.post(
|
||||
|
|
@ -301,3 +365,218 @@ class TestTeamManagementRoutes:
|
|||
client.add_team_member(team_id, member_id)
|
||||
member_key = _generate_key(client, resources, KeyGenerateBody(user_id=member_id, team_id=team_id))
|
||||
return member_id, other_id, member_key, team_id
|
||||
|
||||
|
||||
def _team_admin_editable_fields(client: ManagementClient) -> list[str]:
|
||||
return unwrap(
|
||||
client.proxy.transport.get(
|
||||
"/get/ui_settings",
|
||||
headers=client.proxy.transport.master,
|
||||
params=NoBody(),
|
||||
response_type=UiSettingsRead,
|
||||
)
|
||||
).values.team_admin_editable_team_fields
|
||||
|
||||
|
||||
def _set_team_admin_editable_fields(client: ManagementClient, fields: list[str]) -> None:
|
||||
_ = unwrap(
|
||||
client.proxy.transport.patch(
|
||||
"/update/ui_settings",
|
||||
headers=client.proxy.transport.master,
|
||||
json=TeamAdminEditableFields(team_admin_editable_team_fields=fields),
|
||||
response_type=NoBody,
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
@contextmanager
|
||||
def _team_admins_may_edit(client: ManagementClient, fields: list[str]) -> Generator[None]:
|
||||
"""The allow-list is proxy-wide, so restore whatever was there. Other replicas pick a change up on their
|
||||
config reload, which the wait covers before any team admin call lands on one of them."""
|
||||
original = _team_admin_editable_fields(client)
|
||||
_set_team_admin_editable_fields(client, fields)
|
||||
settle_propagation(time.monotonic())
|
||||
try:
|
||||
yield
|
||||
finally:
|
||||
_set_team_admin_editable_fields(client, original)
|
||||
|
||||
|
||||
@pytest.fixture(scope="class")
|
||||
def no_team_admin_editable_fields(client: ManagementClient) -> Generator[None]:
|
||||
with _team_admins_may_edit(client, []):
|
||||
yield
|
||||
|
||||
|
||||
@pytest.fixture(scope="class")
|
||||
def tpm_limit_editable_by_team_admins(client: ManagementClient) -> Generator[None]:
|
||||
with _team_admins_may_edit(client, ["tpm_limit"]):
|
||||
yield
|
||||
|
||||
|
||||
def _team_with_admin(client: ManagementClient, resources: ResourceManager) -> tuple[str, str]:
|
||||
"""A team with a tpm_limit, and the key of a user who is an admin of that team."""
|
||||
admin_id = _create_user(client, resources, f"e2e-team-admin-{unique_marker()}@example.com")
|
||||
team_id = client.create_team(
|
||||
TeamWithAdminNewBody(
|
||||
team_alias=f"e2e-team-admin-{unique_marker()}",
|
||||
tpm_limit=_TEAM_TPM_LIMIT,
|
||||
members_with_roles=[TeamMemberEntry(role="admin", user_id=admin_id)],
|
||||
)
|
||||
)
|
||||
resources.defer(lambda: client.delete_team(team_id))
|
||||
return team_id, _generate_key(client, resources, KeyGenerateBody(user_id=admin_id))
|
||||
|
||||
|
||||
def _update_team_as(client: ManagementClient, caller_key: str, body: TeamSettingsUpdate) -> StreamingResponse:
|
||||
return client.proxy.transport.send("/team/update", headers=client.proxy.transport.bearer(caller_key), json=body)
|
||||
|
||||
|
||||
@pytest.mark.usefixtures("no_team_admin_editable_fields")
|
||||
class TestTeamAdminWithNoEditableFields:
|
||||
"""No proxy admin has enabled a team field for team admins, which is how every proxy starts."""
|
||||
|
||||
@pytest.mark.covers("mgmt.team.update.team_admin_forbidden_until_enabled")
|
||||
def test_team_admin_cannot_change_any_team_setting(
|
||||
self, client: ManagementClient, resources: ResourceManager
|
||||
) -> None:
|
||||
team_id, admin_key = _team_with_admin(client, resources)
|
||||
access = _read_team(client, team_id, admin_key).team_info.caller_edit_access
|
||||
assert access == CallerEditAccess(kind="team_admin_disabled"), (
|
||||
f"/team/info should tell the team admin that editing is disabled, got {access}"
|
||||
)
|
||||
|
||||
outcome = _update_team_as(client, admin_key, TeamSettingsUpdate(team_id=team_id, tpm_limit=5000))
|
||||
|
||||
assert outcome.status_code == 403, (
|
||||
f"/team/update by a team admin must be 403 while nothing is enabled, got {outcome.status_code}: "
|
||||
f"{outcome.body[:300]}"
|
||||
)
|
||||
assert "cannot edit team settings" in outcome.body, f"403 body should say why, got: {outcome.body[:300]}"
|
||||
tpm_limit = _read_team(client, team_id).team_info.tpm_limit
|
||||
assert tpm_limit == _TEAM_TPM_LIMIT, f"the refused update still changed tpm_limit to {tpm_limit}"
|
||||
|
||||
|
||||
@pytest.mark.usefixtures("tpm_limit_editable_by_team_admins")
|
||||
class TestTeamAdminWithTpmLimitEnabled:
|
||||
"""A proxy admin has enabled tpm_limit, so a team admin may change that setting and no other."""
|
||||
|
||||
@pytest.mark.covers("mgmt.team.update.team_admin_limited_to_enabled_fields")
|
||||
def test_team_admin_saves_the_settings_form_with_a_new_tpm_limit(
|
||||
self, client: ManagementClient, resources: ResourceManager
|
||||
) -> None:
|
||||
team_id, admin_key = _team_with_admin(client, resources)
|
||||
access = _read_team(client, team_id, admin_key).team_info.caller_edit_access
|
||||
assert access == CallerEditAccess(kind="team_admin", editable_fields=["tpm_limit"]), (
|
||||
f"/team/info should list tpm_limit as the team admin's only editable field, got {access}"
|
||||
)
|
||||
before = _read_team(client, team_id).team_info
|
||||
|
||||
outcome = _update_team_as(
|
||||
client,
|
||||
admin_key,
|
||||
TeamSettingsUpdate(team_id=team_id, team_alias=before.team_alias, models=before.models, tpm_limit=5000),
|
||||
)
|
||||
|
||||
assert outcome.status_code == 200, (
|
||||
f"a team admin resending the form with only tpm_limit changed must succeed, got {outcome.status_code}: "
|
||||
f"{outcome.body[:300]}"
|
||||
)
|
||||
after = _poll_team(
|
||||
client, team_id, lambda info: info.tpm_limit == 5000, "/team/info never reflected tpm_limit=5000"
|
||||
)
|
||||
assert after.model_copy(update={"tpm_limit": _TEAM_TPM_LIMIT}) == before, (
|
||||
f"the update changed more than tpm_limit: before {before}, after {after}"
|
||||
)
|
||||
|
||||
@pytest.mark.covers("mgmt.team.update.team_admin_limited_to_enabled_fields")
|
||||
@pytest.mark.parametrize(
|
||||
"change",
|
||||
[
|
||||
pytest.param(TeamSettingsChange(rpm_limit=10), id="rpm_limit"),
|
||||
pytest.param(TeamSettingsChange(max_budget=0.5), id="max_budget"),
|
||||
pytest.param(TeamSettingsChange(team_alias="renamed-by-team-admin"), id="team_alias"),
|
||||
pytest.param(TeamSettingsChange(models=["gemini-2.5-flash"]), id="models"),
|
||||
pytest.param(TeamSettingsChange(budget_duration="1d"), id="budget_duration"),
|
||||
pytest.param(TeamSettingsChange(metadata=TeamCustomMetadata(cost_center="team-admin")), id="metadata"),
|
||||
],
|
||||
)
|
||||
def test_team_admin_cannot_change_a_setting_that_is_not_enabled(
|
||||
self, client: ManagementClient, resources: ResourceManager, change: TeamSettingsChange
|
||||
) -> None:
|
||||
(field,) = change.model_fields_set
|
||||
team_id, admin_key = _team_with_admin(client, resources)
|
||||
before = _read_team(client, team_id).team_info
|
||||
|
||||
outcome = _update_team_as(
|
||||
client,
|
||||
admin_key,
|
||||
TeamSettingsUpdate.model_validate(
|
||||
{**change.model_dump(exclude_unset=True), "team_id": team_id, "tpm_limit": 5000}
|
||||
),
|
||||
)
|
||||
|
||||
assert outcome.status_code == 403, (
|
||||
f"a team admin changing {field} must be 403, got {outcome.status_code}: {outcome.body[:300]}"
|
||||
)
|
||||
assert f"'{field}'" in outcome.body, f"403 body should name {field}, got: {outcome.body[:300]}"
|
||||
after = _read_team(client, team_id).team_info
|
||||
assert after == before, (
|
||||
f"the refused update still wrote to the team, the enabled tpm_limit included: before {before}, "
|
||||
f"after {after}"
|
||||
)
|
||||
|
||||
@pytest.mark.covers("mgmt.team.update.team_admin_resend_keeps_budget_reset")
|
||||
def test_team_admin_resending_the_budget_settings_keeps_the_next_budget_reset(
|
||||
self, client: ManagementClient, resources: ResourceManager
|
||||
) -> None:
|
||||
"""A 120s budget resets at the start of the minute after next. Resending it once the next minute has
|
||||
started would push that reset a minute later, while the stored reset is still a minute out, so the
|
||||
proxy's budget reset job cannot be what moves it."""
|
||||
team_id, admin_key = _team_with_admin(client, resources)
|
||||
_ = unwrap(
|
||||
client.proxy.transport.post(
|
||||
"/team/update",
|
||||
headers=client.proxy.transport.master,
|
||||
json=TeamSettingsUpdate(
|
||||
team_id=team_id,
|
||||
budget_duration="120s",
|
||||
budget_limits=[BudgetWindow(budget_duration="120s", max_budget=5.0)],
|
||||
),
|
||||
response_type=NoBody,
|
||||
)
|
||||
)
|
||||
budgeted = _poll_team(
|
||||
client,
|
||||
team_id,
|
||||
lambda info: info.budget_reset_at is not None and bool(info.budget_limits),
|
||||
"/team/info never reflected the 120s budget the proxy admin set",
|
||||
)
|
||||
assert budgeted.budget_reset_at is not None
|
||||
next_minute = budgeted.budget_reset_at - timedelta(seconds=58)
|
||||
time.sleep(max(0.0, (next_minute - datetime.now(UTC)).total_seconds()))
|
||||
|
||||
outcome = _update_team_as(
|
||||
client,
|
||||
admin_key,
|
||||
TeamSettingsUpdate(
|
||||
team_id=team_id,
|
||||
tpm_limit=5000,
|
||||
budget_duration=budgeted.budget_duration,
|
||||
budget_limits=budgeted.budget_limits,
|
||||
),
|
||||
)
|
||||
|
||||
assert outcome.status_code == 200, (
|
||||
f"resending unchanged budget settings with a new tpm_limit must succeed, got {outcome.status_code}: "
|
||||
f"{outcome.body[:300]}"
|
||||
)
|
||||
after = _poll_team(
|
||||
client, team_id, lambda info: info.tpm_limit == 5000, "/team/info never reflected tpm_limit=5000"
|
||||
)
|
||||
assert after.budget_reset_at == budgeted.budget_reset_at, (
|
||||
f"the team admin pushed the budget reset from {budgeted.budget_reset_at} to {after.budget_reset_at}"
|
||||
)
|
||||
assert after.budget_limits == budgeted.budget_limits, (
|
||||
f"the team admin pushed the budget window resets from {budgeted.budget_limits} to {after.budget_limits}"
|
||||
)
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
114
tests/integration/_support/process.py
Normal file
114
tests/integration/_support/process.py
Normal 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"
|
||||
116
tests/integration/_support/redis_process.py
Normal file
116
tests/integration/_support/redis_process.py
Normal 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()
|
||||
25
tests/integration/_support/sigv4.py
Normal file
25
tests/integration/_support/sigv4.py
Normal 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()
|
||||
|
|
@ -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)
|
||||
|
|
|
|||
113
tests/integration/_support/wire.py
Normal file
113
tests/integration/_support/wire.py
Normal 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}"
|
||||
|
|
@ -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"
|
||||
]
|
||||
}
|
||||
}
|
||||
|
|
|
|||
98
tests/integration/database/test_partition_transactions.py
Normal file
98
tests/integration/database/test_partition_transactions.py
Normal 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,)) == []
|
||||
138
tests/integration/database/test_reader_writer_regeneration.py
Normal file
138
tests/integration/database/test_reader_writer_regeneration.py
Normal 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,)) == []
|
||||
125
tests/integration/database/test_transaction_atomicity.py
Normal file
125
tests/integration/database/test_transaction_atomicity.py
Normal 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,)) == []
|
||||
123
tests/integration/pricing/test_price_precedence.py
Normal file
123
tests/integration/pricing/test_price_precedence.py
Normal 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
|
||||
61
tests/integration/providers/test_anthropic_wire.py
Normal file
61
tests/integration/providers/test_anthropic_wire.py
Normal 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)
|
||||
99
tests/integration/providers/test_bedrock_auth_wire.py
Normal file
99
tests/integration/providers/test_bedrock_auth_wire.py
Normal 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
|
||||
|
|
@ -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)
|
||||
111
tests/integration/providers/test_s3_wire.py
Normal file
111
tests/integration/providers/test_s3_wire.py
Normal 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)
|
||||
98
tests/integration/routing/test_observed_routing.py
Normal file
98
tests/integration/routing/test_observed_routing.py
Normal 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"]
|
||||
59
tests/integration/routing/test_redis_recovery.py
Normal file
59
tests/integration/routing/test_redis_recovery.py
Normal 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() == []
|
||||
234
tests/integration/spend/test_cache_and_quota.py
Normal file
234
tests/integration/spend/test_cache_and_quota.py
Normal 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},
|
||||
]
|
||||
149
tests/integration/streaming/test_stream_contracts.py
Normal file
149
tests/integration/streaming/test_stream_contracts.py
Normal 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
|
||||
|
|
@ -261,7 +261,6 @@ def test_aaparallel_function_call_with_anthropic_thinking(model):
|
|||
|
||||
from litellm.types.utils import ChatCompletionMessageToolCall, Function, Message
|
||||
|
||||
|
||||
_PARALLEL_TOOL_HISTORY_MESSAGES = [
|
||||
{
|
||||
"role": "user",
|
||||
|
|
@ -293,20 +292,11 @@ _PARALLEL_TOOL_HISTORY_MESSAGES = [
|
|||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"model, messages, expect_unsupported_params_error",
|
||||
"model, messages",
|
||||
[
|
||||
# Bedrock Converse still requires modify_params to inject the dummy tool.
|
||||
(
|
||||
"us.anthropic.claude-sonnet-4-5-20250929-v1:0",
|
||||
_PARALLEL_TOOL_HISTORY_MESSAGES,
|
||||
True,
|
||||
),
|
||||
# Anthropic Messages API: dummy tool is injected without modify_params.
|
||||
(
|
||||
"claude-haiku-4-5-20251001",
|
||||
_PARALLEL_TOOL_HISTORY_MESSAGES,
|
||||
False,
|
||||
),
|
||||
# Anthropic Messages API: a dummy tool is injected without modify_params,
|
||||
# so tool history with no tools= completes instead of raising.
|
||||
("claude-haiku-4-5-20251001", _PARALLEL_TOOL_HISTORY_MESSAGES),
|
||||
(
|
||||
"us.anthropic.claude-sonnet-4-5-20250929-v1:0",
|
||||
[
|
||||
|
|
@ -315,7 +305,6 @@ _PARALLEL_TOOL_HISTORY_MESSAGES = [
|
|||
"content": "What's the weather like in San Francisco, Tokyo, and Paris? - give me 3 responses",
|
||||
}
|
||||
],
|
||||
False,
|
||||
),
|
||||
(
|
||||
"claude-haiku-4-5-20251001",
|
||||
|
|
@ -325,48 +314,34 @@ _PARALLEL_TOOL_HISTORY_MESSAGES = [
|
|||
"content": "What's the weather like in San Francisco, Tokyo, and Paris? - give me 3 responses",
|
||||
}
|
||||
],
|
||||
False,
|
||||
),
|
||||
],
|
||||
)
|
||||
def test_parallel_function_call_anthropic_error_msg(
|
||||
model, messages, expect_unsupported_params_error
|
||||
):
|
||||
def test_parallel_function_call_anthropic_error_msg(model, messages):
|
||||
"""
|
||||
Tool history without an explicit ``tools`` param:
|
||||
Tool history without an explicit ``tools`` param must complete, not raise.
|
||||
|
||||
- Bedrock **Converse** still raises ``UnsupportedParamsError`` unless
|
||||
``litellm.modify_params`` is enabled (dummy tool is only added there).
|
||||
- **Anthropic** (and Bedrock Invoke via ``AnthropicConfig.transform_request``)
|
||||
always get a dummy tool so CLIs work with ``modify_params`` left off.
|
||||
|
||||
Reference Issue: https://github.com/BerriAI/litellm/issues/5747, https://github.com/BerriAI/litellm/issues/5388
|
||||
Anthropic (and Bedrock Invoke via ``AnthropicConfig.transform_request``)
|
||||
inject a dummy tool so CLIs work with ``modify_params`` left off. Bedrock
|
||||
Converse's no-raise behavior is covered offline in
|
||||
``tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py``
|
||||
(see #24158, #27138), which needs no live credentials.
|
||||
"""
|
||||
# Ensure modify_params is False so Bedrock Converse path still raises.
|
||||
# Force modify_params off as a clean baseline: it exercises the Anthropic
|
||||
# dummy-tool path, which injects regardless of modify_params
|
||||
# (other tests in this file set it to True and don't reset it)
|
||||
original_modify_params = litellm.modify_params
|
||||
litellm.modify_params = False
|
||||
try:
|
||||
litellm.set_verbose = True
|
||||
|
||||
if expect_unsupported_params_error:
|
||||
with pytest.raises(litellm.UnsupportedParamsError) as e:
|
||||
litellm.completion(
|
||||
model=model,
|
||||
messages=messages,
|
||||
temperature=0.2,
|
||||
seed=22,
|
||||
drop_params=True,
|
||||
)
|
||||
else:
|
||||
second_response = litellm.completion(
|
||||
model=model,
|
||||
messages=messages,
|
||||
temperature=0.2,
|
||||
seed=22,
|
||||
drop_params=True,
|
||||
) # get a new response from the model where it can see the function response
|
||||
print("second response\n", second_response)
|
||||
second_response = litellm.completion(
|
||||
model=model,
|
||||
messages=messages,
|
||||
temperature=0.2,
|
||||
seed=22,
|
||||
drop_params=True,
|
||||
) # get a new response from the model where it can see the function response
|
||||
print("second response\n", second_response)
|
||||
except litellm.InternalServerError as e:
|
||||
print(e)
|
||||
except litellm.RateLimitError as e:
|
||||
|
|
|
|||
|
|
@ -10,12 +10,11 @@ Pins the five helpers
|
|||
|
||||
Driven through /team/new + /team/update.
|
||||
|
||||
Structural finding, updated: /team/new loads the org via `get_org_object`
|
||||
WITH `include_budget_table=True`, so the org max_budget / org tpm / org rpm
|
||||
guards inside `_check_org_team_limits` are live there and are pinned as
|
||||
enforced below. /team/update still loads the org without the budget
|
||||
relation, so its budget guards remain no-ops. The `models` subset guard IS
|
||||
reachable on both because it reads `org_table.models` directly. The
|
||||
Structural finding, updated: /team/new and /team/update both load the org
|
||||
via `get_org_object` WITH `include_budget_table=True`, so the org max_budget /
|
||||
org tpm / org rpm guards inside `_check_org_team_limits` are live on both and
|
||||
are pinned as enforced below. The `models` subset guard reads
|
||||
`org_table.models` directly. The
|
||||
`_check_user_team_limits` guards reach all branches through
|
||||
`user_api_key_dict`, no relation include needed.
|
||||
"""
|
||||
|
|
@ -139,9 +138,8 @@ async def test_check_org_team_limits_models_subset(
|
|||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# _check_org_team_limits — budget / tpm / rpm live on /team/new since its
|
||||
# get_org_object call passes include_budget_table=True. (/team/update still
|
||||
# loads the org without the budget relation, so its guards remain no-ops.)
|
||||
# _check_org_team_limits — budget / tpm / rpm live on /team/new and
|
||||
# /team/update since both get_org_object calls pass include_budget_table=True.
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
_ORG_BUDGET_ENFORCED_SCENARIOS = [
|
||||
|
|
@ -216,6 +214,35 @@ async def test_check_org_team_limits_budget_enforced(
|
|||
assert len(rows) == (1 if expected_status == 200 else 0)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"org_budget,body_extras,expected_status",
|
||||
[(b, c, d) for (_id, b, c, d) in _ORG_BUDGET_ENFORCED_SCENARIOS],
|
||||
ids=[s[0] for s in _ORG_BUDGET_ENFORCED_SCENARIOS],
|
||||
)
|
||||
async def test_check_org_team_limits_budget_enforced_on_update(
|
||||
org_budget,
|
||||
body_extras: Dict[str, Any],
|
||||
expected_status: int,
|
||||
proxy_client,
|
||||
prisma,
|
||||
scratch,
|
||||
world,
|
||||
):
|
||||
org_id = await create_scratch_org(prisma, scratch.prefix, **org_budget)
|
||||
team_id = await create_scratch_team(prisma, scratch.tag("team"), organization_id=org_id)
|
||||
seeder = world.keys[Actor.PROXY_ADMIN].cleartext
|
||||
resp = await proxy_client.post(
|
||||
"/team/update",
|
||||
headers={"Authorization": f"Bearer {seeder}"},
|
||||
json={"team_id": team_id, **body_extras},
|
||||
)
|
||||
assert resp.status_code == expected_status, f"{body_extras!r} → {resp.status_code}: {resp.text}"
|
||||
row = await prisma.db.litellm_teamtable.find_unique(where={"team_id": team_id})
|
||||
assert row is not None
|
||||
persisted = {field: getattr(row, field) for field in body_extras}
|
||||
assert (persisted == body_extras) == (expected_status == 200)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# _check_user_team_limits — fires for standalone (no-org) teams created by
|
||||
# a non-admin caller. Each guard reads from user_api_key_dict / user_obj.
|
||||
|
|
@ -310,65 +337,40 @@ async def test_check_user_team_limits(
|
|||
# /team/update path — budget authority.
|
||||
#
|
||||
# The caller's PERSONAL limits are never applied on update (that compared the
|
||||
# wrong thing). But raising a team's spend ceiling is reserved for proxy admins:
|
||||
# a team admin may keep or LOWER the budget, only a proxy admin may RAISE it.
|
||||
# _check_user_team_limits() only runs on /team/new.
|
||||
# wrong thing). Raising a team's spend ceiling is reserved for proxy admins.
|
||||
# max_budget is not on the team-admin allow-list yet (LIT-5722), so a team
|
||||
# admin is refused in either direction; the raise-only guard underneath the
|
||||
# allow-list is pinned in the unit tests. _check_user_team_limits() only runs
|
||||
# on /team/new.
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
async def test_team_admin_raise_budget_blocked(proxy_client, prisma, scratch):
|
||||
"""A team admin cannot raise the team's budget; the block is NOT based on
|
||||
their personal budget (which here is higher than the requested value)."""
|
||||
caller_cleartext = await _seed_scratch_actor_with_caps(
|
||||
prisma,
|
||||
scratch.prefix,
|
||||
max_budget=100000.0, # generous personal budget; must not matter
|
||||
)
|
||||
creator_user_id = f"{scratch.prefix}-team-creator"
|
||||
@pytest.mark.parametrize(
|
||||
"personal_budget,requested_budget",
|
||||
[(100000.0, 999.0), (10.0, 300.0)],
|
||||
ids=["raise_with_generous_personal_budget", "lower_with_tiny_personal_budget"],
|
||||
)
|
||||
async def test_team_admin_cannot_change_budget_while_max_budget_is_not_editable(
|
||||
proxy_client, prisma, scratch, personal_budget: float, requested_budget: float
|
||||
):
|
||||
caller_cleartext = await _seed_scratch_actor_with_caps(prisma, scratch.prefix, max_budget=personal_budget)
|
||||
team_id = await create_scratch_team(
|
||||
prisma,
|
||||
team_id=scratch.tag("team"),
|
||||
admin_user_ids=[creator_user_id],
|
||||
max_budget=50.0,
|
||||
)
|
||||
# Raise the team budget 50 -> 999 as a team admin.
|
||||
resp = await proxy_client.post(
|
||||
"/team/update",
|
||||
headers={"Authorization": f"Bearer {caller_cleartext}"},
|
||||
json={"team_id": team_id, "max_budget": 999.0},
|
||||
)
|
||||
assert resp.status_code == 403, resp.text
|
||||
|
||||
row = await prisma.db.litellm_teamtable.find_unique(where={"team_id": team_id})
|
||||
assert row is not None
|
||||
assert row.max_budget == 50.0, "team budget must not change on a blocked raise"
|
||||
|
||||
|
||||
async def test_team_admin_lower_budget_allowed(proxy_client, prisma, scratch):
|
||||
"""A team admin may freely lower (or keep) the team's budget."""
|
||||
caller_cleartext = await _seed_scratch_actor_with_caps(
|
||||
prisma,
|
||||
scratch.prefix,
|
||||
max_budget=10.0, # below both the old and new team budget; must not matter
|
||||
)
|
||||
creator_user_id = f"{scratch.prefix}-team-creator"
|
||||
team_id = await create_scratch_team(
|
||||
prisma,
|
||||
team_id=scratch.tag("team"),
|
||||
admin_user_ids=[creator_user_id],
|
||||
admin_user_ids=[f"{scratch.prefix}-team-creator"],
|
||||
max_budget=500.0,
|
||||
)
|
||||
# Lower the team budget 500 -> 300 as a team admin.
|
||||
resp = await proxy_client.post(
|
||||
"/team/update",
|
||||
headers={"Authorization": f"Bearer {caller_cleartext}"},
|
||||
json={"team_id": team_id, "max_budget": 300.0},
|
||||
json={"team_id": team_id, "max_budget": requested_budget},
|
||||
)
|
||||
assert resp.status_code == 200, resp.text
|
||||
assert resp.status_code == 403, resp.text
|
||||
assert "Team admin editable fields" in resp.text, resp.text
|
||||
|
||||
row = await prisma.db.litellm_teamtable.find_unique(where={"team_id": team_id})
|
||||
assert row is not None
|
||||
assert row.max_budget == 300.0, "team admin should be able to lower the budget"
|
||||
assert row.max_budget == 500.0, "a refused update must leave the team budget unchanged"
|
||||
|
||||
|
||||
async def test_proxy_admin_raise_budget_allowed(proxy_client, prisma, scratch):
|
||||
|
|
|
|||
|
|
@ -9,31 +9,31 @@ pytestmark = pytest.mark.asyncio(loop_scope="session")
|
|||
|
||||
|
||||
# POST /team/update — actor x team-shape matrix (shapes built by _seed_target).
|
||||
# Each request carries the team's own organization_id so a non-proxy-admin can
|
||||
# reach the org-scoped branch of the route-permission gate (401 on denial),
|
||||
# which fronts the handler's _verify_team_access. Only PROXY_ADMIN and an
|
||||
# ORG_ADMIN of the team's org pass: an internal_user team admin is filtered by
|
||||
# the route gate before _verify_team_access's team-admin branch is reached.
|
||||
# The route is self-managed (LIT-5722), so every authenticated caller reaches
|
||||
# update_team and denials are the handler's 403, never the route gate's 401.
|
||||
# Only PROXY_ADMIN and an ORG_ADMIN of the team's org pass: a team admin is
|
||||
# admitted by _resolve_team_access but then refused because no team field is
|
||||
# enabled for team admins (team_admin_editable_team_fields defaults to empty).
|
||||
MARKER_ALIAS = "behavior-pin-update-marker-alias"
|
||||
|
||||
_MATRIX = [
|
||||
("alpha/proxy_admin", Actor.PROXY_ADMIN, "alpha", 200),
|
||||
("alpha/org_admin", Actor.ORG_ADMIN, "alpha", 200),
|
||||
("alpha/team_admin", Actor.TEAM_ADMIN, "alpha", 401),
|
||||
("alpha/internal_user", Actor.INTERNAL_USER, "alpha", 401),
|
||||
("alpha/owner", Actor.OWNER, "alpha", 401),
|
||||
("alpha/unrelated_same_org", Actor.UNRELATED_SAME_ORG, "alpha", 401),
|
||||
("alpha/cross_org_user", Actor.CROSS_ORG_USER, "alpha", 401),
|
||||
("alpha/service_account", Actor.SERVICE_ACCOUNT, "alpha", 401),
|
||||
("alpha/org_b_admin", Actor.ORG_B_ADMIN, "alpha", 401),
|
||||
("alpha/team_admin", Actor.TEAM_ADMIN, "alpha", 403),
|
||||
("alpha/internal_user", Actor.INTERNAL_USER, "alpha", 403),
|
||||
("alpha/owner", Actor.OWNER, "alpha", 403),
|
||||
("alpha/unrelated_same_org", Actor.UNRELATED_SAME_ORG, "alpha", 403),
|
||||
("alpha/cross_org_user", Actor.CROSS_ORG_USER, "alpha", 403),
|
||||
("alpha/service_account", Actor.SERVICE_ACCOUNT, "alpha", 403),
|
||||
("alpha/org_b_admin", Actor.ORG_B_ADMIN, "alpha", 403),
|
||||
("beta/proxy_admin", Actor.PROXY_ADMIN, "beta", 200),
|
||||
("beta/org_admin", Actor.ORG_ADMIN, "beta", 401),
|
||||
("beta/team_admin", Actor.TEAM_ADMIN, "beta", 401),
|
||||
("beta/internal_user", Actor.INTERNAL_USER, "beta", 401),
|
||||
("beta/owner", Actor.OWNER, "beta", 401),
|
||||
("beta/unrelated_same_org", Actor.UNRELATED_SAME_ORG, "beta", 401),
|
||||
("beta/cross_org_user", Actor.CROSS_ORG_USER, "beta", 401),
|
||||
("beta/service_account", Actor.SERVICE_ACCOUNT, "beta", 401),
|
||||
("beta/org_admin", Actor.ORG_ADMIN, "beta", 403),
|
||||
("beta/team_admin", Actor.TEAM_ADMIN, "beta", 403),
|
||||
("beta/internal_user", Actor.INTERNAL_USER, "beta", 403),
|
||||
("beta/owner", Actor.OWNER, "beta", 403),
|
||||
("beta/unrelated_same_org", Actor.UNRELATED_SAME_ORG, "beta", 403),
|
||||
("beta/cross_org_user", Actor.CROSS_ORG_USER, "beta", 403),
|
||||
("beta/service_account", Actor.SERVICE_ACCOUNT, "beta", 403),
|
||||
("beta/org_b_admin", Actor.ORG_B_ADMIN, "beta", 200),
|
||||
]
|
||||
|
||||
|
|
@ -110,8 +110,9 @@ async def test_team_update_org_admin_resolved_from_team_without_org_context(
|
|||
):
|
||||
"""With no organization_id in the body the route gate resolves the target
|
||||
team's org from team_id, so an org admin of the team's own org is allowed
|
||||
(200), same as PROXY_ADMIN. A team admin of that same team stays denied
|
||||
(401): the resolution grants org admins access, not team admins."""
|
||||
(200), same as PROXY_ADMIN. A team admin of that same team reaches the
|
||||
handler but is refused (403) until a proxy admin enables fields for team
|
||||
admins, and the response says so."""
|
||||
await _seed_target(prisma, world, "alpha", scratch.prefix)
|
||||
|
||||
allowed_org_admin = await proxy_client.post(
|
||||
|
|
@ -133,21 +134,25 @@ async def test_team_update_org_admin_resolved_from_team_without_org_context(
|
|||
headers={"Authorization": f"Bearer {world.keys[Actor.TEAM_ADMIN].cleartext}"},
|
||||
json={"team_id": scratch.prefix, "team_alias": MARKER_ALIAS},
|
||||
)
|
||||
assert denied_team_admin.status_code == 401, denied_team_admin.text
|
||||
assert denied_team_admin.status_code == 403, denied_team_admin.text
|
||||
assert "cannot edit team settings" in denied_team_admin.text, denied_team_admin.text
|
||||
assert "Team admin editable fields" in denied_team_admin.text, denied_team_admin.text
|
||||
|
||||
|
||||
# Relocation gate — moving a team to a different org. The scratch team starts
|
||||
# in ORG_A; each scenario relocates it to ORG_B. PROXY_ADMIN bypasses;
|
||||
# ORG_B_ADMIN clears the route gate (dest-org admin) but fails
|
||||
# _verify_team_access on the source team (403); the rest fail the route gate
|
||||
# (401). The relocation-*allowed* branch (caller is org admin of both orgs) is
|
||||
# covered by test_team_update_org_relocation_allowed_for_dual_org_admin below.
|
||||
# ORG_B_ADMIN reaches the handler but holds no role on the source team (403);
|
||||
# ORG_ADMIN holds the source team but not the destination org (403 from the
|
||||
# relocation gate); the team admin is refused by the empty field allow-list and
|
||||
# the internal user holds no role at all (403). The relocation-*allowed* branch
|
||||
# (caller is org admin of both orgs) is covered by
|
||||
# test_team_update_org_relocation_allowed_for_dual_org_admin below.
|
||||
_RELOCATION = [
|
||||
("proxy_admin", Actor.PROXY_ADMIN, 200),
|
||||
("org_b_admin", Actor.ORG_B_ADMIN, 403),
|
||||
("org_admin", Actor.ORG_ADMIN, 401),
|
||||
("team_admin", Actor.TEAM_ADMIN, 401),
|
||||
("internal_user", Actor.INTERNAL_USER, 401),
|
||||
("org_admin", Actor.ORG_ADMIN, 403),
|
||||
("team_admin", Actor.TEAM_ADMIN, 403),
|
||||
("internal_user", Actor.INTERNAL_USER, 403),
|
||||
]
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -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"}
|
||||
|
|
|
|||
|
|
@ -19,12 +19,12 @@ to 0 when the only update we saw was the cursor, allowing the
|
|||
text-based fallback to estimate from the real completion text.
|
||||
"""
|
||||
|
||||
|
||||
import pytest
|
||||
|
||||
|
||||
import litellm
|
||||
from litellm.litellm_core_utils.streaming_chunk_builder_utils import ChunkProcessor
|
||||
from litellm.types.utils import (
|
||||
CompletionTokensDetailsWrapper,
|
||||
Delta,
|
||||
ModelResponseStream,
|
||||
StreamingChoices,
|
||||
|
|
@ -35,6 +35,7 @@ from litellm.types.utils import (
|
|||
def _make_chunk(
|
||||
*,
|
||||
content: str = "",
|
||||
reasoning_content: str | None = None,
|
||||
usage: Usage = None,
|
||||
finish_reason: str = None,
|
||||
custom_llm_provider: str = "anthropic",
|
||||
|
|
@ -48,7 +49,7 @@ def _make_chunk(
|
|||
StreamingChoices(
|
||||
finish_reason=finish_reason,
|
||||
index=0,
|
||||
delta=Delta(content=content, role="assistant"),
|
||||
delta=Delta(content=content, role="assistant", reasoning_content=reasoning_content),
|
||||
)
|
||||
],
|
||||
usage=usage,
|
||||
|
|
@ -69,9 +70,7 @@ class TestAnthropicCursorBug:
|
|||
token_counter fallback can estimate from completion text.
|
||||
"""
|
||||
# Anthropic message_start: input_tokens accurate, output_tokens=1 cursor
|
||||
message_start = _make_chunk(
|
||||
usage=Usage(prompt_tokens=1024, completion_tokens=1, total_tokens=1025)
|
||||
)
|
||||
message_start = _make_chunk(usage=Usage(prompt_tokens=1024, completion_tokens=1, total_tokens=1025))
|
||||
# Several content_block_delta chunks (no usage attached)
|
||||
text_chunks = [
|
||||
_make_chunk(content="Hello"),
|
||||
|
|
@ -97,9 +96,7 @@ class TestAnthropicCursorBug:
|
|||
Normal complete stream: message_start cursor=1, then message_delta=3847.
|
||||
Last-wins must give 3847 (the real value).
|
||||
"""
|
||||
message_start = _make_chunk(
|
||||
usage=Usage(prompt_tokens=1024, completion_tokens=1, total_tokens=1025)
|
||||
)
|
||||
message_start = _make_chunk(usage=Usage(prompt_tokens=1024, completion_tokens=1, total_tokens=1025))
|
||||
text_chunks = [_make_chunk(content=t) for t in ["Hello", " world", "!"]]
|
||||
# message_delta with the real cumulative output_tokens
|
||||
message_delta = _make_chunk(
|
||||
|
|
@ -119,19 +116,14 @@ class TestAnthropicCursorBug:
|
|||
End-to-end via calculate_usage(): cursor-only stream + real completion
|
||||
text should produce a token-counter estimate, NOT 1.
|
||||
"""
|
||||
message_start = _make_chunk(
|
||||
usage=Usage(prompt_tokens=1024, completion_tokens=1, total_tokens=1025)
|
||||
)
|
||||
message_start = _make_chunk(usage=Usage(prompt_tokens=1024, completion_tokens=1, total_tokens=1025))
|
||||
# ~50 visible chars ≈ ~12 tokens (anthropic-style tokenizer ballpark)
|
||||
text_chunks = [
|
||||
_make_chunk(content="Based on your question, I think the answer is "),
|
||||
_make_chunk(content="forty-two. Here is my reasoning: "),
|
||||
]
|
||||
chunks = [message_start, *text_chunks]
|
||||
completion_output = (
|
||||
"Based on your question, I think the answer is forty-two. "
|
||||
"Here is my reasoning: "
|
||||
)
|
||||
completion_output = "Based on your question, I think the answer is forty-two. Here is my reasoning: "
|
||||
|
||||
processor = ChunkProcessor(chunks=chunks, messages=[])
|
||||
usage = processor.calculate_usage(
|
||||
|
|
@ -149,9 +141,7 @@ class TestAnthropicCursorBug:
|
|||
|
||||
def test_cache_fields_preserved_from_message_start(self):
|
||||
"""cache_read / cache_creation come from message_start and must survive."""
|
||||
message_start_usage = Usage(
|
||||
prompt_tokens=1024, completion_tokens=1, total_tokens=1025
|
||||
)
|
||||
message_start_usage = Usage(prompt_tokens=1024, completion_tokens=1, total_tokens=1025)
|
||||
# Anthropic puts these in message_start
|
||||
message_start_usage.cache_read_input_tokens = 512
|
||||
message_start_usage.cache_creation_input_tokens = 128
|
||||
|
|
@ -193,9 +183,7 @@ class TestAnthropicCursorBug:
|
|||
on a 1-token string also gives ~1, so billing is still approximately
|
||||
correct. This test pins that the result is sane (1 or 0).
|
||||
"""
|
||||
message_start = _make_chunk(
|
||||
usage=Usage(prompt_tokens=20, completion_tokens=1, total_tokens=21)
|
||||
)
|
||||
message_start = _make_chunk(usage=Usage(prompt_tokens=20, completion_tokens=1, total_tokens=21))
|
||||
text_chunk = _make_chunk(content="Yes.")
|
||||
# Anthropic's message_delta also gives output_tokens=1 in this case
|
||||
message_delta = _make_chunk(
|
||||
|
|
@ -231,9 +219,7 @@ class TestAnthropicCursorBug:
|
|||
must fire so token_counter estimates from completion text instead of
|
||||
billing the placeholder.
|
||||
"""
|
||||
message_start_usage = Usage(
|
||||
prompt_tokens=1024, completion_tokens=1, total_tokens=1025
|
||||
)
|
||||
message_start_usage = Usage(prompt_tokens=1024, completion_tokens=1, total_tokens=1025)
|
||||
message_start_usage.cache_read_input_tokens = 4096
|
||||
message_start = _make_chunk(usage=message_start_usage)
|
||||
# Subsequent chunks with cache fields but no completion_tokens
|
||||
|
|
@ -253,6 +239,114 @@ class TestAnthropicCursorBug:
|
|||
"Reset to 0 forces token_counter fallback."
|
||||
)
|
||||
|
||||
@pytest.mark.parametrize("placeholder", [1, 3, 8])
|
||||
def test_interrupted_reasoning_only_stream_estimates_from_reasoning(self, placeholder: int):
|
||||
message_start = _make_chunk(
|
||||
usage=Usage(
|
||||
prompt_tokens=100,
|
||||
completion_tokens=placeholder,
|
||||
total_tokens=100 + placeholder,
|
||||
completion_tokens_details=CompletionTokensDetailsWrapper(reasoning_tokens=0, text_tokens=placeholder),
|
||||
)
|
||||
)
|
||||
reasoning_text = "Let me work through the scheduling constraints step by step. " * 40
|
||||
reasoning_chunks = [
|
||||
_make_chunk(reasoning_content=reasoning_text[i : i + 50]) for i in range(0, len(reasoning_text), 50)
|
||||
]
|
||||
|
||||
response = litellm.stream_chunk_builder(
|
||||
chunks=[message_start, *reasoning_chunks],
|
||||
messages=[{"role": "user", "content": "Plan the schedule."}],
|
||||
)
|
||||
|
||||
assert response.choices[0].message.reasoning_content == reasoning_text
|
||||
reasoning_tokens = response.usage.completion_tokens_details.reasoning_tokens
|
||||
assert reasoning_tokens > placeholder
|
||||
assert response.usage.completion_tokens == reasoning_tokens, (
|
||||
f"Expected completion_tokens to be the reasoning estimate, got "
|
||||
f"completion_tokens={response.usage.completion_tokens} reasoning_tokens={reasoning_tokens}"
|
||||
)
|
||||
assert response.usage.total_tokens == response.usage.prompt_tokens + reasoning_tokens
|
||||
details = response.usage.completion_tokens_details
|
||||
assert details.text_tokens + details.reasoning_tokens == response.usage.completion_tokens
|
||||
|
||||
def test_fallback_counts_reasoning_and_text_together(self):
|
||||
reasoning = "First I should check whether the input is sorted. " * 10
|
||||
text = "The list is already sorted, so no work is needed."
|
||||
chunks = [_make_chunk(reasoning_content=reasoning), _make_chunk(content=text)]
|
||||
|
||||
response = litellm.stream_chunk_builder(chunks=chunks, messages=[{"role": "user", "content": "Sort it."}])
|
||||
|
||||
text_only = litellm.token_counter(model="claude-sonnet-4-6", text=text, count_response_tokens=True)
|
||||
details = response.usage.completion_tokens_details
|
||||
assert details.reasoning_tokens > 0
|
||||
assert response.usage.completion_tokens == text_only + details.reasoning_tokens
|
||||
assert details.text_tokens == text_only
|
||||
|
||||
def test_lone_usage_event_with_finish_reason_is_trusted(self):
|
||||
chunks = [
|
||||
_make_chunk(content="Yes, "),
|
||||
_make_chunk(content="that works."),
|
||||
_make_chunk(
|
||||
usage=Usage(prompt_tokens=20, completion_tokens=5, total_tokens=25),
|
||||
finish_reason="stop",
|
||||
),
|
||||
]
|
||||
processor = ChunkProcessor(chunks=chunks, messages=[])
|
||||
result = processor._calculate_usage_per_chunk(chunks=chunks)
|
||||
assert result["completion_tokens"] == 5
|
||||
|
||||
def test_dict_chunks_with_finish_reason_are_trusted(self):
|
||||
chunks = [
|
||||
{
|
||||
"_hidden_params": {"custom_llm_provider": "anthropic"},
|
||||
"choices": [{"delta": {"content": "Yes, "}, "finish_reason": None}],
|
||||
},
|
||||
{
|
||||
"_hidden_params": {"custom_llm_provider": "anthropic"},
|
||||
"choices": [{"delta": {"content": "that works."}, "finish_reason": "stop"}],
|
||||
"usage": Usage(prompt_tokens=20, completion_tokens=5, total_tokens=25),
|
||||
},
|
||||
]
|
||||
processor = ChunkProcessor(chunks=chunks, messages=[])
|
||||
result = processor._calculate_usage_per_chunk(chunks=chunks)
|
||||
assert result["completion_tokens"] == 5
|
||||
|
||||
def test_dict_chunks_without_finish_reason_reset_placeholder(self):
|
||||
chunks = [
|
||||
{
|
||||
"_hidden_params": {"custom_llm_provider": "anthropic"},
|
||||
"choices": [],
|
||||
"usage": Usage(prompt_tokens=20, completion_tokens=1, total_tokens=21),
|
||||
},
|
||||
{
|
||||
"_hidden_params": {"custom_llm_provider": "anthropic"},
|
||||
"choices": [{"delta": {"content": "partial"}, "finish_reason": None}],
|
||||
},
|
||||
]
|
||||
processor = ChunkProcessor(chunks=chunks, messages=[])
|
||||
result = processor._calculate_usage_per_chunk(chunks=chunks)
|
||||
assert result["completion_tokens"] == 0
|
||||
assert result["completion_tokens_details"] is None
|
||||
|
||||
def test_estimated_reasoning_is_capped_to_trusted_completion_total(self):
|
||||
chunks = [
|
||||
_make_chunk(reasoning_content="Let me reason about this carefully and at length. " * 20),
|
||||
_make_chunk(
|
||||
finish_reason="stop",
|
||||
usage=Usage(prompt_tokens=20, completion_tokens=5, total_tokens=25),
|
||||
),
|
||||
]
|
||||
response = litellm.stream_chunk_builder(
|
||||
chunks=chunks,
|
||||
messages=[{"role": "user", "content": "Go."}],
|
||||
)
|
||||
details = response.usage.completion_tokens_details
|
||||
assert response.usage.completion_tokens == 5
|
||||
assert details.reasoning_tokens <= response.usage.completion_tokens
|
||||
assert details.reasoning_tokens + details.text_tokens == response.usage.completion_tokens
|
||||
assert details.text_tokens >= 0
|
||||
|
||||
|
||||
class TestProviderGuard:
|
||||
"""Class A: the cursor-reset heuristic must NOT silently affect non-Anthropic
|
||||
|
|
@ -297,11 +391,12 @@ class TestNonAnthropicStreamingIntact:
|
|||
"""Make sure providers without cursor pattern still work."""
|
||||
|
||||
def test_completion_tokens_above_one_never_resets(self):
|
||||
"""Any chunk reporting completion_tokens > 1 sets saw_non_cursor
|
||||
and prevents the reset."""
|
||||
"""A non-Anthropic provider reporting completion_tokens > 1 from a
|
||||
single usage event keeps that value."""
|
||||
chunks = [
|
||||
_make_chunk(
|
||||
usage=Usage(prompt_tokens=10, completion_tokens=5, total_tokens=15)
|
||||
usage=Usage(prompt_tokens=10, completion_tokens=5, total_tokens=15),
|
||||
custom_llm_provider="openai",
|
||||
),
|
||||
]
|
||||
processor = ChunkProcessor(chunks=chunks, messages=[])
|
||||
|
|
|
|||
|
|
@ -677,3 +677,14 @@ def test_azure_responses_gpt6_astra_rejects_temperature_while_reasoning(local_mo
|
|||
model="gpt-6-astra",
|
||||
drop_params=False,
|
||||
)
|
||||
|
||||
|
||||
def test_azure_responses_sends_the_deployment_name_when_azure_ai_prefix_survives_provider_remap():
|
||||
request = AzureOpenAIResponsesAPIConfig().transform_responses_api_request(
|
||||
model="azure_ai/gpt-5.4-nano",
|
||||
input="hi",
|
||||
response_api_optional_request_params={},
|
||||
litellm_params=GenericLiteLLMParams(),
|
||||
headers={},
|
||||
)
|
||||
assert request["model"] == "gpt-5.4-nano"
|
||||
|
|
|
|||
|
|
@ -0,0 +1,311 @@
|
|||
import json
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
import respx
|
||||
|
||||
import litellm
|
||||
from litellm.llms.azure_ai.responses.transformation import AzureAIResponsesAPIConfig
|
||||
from litellm.responses.main import _will_bridge_to_chat_completions
|
||||
from litellm.types.router import GenericLiteLLMParams
|
||||
from litellm.utils import ProviderConfigManager
|
||||
|
||||
FOUNDRY_PROJECT_BASE = "https://res.services.ai.azure.com/api/projects/proj"
|
||||
FOUNDRY_RESPONSES_URL = f"{FOUNDRY_PROJECT_BASE}/openai/v1/responses"
|
||||
SERVERLESS_BASE = "https://endpoint.eastus.models.ai.azure.com"
|
||||
WEATHER_TOOL = {
|
||||
"type": "function",
|
||||
"name": "get_weather",
|
||||
"description": "Get weather",
|
||||
"parameters": {"type": "object", "properties": {"city": {"type": "string"}}, "required": ["city"]},
|
||||
}
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def clear_azure_ai_env(monkeypatch):
|
||||
monkeypatch.setattr(litellm, "api_base", None)
|
||||
monkeypatch.setattr(litellm, "api_key", None)
|
||||
monkeypatch.setattr(litellm, "enable_azure_ad_token_refresh", False)
|
||||
monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
|
||||
for env_var in (
|
||||
"AZURE_AI_API_BASE",
|
||||
"AZURE_AI_API_KEY",
|
||||
"AZURE_AD_TOKEN",
|
||||
"AZURE_TENANT_ID",
|
||||
"AZURE_CLIENT_ID",
|
||||
"AZURE_CLIENT_SECRET",
|
||||
):
|
||||
monkeypatch.delenv(env_var, raising=False)
|
||||
|
||||
|
||||
def _responses_payload(model: str) -> dict:
|
||||
return {
|
||||
"id": "resp_123",
|
||||
"object": "response",
|
||||
"created_at": 1741369938,
|
||||
"status": "completed",
|
||||
"model": model,
|
||||
"output": [],
|
||||
"parallel_tool_calls": False,
|
||||
"usage": {"input_tokens": 1, "output_tokens": 1, "total_tokens": 2},
|
||||
"error": None,
|
||||
"tool_choice": "auto",
|
||||
"tools": [],
|
||||
"metadata": None,
|
||||
"temperature": None,
|
||||
"top_p": None,
|
||||
"max_output_tokens": None,
|
||||
"previous_response_id": None,
|
||||
"reasoning": None,
|
||||
"truncation": None,
|
||||
"instructions": None,
|
||||
"incomplete_details": None,
|
||||
"user": None,
|
||||
}
|
||||
|
||||
|
||||
def _chat_completion_payload(model: str) -> dict:
|
||||
return {
|
||||
"id": "chatcmpl-123",
|
||||
"object": "chat.completion",
|
||||
"created": 1741369938,
|
||||
"model": model,
|
||||
"choices": [{"index": 0, "message": {"role": "assistant", "content": "hi"}, "finish_reason": "stop"}],
|
||||
"usage": {"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2},
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.parametrize("model", ["gpt-5.6-luna-20260710154139", "gpt-5.5-20260504143601", "DeepSeek-R1-0528", None])
|
||||
@pytest.mark.parametrize(
|
||||
"api_base", [FOUNDRY_PROJECT_BASE, "https://res.services.ai.azure.com", "https://res.openai.azure.com"]
|
||||
)
|
||||
def test_azure_openai_v1_hosts_resolve_native_config(model, api_base):
|
||||
config = ProviderConfigManager.get_provider_responses_api_config(provider="azure_ai", model=model, api_base=api_base)
|
||||
assert isinstance(config, AzureAIResponsesAPIConfig)
|
||||
|
||||
|
||||
def test_api_base_from_env_resolves_native_config(monkeypatch):
|
||||
monkeypatch.setenv("AZURE_AI_API_BASE", FOUNDRY_PROJECT_BASE)
|
||||
config = ProviderConfigManager.get_provider_responses_api_config(provider="azure_ai", model="gpt-5.6-luna", api_base=None)
|
||||
assert isinstance(config, AzureAIResponsesAPIConfig)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("model", ["gpt-5.6-luna", None])
|
||||
@pytest.mark.parametrize(
|
||||
"api_base",
|
||||
[SERVERLESS_BASE, "https://endpoint.eastus.inference.ml.azure.com/score", "https://res.cognitiveservices.azure.com"],
|
||||
)
|
||||
def test_other_hosts_keep_chat_bridge(model, api_base):
|
||||
config = ProviderConfigManager.get_provider_responses_api_config(provider="azure_ai", model=model, api_base=api_base)
|
||||
assert config is None
|
||||
|
||||
|
||||
@pytest.mark.parametrize("model", ["claude-3-5-sonnet", "model_router/gpt-5", "agents/my-agent"])
|
||||
def test_non_openai_surfaces_keep_chat_bridge(model):
|
||||
config = ProviderConfigManager.get_provider_responses_api_config(
|
||||
provider="azure_ai", model=model, api_base=FOUNDRY_PROJECT_BASE
|
||||
)
|
||||
assert config is None
|
||||
|
||||
|
||||
@pytest.mark.parametrize("api_base,bridged", [(FOUNDRY_PROJECT_BASE, False), (SERVERLESS_BASE, True)])
|
||||
def test_will_bridge_to_chat_completions_follows_host(api_base, bridged):
|
||||
assert _will_bridge_to_chat_completions("gpt-5.6-luna", "azure_ai", False, None, api_base) is bridged
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"api_base,expected",
|
||||
[
|
||||
(FOUNDRY_PROJECT_BASE, FOUNDRY_RESPONSES_URL),
|
||||
(f"{FOUNDRY_PROJECT_BASE}/", FOUNDRY_RESPONSES_URL),
|
||||
(f"{FOUNDRY_PROJECT_BASE}/openai/v1", FOUNDRY_RESPONSES_URL),
|
||||
(FOUNDRY_RESPONSES_URL, FOUNDRY_RESPONSES_URL),
|
||||
("https://res.services.ai.azure.com", "https://res.services.ai.azure.com/openai/v1/responses"),
|
||||
("https://res.services.ai.azure.com/models", "https://res.services.ai.azure.com/openai/v1/responses"),
|
||||
(
|
||||
"https://res.services.ai.azure.com/models/chat/completions?api-version=2024-05-01-preview",
|
||||
"https://res.services.ai.azure.com/openai/v1/responses",
|
||||
),
|
||||
("https://res.openai.azure.com", "https://res.openai.azure.com/openai/v1/responses"),
|
||||
(
|
||||
"https://res.openai.azure.com/openai/deployments/gpt-5?api-version=2025-04-01-preview",
|
||||
"https://res.openai.azure.com/openai/v1/responses",
|
||||
),
|
||||
],
|
||||
)
|
||||
def test_get_complete_url(api_base, expected):
|
||||
assert AzureAIResponsesAPIConfig().get_complete_url(api_base=api_base, litellm_params={}) == expected
|
||||
|
||||
|
||||
def test_get_complete_url_ignores_api_version():
|
||||
url = AzureAIResponsesAPIConfig().get_complete_url(
|
||||
api_base=FOUNDRY_PROJECT_BASE, litellm_params={"api_version": "2025-04-01-preview"}
|
||||
)
|
||||
assert url == FOUNDRY_RESPONSES_URL
|
||||
|
||||
|
||||
def test_get_complete_url_uses_env_api_base(monkeypatch):
|
||||
monkeypatch.setenv("AZURE_AI_API_BASE", FOUNDRY_PROJECT_BASE)
|
||||
assert AzureAIResponsesAPIConfig().get_complete_url(api_base=None, litellm_params={}) == FOUNDRY_RESPONSES_URL
|
||||
|
||||
|
||||
def test_get_complete_url_raises_without_api_base():
|
||||
with pytest.raises(ValueError, match="AZURE_AI_API_BASE"):
|
||||
AzureAIResponsesAPIConfig().get_complete_url(api_base=None, litellm_params={})
|
||||
|
||||
|
||||
def test_native_websocket_stays_off():
|
||||
assert AzureAIResponsesAPIConfig().supports_native_websocket() is False
|
||||
|
||||
|
||||
def test_validate_environment_sends_api_key_header():
|
||||
headers = AzureAIResponsesAPIConfig().validate_environment(
|
||||
headers={"x-custom": "1"},
|
||||
model="gpt-5.6-luna",
|
||||
litellm_params=GenericLiteLLMParams(api_key="secret", api_base=FOUNDRY_PROJECT_BASE),
|
||||
)
|
||||
assert headers == {"x-custom": "1", "api-key": "secret", "Content-Type": "application/json"}
|
||||
|
||||
|
||||
def test_validate_environment_reads_api_key_from_env(monkeypatch):
|
||||
monkeypatch.setenv("AZURE_AI_API_KEY", "env-secret")
|
||||
headers = AzureAIResponsesAPIConfig().validate_environment(
|
||||
headers={}, model="gpt-5.6-luna", litellm_params=GenericLiteLLMParams(api_base=FOUNDRY_PROJECT_BASE)
|
||||
)
|
||||
assert headers["api-key"] == "env-secret"
|
||||
|
||||
|
||||
def test_validate_environment_uses_entra_token_without_api_key():
|
||||
headers = AzureAIResponsesAPIConfig().validate_environment(
|
||||
headers={},
|
||||
model="gpt-5.6-luna",
|
||||
litellm_params=GenericLiteLLMParams(azure_ad_token="entra-token", api_base=FOUNDRY_PROJECT_BASE),
|
||||
)
|
||||
assert headers["Authorization"] == "Bearer entra-token"
|
||||
assert "api-key" not in headers
|
||||
|
||||
|
||||
def test_validate_environment_raises_without_credentials():
|
||||
with pytest.raises(ValueError, match="AZURE_AI_API_KEY"):
|
||||
AzureAIResponsesAPIConfig().validate_environment(
|
||||
headers={}, model="gpt-5.6-luna", litellm_params=GenericLiteLLMParams(api_base=FOUNDRY_PROJECT_BASE)
|
||||
)
|
||||
|
||||
|
||||
NATIVE_RESPONSES_CASES = [
|
||||
("azure_ai/gpt-5.6-luna-20260710154139", FOUNDRY_PROJECT_BASE, FOUNDRY_RESPONSES_URL, "gpt-5.6-luna-20260710154139"),
|
||||
(
|
||||
"azure_ai/gpt-5.6-luna",
|
||||
"https://res.services.ai.azure.com/models",
|
||||
"https://res.services.ai.azure.com/openai/v1/responses",
|
||||
"gpt-5.6-luna",
|
||||
),
|
||||
(
|
||||
"azure_ai/gpt-5.6-sol",
|
||||
"https://res.services.ai.azure.com",
|
||||
"https://res.services.ai.azure.com/openai/v1/responses",
|
||||
"gpt-5.6-sol",
|
||||
),
|
||||
(
|
||||
"azure_ai/gpt-5.6-luna-20260710154139",
|
||||
"https://res.openai.azure.com",
|
||||
"https://res.openai.azure.com/openai/v1/responses",
|
||||
"gpt-5.6-luna-20260710154139",
|
||||
),
|
||||
(
|
||||
"azure_ai/gpt-5.6-sol",
|
||||
"https://res.openai.azure.com",
|
||||
"https://res.openai.azure.com/openai/v1/responses",
|
||||
"gpt-5.6-sol",
|
||||
),
|
||||
]
|
||||
|
||||
|
||||
def _assert_native_responses_request(route, expected_url, expected_model):
|
||||
request = route.calls.last.request
|
||||
body = json.loads(request.content)
|
||||
assert f"{request.url.scheme}://{request.url.host}{request.url.path}" == expected_url
|
||||
assert request.headers["api-key"] == "fake-key"
|
||||
assert body["model"] == expected_model
|
||||
assert body["input"] == "What is the weather in SF?"
|
||||
assert "messages" not in body
|
||||
assert body["reasoning"] == {"effort": "high"}
|
||||
assert body["tools"] == [WEATHER_TOOL]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@respx.mock
|
||||
@pytest.mark.parametrize("model,api_base,expected_url,expected_model", NATIVE_RESPONSES_CASES)
|
||||
async def test_aresponses_sends_reasoning_and_tools_to_native_endpoint(model, api_base, expected_url, expected_model):
|
||||
route = respx.post(url__regex=r".*/openai/v1/responses(\?.*)?$").mock(
|
||||
return_value=httpx.Response(200, json=_responses_payload(expected_model))
|
||||
)
|
||||
|
||||
await litellm.aresponses(
|
||||
model=model,
|
||||
input="What is the weather in SF?",
|
||||
reasoning_effort="high",
|
||||
tools=[WEATHER_TOOL],
|
||||
api_base=api_base,
|
||||
api_key="fake-key",
|
||||
)
|
||||
|
||||
_assert_native_responses_request(route, expected_url, expected_model)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@respx.mock
|
||||
async def test_aresponses_catalog_name_remapped_to_azure_sends_bare_deployment_name(monkeypatch):
|
||||
monkeypatch.setenv("AZURE_AI_API_BASE", "https://res.openai.azure.com")
|
||||
route = respx.post(url__regex=r".*/openai/v1/responses(\?.*)?$").mock(
|
||||
return_value=httpx.Response(200, json=_responses_payload("gpt-5.4-nano"))
|
||||
)
|
||||
|
||||
await litellm.aresponses(
|
||||
model="azure_ai/gpt-5.4-nano",
|
||||
input="What is the weather in SF?",
|
||||
api_base="https://res.openai.azure.com",
|
||||
api_key="fake-key",
|
||||
)
|
||||
|
||||
assert json.loads(route.calls.last.request.content)["model"] == "gpt-5.4-nano"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@respx.mock
|
||||
@pytest.mark.parametrize("model,api_base,expected_url,expected_model", NATIVE_RESPONSES_CASES)
|
||||
async def test_router_aresponses_sends_bare_deployment_name(model, api_base, expected_url, expected_model):
|
||||
route = respx.post(url__regex=r".*/openai/v1/responses(\?.*)?$").mock(
|
||||
return_value=httpx.Response(200, json=_responses_payload(expected_model))
|
||||
)
|
||||
router = litellm.Router(
|
||||
model_list=[{"model_name": "gpt-5.6", "litellm_params": {"model": model, "api_base": api_base, "api_key": "fake-key"}}],
|
||||
num_retries=0,
|
||||
)
|
||||
|
||||
await router.aresponses(
|
||||
model="gpt-5.6", input="What is the weather in SF?", reasoning={"effort": "high"}, tools=[WEATHER_TOOL]
|
||||
)
|
||||
|
||||
_assert_native_responses_request(route, expected_url, expected_model)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@respx.mock
|
||||
async def test_aresponses_serverless_host_stays_on_chat_bridge():
|
||||
chat_route = respx.post(url__regex=r".*/chat/completions$").mock(
|
||||
return_value=httpx.Response(200, json=_chat_completion_payload("gpt-5.6-luna"))
|
||||
)
|
||||
responses_route = respx.post(url__regex=r".*/responses$")
|
||||
|
||||
await litellm.aresponses(
|
||||
model="azure_ai/gpt-5.6-luna-20260710154139",
|
||||
input="What is the weather in SF?",
|
||||
tools=[WEATHER_TOOL],
|
||||
api_base=SERVERLESS_BASE,
|
||||
api_key="fake-key",
|
||||
)
|
||||
|
||||
assert chat_route.called
|
||||
assert not responses_route.called
|
||||
assert chat_route.calls.last.request.headers["Authorization"] == "Bearer fake-key"
|
||||
|
|
@ -6534,6 +6534,446 @@ async def test_grounding_source_and_query_rendered_as_text():
|
|||
assert {"text": "What is the capital of Japan?"} in user_content
|
||||
|
||||
|
||||
def _orphaned_tool_history_messages():
|
||||
return [
|
||||
{"role": "user", "content": "What's the weather in Paris?"},
|
||||
{
|
||||
"role": "assistant",
|
||||
"content": None,
|
||||
"tool_calls": [
|
||||
{
|
||||
"id": "call_abc",
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": "get_weather",
|
||||
"arguments": '{"city": "Paris"}',
|
||||
},
|
||||
}
|
||||
],
|
||||
},
|
||||
{
|
||||
"role": "tool",
|
||||
"tool_call_id": "call_abc",
|
||||
"content": "Sunny, 25C",
|
||||
},
|
||||
{"role": "user", "content": "Summarize our conversation so far."},
|
||||
]
|
||||
|
||||
|
||||
def test_neutralize_orphaned_tool_blocks_rewrites_when_no_tools():
|
||||
"""No tools= but history has tool blocks: assistant tool_calls and the tool
|
||||
result must be rewritten to text, with the structured tool fields gone and
|
||||
tool_call_id preserved, so Bedrock accepts the request without a toolConfig
|
||||
(#24158, #27138)."""
|
||||
messages = _orphaned_tool_history_messages()
|
||||
|
||||
result = AmazonConverseConfig._neutralize_orphaned_tool_blocks(
|
||||
messages, optional_params={}
|
||||
)
|
||||
|
||||
serialized = json.dumps(result)
|
||||
assert "tool_calls" not in serialized
|
||||
assert not any(m.get("role") in ("tool", "function") for m in result)
|
||||
assert "get_weather" in serialized
|
||||
# The arguments string contains quotes; after json.dumps the literal
|
||||
# '{"city": "Paris"}' is escaped, so assert on quote-free tokens that survive.
|
||||
assert "city" in serialized and "Paris" in serialized
|
||||
assert "Sunny, 25C" in serialized
|
||||
assert "[tool call call_abc: get_weather(" in result[1]["content"]
|
||||
assert "[tool result for call_abc: Sunny, 25C]" in result[2]["content"]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("tools_value", [[], None])
|
||||
def test_neutralize_orphaned_tool_blocks_rewrites_when_tools_empty(tools_value):
|
||||
"""tools=[] and tools=None are 'no usable tools'; the gate must be on
|
||||
truthiness, not key presence, or these slip through and still emit
|
||||
structured tool blocks with no toolConfig."""
|
||||
messages = _orphaned_tool_history_messages()
|
||||
|
||||
result = AmazonConverseConfig._neutralize_orphaned_tool_blocks(
|
||||
messages, optional_params={"tools": tools_value}
|
||||
)
|
||||
|
||||
serialized = json.dumps(result)
|
||||
assert "tool_calls" not in serialized
|
||||
assert "get_weather" in serialized
|
||||
|
||||
|
||||
def test_neutralize_orphaned_tool_blocks_rewrites_tool_result_only_history():
|
||||
"""A role:"tool"-only history (no assistant tool_calls) must also be
|
||||
neutralized; has_tool_call_blocks misses this, but the factory still emits a
|
||||
lone toolResult with no toolConfig."""
|
||||
messages = [
|
||||
{"role": "user", "content": "hi"},
|
||||
{"role": "tool", "tool_call_id": "call_xyz", "content": "lookup result"},
|
||||
]
|
||||
|
||||
result = AmazonConverseConfig._neutralize_orphaned_tool_blocks(
|
||||
messages, optional_params={}
|
||||
)
|
||||
|
||||
assert not any(m.get("role") in ("tool", "function") for m in result)
|
||||
serialized = json.dumps(result)
|
||||
assert "lookup result" in serialized
|
||||
assert "call_xyz" in serialized
|
||||
|
||||
|
||||
def test_neutralize_orphaned_tool_blocks_non_text_result_marked_not_empty():
|
||||
"""Non-text tool-result payloads (image/file) collapse to an explicit
|
||||
marker, never an empty string (Bedrock rejects empty text blocks) and never
|
||||
a silent drop."""
|
||||
messages = [
|
||||
{"role": "user", "content": "hi"},
|
||||
{
|
||||
"role": "assistant",
|
||||
"content": None,
|
||||
"tool_calls": [
|
||||
{
|
||||
"id": "c1",
|
||||
"type": "function",
|
||||
"function": {"name": "render", "arguments": "{}"},
|
||||
}
|
||||
],
|
||||
},
|
||||
{
|
||||
"role": "tool",
|
||||
"tool_call_id": "c1",
|
||||
"content": [
|
||||
{
|
||||
"type": "image_url",
|
||||
"image_url": {"url": "data:image/png;base64,AAAA"},
|
||||
}
|
||||
],
|
||||
},
|
||||
]
|
||||
|
||||
result = AmazonConverseConfig._neutralize_orphaned_tool_blocks(
|
||||
messages, optional_params={}
|
||||
)
|
||||
|
||||
rewritten = next(
|
||||
m for m in result if m.get("role") == "user" and m is not messages[0]
|
||||
)
|
||||
text = rewritten["content"]
|
||||
assert text.strip() # never empty
|
||||
assert "non-text tool result omitted" in text
|
||||
|
||||
|
||||
def test_neutralize_orphaned_tool_blocks_noop_when_tools_present():
|
||||
"""When a non-empty tools= is provided, tool blocks are legitimate and must
|
||||
be left untouched (returns the same object, no rewriting)."""
|
||||
messages = _orphaned_tool_history_messages()
|
||||
|
||||
result = AmazonConverseConfig._neutralize_orphaned_tool_blocks(
|
||||
messages,
|
||||
optional_params={"tools": [{"type": "function", "function": {"name": "x"}}]},
|
||||
)
|
||||
|
||||
assert result is messages
|
||||
|
||||
|
||||
def test_neutralize_orphaned_tool_blocks_noop_when_no_tool_history():
|
||||
"""Plain conversation with no tool blocks is returned unchanged."""
|
||||
messages = [{"role": "user", "content": "hi"}]
|
||||
|
||||
result = AmazonConverseConfig._neutralize_orphaned_tool_blocks(
|
||||
messages, optional_params={}
|
||||
)
|
||||
|
||||
assert result is messages
|
||||
|
||||
|
||||
def test_neutralize_orphaned_tool_blocks_logs_warning(caplog):
|
||||
"""Neutralization must surface at WARNING level so a developer who forgot
|
||||
tools= sees it instead of a silent degrade."""
|
||||
messages = _orphaned_tool_history_messages()
|
||||
|
||||
with caplog.at_level("WARNING"):
|
||||
AmazonConverseConfig._neutralize_orphaned_tool_blocks(
|
||||
messages, optional_params={}
|
||||
)
|
||||
|
||||
assert any(
|
||||
"neutralizing orphaned tool blocks" in record.getMessage()
|
||||
for record in caplog.records
|
||||
)
|
||||
|
||||
|
||||
def _assert_no_structured_tool_blocks(result):
|
||||
"""A valid Bedrock body for a neutralized request has no tool config AND no
|
||||
structured tool blocks in messages. Checking only toolConfig is insufficient:
|
||||
deleting the raise without rewriting still leaves toolUse/toolResult, the
|
||||
exact shape Bedrock rejects."""
|
||||
assert "toolConfig" not in result
|
||||
serialized = json.dumps(result)
|
||||
assert "toolUse" not in serialized
|
||||
assert "toolResult" not in serialized
|
||||
|
||||
|
||||
def test_transform_request_no_tools_with_tool_history_succeeds_24158(monkeypatch):
|
||||
"""#24158: a compaction-style call (tool blocks in history, no tools=) must
|
||||
not raise and must send no toolConfig or structured tool blocks, on
|
||||
default settings."""
|
||||
monkeypatch.setattr(litellm, "modify_params", False)
|
||||
config = AmazonConverseConfig()
|
||||
|
||||
result = config.transform_request(
|
||||
model="us.anthropic.claude-opus-4-5-20251101-v1:0",
|
||||
messages=_orphaned_tool_history_messages(),
|
||||
optional_params={},
|
||||
litellm_params={},
|
||||
headers={},
|
||||
)
|
||||
|
||||
_assert_no_structured_tool_blocks(result)
|
||||
serialized = json.dumps(result)
|
||||
assert "get_weather" in serialized
|
||||
assert "Sunny, 25C" in serialized
|
||||
|
||||
|
||||
def test_transform_request_tool_unsupported_model_no_toolconfig_27138(monkeypatch):
|
||||
"""#27138: a tool-incapable model with tool blocks in history and no tools=
|
||||
must not get a toolConfig/toolUse/toolResult injected (which Bedrock would
|
||||
400 on), even with modify_params on."""
|
||||
monkeypatch.setattr(litellm, "modify_params", True)
|
||||
config = AmazonConverseConfig()
|
||||
|
||||
result = config.transform_request(
|
||||
model="meta.llama3-2-3b-instruct-v1:0",
|
||||
messages=_orphaned_tool_history_messages(),
|
||||
optional_params={},
|
||||
litellm_params={},
|
||||
headers={},
|
||||
)
|
||||
|
||||
_assert_no_structured_tool_blocks(result)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("tools_value", [[], None])
|
||||
def test_transform_request_empty_tools_with_tool_history(monkeypatch, tools_value):
|
||||
"""tools=[] / tools=None must be neutralized like no tools at all; a
|
||||
key-presence gate would skip them and emit toolUse/toolResult with no
|
||||
toolConfig."""
|
||||
monkeypatch.setattr(litellm, "modify_params", False)
|
||||
config = AmazonConverseConfig()
|
||||
|
||||
result = config.transform_request(
|
||||
model="us.anthropic.claude-opus-4-5-20251101-v1:0",
|
||||
messages=_orphaned_tool_history_messages(),
|
||||
optional_params={"tools": tools_value},
|
||||
litellm_params={},
|
||||
headers={},
|
||||
)
|
||||
|
||||
_assert_no_structured_tool_blocks(result)
|
||||
|
||||
|
||||
def test_transform_request_tool_result_only_history(monkeypatch):
|
||||
"""A role:"tool"-only history (no assistant tool_calls) currently emits a
|
||||
lone toolResult with no toolConfig; it must be neutralized."""
|
||||
monkeypatch.setattr(litellm, "modify_params", False)
|
||||
config = AmazonConverseConfig()
|
||||
|
||||
result = config.transform_request(
|
||||
model="us.anthropic.claude-opus-4-5-20251101-v1:0",
|
||||
messages=[
|
||||
{"role": "user", "content": "hi"},
|
||||
{"role": "tool", "tool_call_id": "call_xyz", "content": "lookup result"},
|
||||
],
|
||||
optional_params={},
|
||||
litellm_params={},
|
||||
headers={},
|
||||
)
|
||||
|
||||
_assert_no_structured_tool_blocks(result)
|
||||
assert "lookup result" in json.dumps(result)
|
||||
|
||||
|
||||
def test_transform_request_neutralized_tool_output_is_guarded(monkeypatch):
|
||||
"""With guardrailConfig present, a neutralized tool result that becomes the
|
||||
trailing user turn must be emitted as guardContent, not plain text, so
|
||||
untrusted tool output does not bypass the guardrail (neutralize must run
|
||||
before guarded-text conversion)."""
|
||||
monkeypatch.setattr(litellm, "modify_params", False)
|
||||
config = AmazonConverseConfig()
|
||||
|
||||
result = config.transform_request(
|
||||
model="us.anthropic.claude-opus-4-5-20251101-v1:0",
|
||||
messages=[
|
||||
{"role": "user", "content": "look it up"},
|
||||
{
|
||||
"role": "assistant",
|
||||
"content": None,
|
||||
"tool_calls": [
|
||||
{
|
||||
"id": "c1",
|
||||
"type": "function",
|
||||
"function": {"name": "lookup", "arguments": "{}"},
|
||||
}
|
||||
],
|
||||
},
|
||||
{"role": "tool", "tool_call_id": "c1", "content": "secret tool output"},
|
||||
],
|
||||
optional_params={
|
||||
"guardrailConfig": {"guardrailIdentifier": "gid", "guardrailVersion": "1"}
|
||||
},
|
||||
litellm_params={},
|
||||
headers={},
|
||||
)
|
||||
|
||||
_assert_no_structured_tool_blocks(result)
|
||||
serialized = json.dumps(result)
|
||||
assert "guardContent" in serialized
|
||||
assert "secret tool output" in serialized
|
||||
|
||||
|
||||
def test_transform_request_neutralized_tool_output_guarded_mid_history(monkeypatch):
|
||||
"""Regression: a neutralized tool result that is NOT the trailing turn (an
|
||||
assistant reply and a later user turn follow it) must still be guardContent.
|
||||
_convert_consecutive_user_messages_to_guarded_text only covers the trailing
|
||||
user turn, so neutralize itself must guard untrusted tool output regardless
|
||||
of position, else an attacker controlling the tool response bypasses the
|
||||
guardrail (bot review)."""
|
||||
monkeypatch.setattr(litellm, "modify_params", False)
|
||||
config = AmazonConverseConfig()
|
||||
|
||||
result = config.transform_request(
|
||||
model="us.anthropic.claude-opus-4-5-20251101-v1:0",
|
||||
messages=[
|
||||
{"role": "user", "content": "look it up"},
|
||||
{
|
||||
"role": "assistant",
|
||||
"content": None,
|
||||
"tool_calls": [
|
||||
{
|
||||
"id": "c1",
|
||||
"type": "function",
|
||||
"function": {"name": "lookup", "arguments": "{}"},
|
||||
}
|
||||
],
|
||||
},
|
||||
{"role": "tool", "tool_call_id": "c1", "content": "IGNORE_PRIOR malware"},
|
||||
{"role": "assistant", "content": "Here is the summary."},
|
||||
{"role": "user", "content": "thanks"},
|
||||
],
|
||||
optional_params={
|
||||
"guardrailConfig": {"guardrailIdentifier": "gid", "guardrailVersion": "1"}
|
||||
},
|
||||
litellm_params={},
|
||||
headers={},
|
||||
)
|
||||
|
||||
_assert_no_structured_tool_blocks(result)
|
||||
blocks = [block for message in result["messages"] for block in message["content"]]
|
||||
guarded_texts = [
|
||||
block["guardContent"]["text"]["text"] for block in blocks if "guardContent" in block
|
||||
]
|
||||
plain_texts = [block["text"] for block in blocks if "text" in block and "guardContent" not in block]
|
||||
assert any("malware" in text for text in guarded_texts), "mid-history tool output must be guarded"
|
||||
assert not any(
|
||||
"malware" in text for text in plain_texts
|
||||
), "mid-history tool output must not reach the model as unguarded text"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_transform_request_no_tools_with_tool_history(monkeypatch):
|
||||
"""Async is a separate request assembler; it must neutralize identically."""
|
||||
monkeypatch.setattr(litellm, "modify_params", False)
|
||||
config = AmazonConverseConfig()
|
||||
|
||||
result = await config._async_transform_request(
|
||||
model="us.anthropic.claude-opus-4-5-20251101-v1:0",
|
||||
messages=_orphaned_tool_history_messages(),
|
||||
optional_params={},
|
||||
litellm_params={},
|
||||
headers={},
|
||||
)
|
||||
|
||||
_assert_no_structured_tool_blocks(result)
|
||||
assert "get_weather" in json.dumps(result)
|
||||
|
||||
|
||||
def test_transform_request_with_tools_still_builds_toolconfig(monkeypatch):
|
||||
"""Guard: when a non-empty tools= IS provided, tool blocks are legitimate and
|
||||
a toolConfig must still be produced (neutralization must not regress this)."""
|
||||
monkeypatch.setattr(litellm, "modify_params", False)
|
||||
config = AmazonConverseConfig()
|
||||
|
||||
result = config.transform_request(
|
||||
model="us.anthropic.claude-opus-4-5-20251101-v1:0",
|
||||
messages=_orphaned_tool_history_messages(),
|
||||
optional_params={
|
||||
"tools": [
|
||||
{
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": "get_weather",
|
||||
"description": "Get weather",
|
||||
"parameters": {"type": "object", "properties": {}},
|
||||
},
|
||||
}
|
||||
]
|
||||
},
|
||||
litellm_params={},
|
||||
headers={},
|
||||
)
|
||||
|
||||
assert "toolConfig" in result
|
||||
|
||||
|
||||
def test_transform_request_flag_off_restores_raise(monkeypatch):
|
||||
"""Opt-out: with bedrock_neutralize_orphaned_tool_blocks=False and
|
||||
modify_params=False, the legacy UnsupportedParamsError contract is restored."""
|
||||
monkeypatch.setattr(litellm, "bedrock_neutralize_orphaned_tool_blocks", False)
|
||||
monkeypatch.setattr(litellm, "modify_params", False)
|
||||
config = AmazonConverseConfig()
|
||||
|
||||
with pytest.raises(litellm.utils.UnsupportedParamsError, match="without `tools="):
|
||||
config.transform_request(
|
||||
model="us.anthropic.claude-opus-4-5-20251101-v1:0",
|
||||
messages=_orphaned_tool_history_messages(),
|
||||
optional_params={},
|
||||
litellm_params={},
|
||||
headers={},
|
||||
)
|
||||
|
||||
|
||||
def test_transform_request_flag_off_with_modify_params_restores_dummy_tool(monkeypatch):
|
||||
"""Opt-out: with the flag off and modify_params=True, the legacy dummy-tool
|
||||
injection is restored (a toolConfig is produced, not neutralized text)."""
|
||||
monkeypatch.setattr(litellm, "bedrock_neutralize_orphaned_tool_blocks", False)
|
||||
monkeypatch.setattr(litellm, "modify_params", True)
|
||||
config = AmazonConverseConfig()
|
||||
|
||||
result = config.transform_request(
|
||||
model="us.anthropic.claude-opus-4-5-20251101-v1:0",
|
||||
messages=_orphaned_tool_history_messages(),
|
||||
optional_params={},
|
||||
litellm_params={},
|
||||
headers={},
|
||||
)
|
||||
|
||||
assert "toolConfig" in result
|
||||
assert "dummy_tool" in json.dumps(result)
|
||||
|
||||
|
||||
def test_transform_request_flag_on_is_default(monkeypatch):
|
||||
"""Default-on: without touching the flag, neutralization is the behavior."""
|
||||
monkeypatch.setattr(litellm, "modify_params", False)
|
||||
config = AmazonConverseConfig()
|
||||
|
||||
assert litellm.bedrock_neutralize_orphaned_tool_blocks is True
|
||||
result = config.transform_request(
|
||||
model="us.anthropic.claude-opus-4-5-20251101-v1:0",
|
||||
messages=_orphaned_tool_history_messages(),
|
||||
optional_params={},
|
||||
litellm_params={},
|
||||
headers={},
|
||||
)
|
||||
|
||||
_assert_no_structured_tool_blocks(result)
|
||||
|
||||
|
||||
def _agentic_messages_with_ttl(ttl_target: str):
|
||||
"""A tool-loop conversation with `ttl: 1h` cache_control at `ttl_target`:
|
||||
'user', 'tool_call' (per-tool-call, on the assistant's tool call), or
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -8,7 +8,7 @@ from litellm.litellm_core_utils.get_model_cost_map import get_model_cost_map
|
|||
from litellm.llms.openai.chat.gpt_5_transformation import OpenAIGPT5Config
|
||||
from litellm.llms.openai.openai import OpenAIConfig
|
||||
from litellm.utils import (
|
||||
_is_explicitly_disabled_factory,
|
||||
is_explicitly_disabled_factory,
|
||||
peek_reasoning_summary_aliases,
|
||||
strip_reasoning_summary_aliases_from_optional_params,
|
||||
)
|
||||
|
|
@ -524,19 +524,19 @@ def test_gpt5_minimal_explicitly_disabled_check(gpt5_config: OpenAIGPT5Config):
|
|||
|
||||
|
||||
def test_is_explicitly_disabled_factory_minimal():
|
||||
"""_is_explicitly_disabled_factory returns True only for explicit False entries.
|
||||
"""is_explicitly_disabled_factory returns True only for explicit False entries.
|
||||
|
||||
Verifies the shared helper used by _is_reasoning_effort_level_explicitly_disabled
|
||||
directly — so future changes to the helper are caught without going through the
|
||||
method wrapper.
|
||||
"""
|
||||
key = "supports_minimal_reasoning_effort"
|
||||
assert _is_explicitly_disabled_factory("gpt-5.4-mini", None, key)
|
||||
assert _is_explicitly_disabled_factory("gpt-5.4-nano", None, key)
|
||||
assert _is_explicitly_disabled_factory("openai/gpt-5.4-mini", None, key)
|
||||
assert _is_explicitly_disabled_factory("gpt-5.4", None, key)
|
||||
assert _is_explicitly_disabled_factory("gpt-5.4-pro", None, key)
|
||||
assert not _is_explicitly_disabled_factory("gpt-5.4-turbo-preview", None, key)
|
||||
assert is_explicitly_disabled_factory("gpt-5.4-mini", None, key)
|
||||
assert is_explicitly_disabled_factory("gpt-5.4-nano", None, key)
|
||||
assert is_explicitly_disabled_factory("openai/gpt-5.4-mini", None, key)
|
||||
assert is_explicitly_disabled_factory("gpt-5.4", None, key)
|
||||
assert is_explicitly_disabled_factory("gpt-5.4-pro", None, key)
|
||||
assert not is_explicitly_disabled_factory("gpt-5.4-turbo-preview", None, key)
|
||||
|
||||
|
||||
def test_gpt5_unknown_model_passes_through_minimal(config: OpenAIConfig):
|
||||
|
|
|
|||
|
|
@ -5,11 +5,14 @@ from copy import deepcopy
|
|||
from typing import Final, List, cast
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
from pydantic import BaseModel
|
||||
|
||||
import litellm
|
||||
from litellm import ModelResponse, completion
|
||||
from litellm.llms.anthropic.experimental_pass_through.messages import handler as anthropic_messages_handler
|
||||
from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler
|
||||
from litellm.llms.gemini.chat.transformation import GoogleAIStudioGeminiConfig
|
||||
from litellm.llms.vertex_ai.common_utils import VertexAIError
|
||||
from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import (
|
||||
|
|
@ -2678,6 +2681,118 @@ def test_reasoning_effort_maps_to_thinking_level_gemini_3():
|
|||
assert result["thinkingConfig"]["includeThoughts"] is False
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"model",
|
||||
[
|
||||
"gemini-3.7-flash",
|
||||
"vertex_ai/gemini-3.8-flash",
|
||||
"gemini/gemini-3.8-flash",
|
||||
],
|
||||
)
|
||||
@pytest.mark.parametrize(
|
||||
("reasoning_effort", "include_thoughts"),
|
||||
[("minimal", True), ("none", False), ("disable", False)],
|
||||
)
|
||||
def test_gemini_37_38_flash_floor_minimal_thinking_level(
|
||||
local_model_cost_map, model, reasoning_effort, include_thoughts
|
||||
):
|
||||
result = VertexGeminiConfig._map_reasoning_effort_to_thinking_level(
|
||||
reasoning_effort, model
|
||||
)
|
||||
|
||||
assert result["thinkingLevel"] == "low"
|
||||
assert result["includeThoughts"] is include_thoughts
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("model", "reasoning_effort", "expected_level", "include_thoughts"),
|
||||
[
|
||||
("gemini-3-flash-preview", "minimal", "minimal", True),
|
||||
("gemini-3-flash-preview", "none", "minimal", False),
|
||||
("gemini-3-flash-preview", "disable", "minimal", False),
|
||||
("gemini-3.6-flash", "minimal", "minimal", True),
|
||||
("gemini-3.6-flash", "none", "minimal", False),
|
||||
("gemini-3.6-flash", "disable", "minimal", False),
|
||||
("gemini-3.5-flash", "minimal", "minimal", True),
|
||||
("gemini-3.5-flash", "none", "minimal", False),
|
||||
("gemini-3.5-flash", "disable", "minimal", False),
|
||||
("gemini-3.8-flash", "medium", "medium", True),
|
||||
],
|
||||
)
|
||||
def test_gemini_flash_minimal_thinking_support(
|
||||
local_model_cost_map, model, reasoning_effort, expected_level, include_thoughts
|
||||
):
|
||||
result = VertexGeminiConfig._map_reasoning_effort_to_thinking_level(
|
||||
reasoning_effort, model
|
||||
)
|
||||
|
||||
assert result["thinkingLevel"] == expected_level
|
||||
assert result["includeThoughts"] is include_thoughts
|
||||
|
||||
|
||||
def test_gemini_38_flash_feature_flag_uses_low_thinking_level(local_model_cost_map, monkeypatch):
|
||||
monkeypatch.setattr(litellm, "enable_gemini_default_thinking_level_low", True)
|
||||
thinking_param = {"type": "enabled", "budget_tokens": 1024}
|
||||
|
||||
result_38 = VertexGeminiConfig._map_thinking_param(
|
||||
thinking_param, model="gemini-3.8-flash"
|
||||
)
|
||||
result_36 = VertexGeminiConfig._map_thinking_param(
|
||||
thinking_param, model="gemini-3.6-flash"
|
||||
)
|
||||
|
||||
assert result_38["thinkingLevel"] == "low"
|
||||
assert result_36["thinkingLevel"] == "minimal"
|
||||
|
||||
|
||||
def test_gemini_38_flash_public_reasoning_effort_none_uses_low(local_model_cost_map):
|
||||
result = VertexGeminiConfig().map_openai_params(
|
||||
non_default_params={"reasoning_effort": "none"},
|
||||
optional_params={},
|
||||
model="gemini-3.8-flash",
|
||||
drop_params=False,
|
||||
)
|
||||
|
||||
assert result["thinkingConfig"] == {
|
||||
"thinkingLevel": "low",
|
||||
"includeThoughts": False,
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_gemini_38_flash_messages_bridge_thinking_disabled_sends_low_thinking_level(local_model_cost_map):
|
||||
captured: dict[str, dict] = {}
|
||||
|
||||
def upstream(request: httpx.Request) -> httpx.Response:
|
||||
captured["body"] = json.loads(request.content)
|
||||
return httpx.Response(
|
||||
200,
|
||||
json={
|
||||
"candidates": [{"content": {"parts": [{"text": "hi"}], "role": "model"}, "finishReason": "STOP"}],
|
||||
"usageMetadata": {"promptTokenCount": 1, "candidatesTokenCount": 1, "totalTokenCount": 2},
|
||||
},
|
||||
request=request,
|
||||
)
|
||||
|
||||
client = AsyncHTTPHandler()
|
||||
client.client = httpx.AsyncClient(transport=httpx.MockTransport(upstream))
|
||||
|
||||
await anthropic_messages_handler.anthropic_messages(
|
||||
max_tokens=16,
|
||||
messages=[{"role": "user", "content": "hi"}],
|
||||
model="gemini/gemini-3.8-flash",
|
||||
custom_llm_provider="gemini",
|
||||
thinking={"type": "disabled"},
|
||||
api_key="fake-gemini-key",
|
||||
client=client,
|
||||
)
|
||||
|
||||
assert captured["body"]["generationConfig"]["thinkingConfig"] == {
|
||||
"thinkingLevel": "low",
|
||||
"includeThoughts": False,
|
||||
}
|
||||
|
||||
|
||||
def test_reasoning_effort_dict_format_gemini_3():
|
||||
"""
|
||||
Test that reasoning_effort works when passed as dict format from OpenAI Agents SDK.
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
||||
|
|
|
|||
|
|
@ -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),
|
||||
|
|
|
|||
|
|
@ -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),
|
||||
|
|
|
|||
|
|
@ -2892,45 +2892,49 @@ def test_team_update_gate_allows_org_admin_with_resolved_org():
|
|||
)
|
||||
|
||||
|
||||
def test_team_update_gate_rejects_without_org_context():
|
||||
"""Without organization_id (i.e. resolution found no org, or a non-org-admin),
|
||||
the gate still rejects /team/update — the fix adds no blanket allow. Guards
|
||||
against re-widening the route (e.g. dropping it into self_managed_routes)."""
|
||||
def test_team_update_gate_admits_internal_user_without_org_context(): # test-quality-ok: the gate's only success signal is not raising; the handler's team-admin 403s are pinned in test_team_endpoints
|
||||
"""/team/update is self-managed (LIT-5722): the coarse gate admits any authenticated
|
||||
caller and update_team resolves proxy, org or team admin itself, then filters team admins
|
||||
through the team_admin_editable_team_fields setting. Before that the gate 401'd every
|
||||
team admin, which left the handler's team-admin branch unreachable."""
|
||||
user_obj = LiteLLM_UserTable(
|
||||
user_id="team-admin-user",
|
||||
user_role=LitellmUserRoles.INTERNAL_USER.value,
|
||||
organization_memberships=None,
|
||||
)
|
||||
valid_token = UserAPIKeyAuth(user_id="team-admin-user", user_role=LitellmUserRoles.INTERNAL_USER.value)
|
||||
request = MagicMock(spec=Request)
|
||||
request.method = "POST"
|
||||
request.query_params = {}
|
||||
|
||||
RouteChecks.non_proxy_admin_allowed_routes_check(
|
||||
user_obj=user_obj,
|
||||
_user_role=LitellmUserRoles.INTERNAL_USER.value,
|
||||
route="/team/update",
|
||||
request=request,
|
||||
valid_token=valid_token,
|
||||
request_data={"team_id": "team-1", "max_budget": 42},
|
||||
)
|
||||
|
||||
|
||||
def test_team_update_gate_defers_cross_org_admin_to_the_handler(): # test-quality-ok: the gate's only success signal is not raising; the handler's 403 it defers to is pinned in test_team_endpoints
|
||||
"""An org admin of a DIFFERENT org clears the coarse gate like any internal user;
|
||||
update_team's _resolve_team_access finds no role on the team and 403s (pinned in
|
||||
test_team_endpoints), so there is still no cross-org escalation."""
|
||||
user_obj = _make_org_admin_user("org-1")
|
||||
valid_token = UserAPIKeyAuth(user_id="org-admin-user", user_role=LitellmUserRoles.INTERNAL_USER.value)
|
||||
request = MagicMock(spec=Request)
|
||||
request.method = "POST"
|
||||
request.query_params = {}
|
||||
|
||||
with pytest.raises(Exception, match="Only proxy admin can be used to generate"):
|
||||
RouteChecks.non_proxy_admin_allowed_routes_check(
|
||||
user_obj=user_obj,
|
||||
_user_role=LitellmUserRoles.INTERNAL_USER.value,
|
||||
route="/team/update",
|
||||
request=request,
|
||||
valid_token=valid_token,
|
||||
request_data={"team_id": "team-1", "max_budget": 42},
|
||||
)
|
||||
|
||||
|
||||
def test_team_update_gate_rejects_cross_org_admin_with_resolved_org():
|
||||
"""Even after the target team's org is resolved, an org admin of a DIFFERENT
|
||||
org is rejected at the gate (no cross-org escalation)."""
|
||||
user_obj = _make_org_admin_user("org-1")
|
||||
valid_token = UserAPIKeyAuth(user_id="org-admin-user", user_role=LitellmUserRoles.INTERNAL_USER.value)
|
||||
request = MagicMock(spec=Request)
|
||||
request.method = "POST"
|
||||
request.query_params = {}
|
||||
|
||||
with pytest.raises(Exception, match="Only proxy admin can be used to generate"):
|
||||
RouteChecks.non_proxy_admin_allowed_routes_check(
|
||||
user_obj=user_obj,
|
||||
_user_role=LitellmUserRoles.INTERNAL_USER.value,
|
||||
route="/team/update",
|
||||
request=request,
|
||||
valid_token=valid_token,
|
||||
request_data={"team_id": "team-1", "organization_id": "org-2"},
|
||||
)
|
||||
RouteChecks.non_proxy_admin_allowed_routes_check(
|
||||
user_obj=user_obj,
|
||||
_user_role=LitellmUserRoles.INTERNAL_USER.value,
|
||||
route="/team/update",
|
||||
request=request,
|
||||
valid_token=valid_token,
|
||||
request_data={"team_id": "team-1", "organization_id": "org-2"},
|
||||
)
|
||||
|
||||
|
||||
# ── PATCH /team/{team_id}: same org-context + role reach as POST /team/update ──
|
||||
|
|
@ -2993,23 +2997,6 @@ async def test_add_team_org_context_noop_for_static_team_route():
|
|||
assert out == body
|
||||
|
||||
|
||||
def test_patch_team_route_has_same_reach_as_team_update():
|
||||
"""/team/{team_id} is reachable by org admins (in org_admin_allowed_routes) but
|
||||
NOT by regular internal users or the role-agnostic self_managed_routes — the
|
||||
latter would open /team/new (the collision footgun) to any authenticated user."""
|
||||
from litellm.proxy._types import LiteLLMRoutes
|
||||
|
||||
assert RouteChecks.check_route_access(
|
||||
route="/team/abc-123", allowed_routes=LiteLLMRoutes.org_admin_allowed_routes.value
|
||||
)
|
||||
assert not RouteChecks.check_route_access(
|
||||
route="/team/abc-123", allowed_routes=LiteLLMRoutes.internal_user_routes.value
|
||||
)
|
||||
assert not RouteChecks.check_route_access(
|
||||
route="/team/abc-123", allowed_routes=LiteLLMRoutes.self_managed_routes.value
|
||||
)
|
||||
|
||||
|
||||
def _patch_team_request() -> MagicMock:
|
||||
request = MagicMock(spec=Request)
|
||||
request.method = "PATCH"
|
||||
|
|
@ -3897,7 +3884,6 @@ def test_team_disable_logging_stays_proxy_admin_only():
|
|||
"route",
|
||||
[
|
||||
"/team/06bda574-5ca9-43d3-beb8-3b23c2f17112",
|
||||
"/team/update",
|
||||
"/team/06bda574-5ca9-43d3-beb8-3b23c2f17112/model/add",
|
||||
],
|
||||
)
|
||||
|
|
|
|||
|
|
@ -0,0 +1,138 @@
|
|||
import pytest
|
||||
from fastapi import HTTPException
|
||||
|
||||
from litellm.proxy._types import LiteLLM_ModelTable, LiteLLM_TeamTable, UpdateTeamRequest
|
||||
from litellm.proxy.management_endpoints.team_admin_field_permissions import (
|
||||
TeamAdminEditAllowed,
|
||||
TeamAdminEditingDisabled,
|
||||
TeamAdminFieldNotPermitted,
|
||||
changed_team_fields,
|
||||
resolve_team_admin_editable_fields,
|
||||
team_admin_edit_verdict,
|
||||
team_admin_request_or_raise,
|
||||
)
|
||||
|
||||
_SUPPORTED = frozenset({"tpm_limit", "rpm_limit", "team_alias"})
|
||||
|
||||
|
||||
def _team(**overrides):
|
||||
return LiteLLM_TeamTable(team_id="team-1", **overrides)
|
||||
|
||||
|
||||
class TestResolveTeamAdminEditableFields:
|
||||
def test_missing_setting_means_nothing_editable(self):
|
||||
assert resolve_team_admin_editable_fields({}, _SUPPORTED) == frozenset()
|
||||
|
||||
def test_keeps_only_supported_names(self):
|
||||
configured = {"team_admin_editable_team_fields": ["tpm_limit", "blocked", "organization_id"]}
|
||||
assert resolve_team_admin_editable_fields(configured, _SUPPORTED) == frozenset({"tpm_limit"})
|
||||
|
||||
@pytest.mark.parametrize("raw", ["tpm_limit", 7, {"tpm_limit": True}, [1, 2]])
|
||||
def test_malformed_setting_fails_closed(self, raw):
|
||||
assert resolve_team_admin_editable_fields({"team_admin_editable_team_fields": raw}, _SUPPORTED) == frozenset()
|
||||
|
||||
|
||||
class TestChangedTeamFields:
|
||||
def test_team_id_alone_changes_nothing(self):
|
||||
assert changed_team_fields(UpdateTeamRequest(team_id="team-1"), _team()) == frozenset()
|
||||
|
||||
def test_column_echoing_stored_value_is_not_a_change(self):
|
||||
data = UpdateTeamRequest(team_id="team-1", tpm_limit=5, team_alias="alpha", max_budget=None)
|
||||
assert changed_team_fields(data, _team(tpm_limit=5, team_alias="alpha")) == frozenset()
|
||||
|
||||
def test_column_with_different_value_is_a_change(self):
|
||||
data = UpdateTeamRequest(team_id="team-1", tpm_limit=6, team_alias="alpha")
|
||||
assert changed_team_fields(data, _team(tpm_limit=5, team_alias="alpha")) == frozenset({"tpm_limit"})
|
||||
|
||||
def test_explicit_null_clearing_a_stored_column_is_a_change(self):
|
||||
data = UpdateTeamRequest(team_id="team-1", max_budget=None)
|
||||
assert changed_team_fields(data, _team(max_budget=30.0)) == frozenset({"max_budget"})
|
||||
|
||||
def test_folded_field_sent_top_level_is_named_not_metadata(self):
|
||||
data = UpdateTeamRequest(team_id="team-1", guardrails=["b"])
|
||||
assert changed_team_fields(data, _team(metadata={"guardrails": ["a"]})) == frozenset({"guardrails"})
|
||||
|
||||
def test_folded_field_sent_inside_metadata_is_named_not_metadata(self):
|
||||
data = UpdateTeamRequest(team_id="team-1", metadata={"guardrails": ["b"]})
|
||||
assert changed_team_fields(data, _team(metadata={"guardrails": ["a"]})) == frozenset({"guardrails"})
|
||||
|
||||
def test_custom_metadata_key_change_is_attributed_to_metadata(self):
|
||||
data = UpdateTeamRequest(team_id="team-1", metadata={"guardrails": ["a"], "cost_center": "b"})
|
||||
existing = _team(metadata={"guardrails": ["a"], "cost_center": "a"})
|
||||
assert changed_team_fields(data, existing) == frozenset({"metadata"})
|
||||
|
||||
def test_metadata_echo_with_top_level_override_only_names_the_override(self):
|
||||
data = UpdateTeamRequest(team_id="team-1", guardrails=["b"], metadata={"guardrails": ["a"], "cost_center": "a"})
|
||||
existing = _team(metadata={"guardrails": ["a"], "cost_center": "a"})
|
||||
assert changed_team_fields(data, existing) == frozenset({"guardrails"})
|
||||
|
||||
def test_dropping_a_stored_key_from_submitted_metadata_is_a_change(self):
|
||||
data = UpdateTeamRequest(team_id="team-1", metadata={"cost_center": "a"})
|
||||
existing = _team(metadata={"cost_center": "a", "tags": ["x"], "logging": [{"callback": "langfuse"}]})
|
||||
assert changed_team_fields(data, existing) == frozenset({"tags", "logging"})
|
||||
|
||||
def test_server_managed_metadata_key_is_ignored(self):
|
||||
data = UpdateTeamRequest(team_id="team-1", metadata={"cost_center": "a"})
|
||||
existing = _team(metadata={"cost_center": "a", "team_member_budget_id": "budget-1"})
|
||||
assert changed_team_fields(data, existing) == frozenset()
|
||||
|
||||
def test_model_aliases_compare_against_the_model_table(self):
|
||||
table = LiteLLM_ModelTable(model_aliases='{"fast": "gpt-4o-mini"}', created_by="a", updated_by="a")
|
||||
same = UpdateTeamRequest(team_id="team-1", model_aliases={"fast": "gpt-4o-mini"})
|
||||
different = UpdateTeamRequest(team_id="team-1", model_aliases={"fast": "gpt-4o"})
|
||||
assert changed_team_fields(same, _team(litellm_model_table=table)) == frozenset()
|
||||
assert changed_team_fields(different, _team(litellm_model_table=table)) == frozenset({"model_aliases"})
|
||||
|
||||
def test_empty_model_aliases_against_no_model_table_is_not_a_change(self):
|
||||
assert changed_team_fields(UpdateTeamRequest(team_id="team-1", model_aliases={}), _team()) == frozenset()
|
||||
|
||||
def test_field_without_a_stored_counterpart_counts_as_changed_when_sent(self):
|
||||
data = UpdateTeamRequest(team_id="team-1", team_member_budget=10.0)
|
||||
assert changed_team_fields(data, _team()) == frozenset({"team_member_budget"})
|
||||
|
||||
|
||||
class TestTeamAdminEditVerdict:
|
||||
def test_no_permitted_fields_disables_editing_even_for_a_no_op(self):
|
||||
verdict = team_admin_edit_verdict(UpdateTeamRequest(team_id="team-1"), _team(), frozenset())
|
||||
assert verdict == TeamAdminEditingDisabled()
|
||||
|
||||
def test_allowed_request_keeps_only_the_changed_fields(self):
|
||||
data = UpdateTeamRequest(team_id="team-1", tpm_limit=6, team_alias="alpha", budget_duration="30d")
|
||||
existing = _team(team_alias="alpha", budget_duration="30d")
|
||||
verdict = team_admin_edit_verdict(data, existing, frozenset({"tpm_limit"}))
|
||||
assert isinstance(verdict, TeamAdminEditAllowed)
|
||||
assert verdict.request.model_dump(exclude_unset=True) == {"team_id": "team-1", "tpm_limit": 6}
|
||||
|
||||
def test_permitted_field_changed_inside_metadata_keeps_the_metadata(self):
|
||||
data = UpdateTeamRequest(team_id="team-1", metadata={"guardrails": ["b"]}, team_alias="alpha")
|
||||
existing = _team(team_alias="alpha", metadata={"guardrails": ["a"]})
|
||||
verdict = team_admin_edit_verdict(data, existing, frozenset({"guardrails"}))
|
||||
assert isinstance(verdict, TeamAdminEditAllowed)
|
||||
assert verdict.request.model_dump(exclude_unset=True) == {
|
||||
"team_id": "team-1",
|
||||
"metadata": {"guardrails": ["b"]},
|
||||
}
|
||||
|
||||
def test_first_blocked_field_in_sorted_order_is_reported(self):
|
||||
data = UpdateTeamRequest(team_id="team-1", tpm_limit=6, rpm_limit=6, blocked=True)
|
||||
verdict = team_admin_edit_verdict(data, _team(), frozenset({"tpm_limit"}))
|
||||
assert verdict == TeamAdminFieldNotPermitted(field="blocked")
|
||||
|
||||
|
||||
class TestTeamAdminRequestOrRaise:
|
||||
def test_allowed_hands_back_its_request(self):
|
||||
request = UpdateTeamRequest(team_id="team-1", tpm_limit=6)
|
||||
assert team_admin_request_or_raise(TeamAdminEditAllowed(request=request)) is request
|
||||
|
||||
def test_disabled_is_a_403_pointing_at_the_proxy_admin(self):
|
||||
with pytest.raises(HTTPException) as exc:
|
||||
team_admin_request_or_raise(TeamAdminEditingDisabled())
|
||||
assert exc.value.status_code == 403
|
||||
assert "cannot edit team settings" in exc.value.detail
|
||||
assert "Settings > UI > Team admin editable fields" in exc.value.detail
|
||||
|
||||
def test_field_not_permitted_is_a_403_naming_the_field(self):
|
||||
with pytest.raises(HTTPException) as exc:
|
||||
team_admin_request_or_raise(TeamAdminFieldNotPermitted(field="blocked"))
|
||||
assert exc.value.status_code == 403
|
||||
assert "'blocked'" in exc.value.detail
|
||||
|
|
@ -1,6 +1,6 @@
|
|||
import asyncio
|
||||
import json
|
||||
from contextlib import asynccontextmanager
|
||||
from contextlib import asynccontextmanager, contextmanager
|
||||
from datetime import datetime, timezone
|
||||
from types import SimpleNamespace
|
||||
from typing import Final, Optional, cast
|
||||
|
|
@ -76,6 +76,31 @@ from litellm.types.proxy.management_endpoints.team_endpoints import (
|
|||
client = TestClient(app)
|
||||
|
||||
|
||||
@contextmanager
|
||||
def _team_admin_may_edit(*fields: str):
|
||||
"""Let team admins change ``fields`` on /team/update for the duration of the block.
|
||||
|
||||
The registry only lists the fields shipped so far (LIT-5722 adds them one PR at a time), so tests that
|
||||
exercise the gates layered underneath the allow-list widen it here instead of asserting the early 403."""
|
||||
with (
|
||||
patch( # test-quality-ok: the registry is a module constant update_team reads directly; no seam to inject
|
||||
"litellm.proxy.management_endpoints.team_endpoints.SUPPORTED_TEAM_ADMIN_EDITABLE_TEAM_FIELDS",
|
||||
frozenset(fields),
|
||||
),
|
||||
patch("litellm.proxy.proxy_server.general_settings", {"team_admin_editable_team_fields": list(fields)}), # test-quality-ok: update_team reads general_settings as a proxy_server module global
|
||||
):
|
||||
yield
|
||||
|
||||
|
||||
def _not_org_admin():
|
||||
"""update_team asks whether the caller administers the team's org before it settles for team admin;
|
||||
a MagicMock prisma cannot answer that lookup, so pin it to False."""
|
||||
return patch( # test-quality-ok: the org-admin lookup needs a real prisma client this file's MagicMock cannot provide
|
||||
"litellm.proxy.management_endpoints.team_endpoints._is_user_org_admin_for_team",
|
||||
AsyncMock(return_value=False),
|
||||
)
|
||||
|
||||
|
||||
def _wire_team_create_tx(prisma_client):
|
||||
"""`/team/new` inserts the team and mirrors it onto the access groups in one transaction,
|
||||
so a mocked client has to hand its team table back out of `db.tx()`.
|
||||
|
|
@ -6393,6 +6418,7 @@ async def test_update_team_standalone_budget_raise_blocked_for_team_admin():
|
|||
dummy_request = MagicMock(spec=Request)
|
||||
|
||||
with (
|
||||
_team_admin_may_edit("max_budget"),
|
||||
patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma,
|
||||
patch("litellm.proxy.proxy_server.user_api_key_cache") as mock_cache,
|
||||
patch("litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin"),
|
||||
|
|
@ -6549,6 +6575,7 @@ async def test_update_team_standalone_budget_removal_blocked_for_team_admin():
|
|||
dummy_request = MagicMock(spec=Request)
|
||||
|
||||
with (
|
||||
_team_admin_may_edit("max_budget"),
|
||||
patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma,
|
||||
patch("litellm.proxy.proxy_server.user_api_key_cache") as mock_cache,
|
||||
patch("litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin"),
|
||||
|
|
@ -6618,6 +6645,7 @@ async def test_update_team_standalone_uncapped_team_admin_sets_finite_allowed(
|
|||
dummy_request = MagicMock(spec=Request)
|
||||
|
||||
with (
|
||||
_team_admin_may_edit("max_budget"),
|
||||
patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma,
|
||||
patch("litellm.proxy.proxy_server.user_api_key_cache") as mock_cache,
|
||||
patch("litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin"),
|
||||
|
|
@ -6712,6 +6740,7 @@ async def test_update_team_standalone_unchanged_budget_allowed(
|
|||
dummy_request = MagicMock(spec=Request)
|
||||
|
||||
with (
|
||||
_team_admin_may_edit("max_budget", "tpm_limit"),
|
||||
patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma,
|
||||
patch("litellm.proxy.proxy_server.user_api_key_cache") as mock_cache,
|
||||
patch("litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin"),
|
||||
|
|
@ -6810,6 +6839,7 @@ async def test_update_team_standalone_lower_budget_allowed(
|
|||
dummy_request = MagicMock(spec=Request)
|
||||
|
||||
with (
|
||||
_team_admin_may_edit("max_budget"),
|
||||
patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma,
|
||||
patch("litellm.proxy.proxy_server.user_api_key_cache") as mock_cache,
|
||||
patch("litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin"),
|
||||
|
|
@ -6912,6 +6942,8 @@ async def test_update_team_org_scoped_budget_exceeds_org_limit():
|
|||
mock_org.litellm_budget_table = mock_budget_table
|
||||
|
||||
with (
|
||||
_team_admin_may_edit("max_budget"),
|
||||
_not_org_admin(),
|
||||
patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma,
|
||||
patch("litellm.proxy.proxy_server.user_api_key_cache") as mock_cache,
|
||||
patch("litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin"),
|
||||
|
|
@ -6992,6 +7024,7 @@ async def test_update_team_standalone_models_not_gated_by_user_limit(
|
|||
dummy_request = MagicMock(spec=Request)
|
||||
|
||||
with (
|
||||
_team_admin_may_edit("models"),
|
||||
patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma,
|
||||
patch("litellm.proxy.proxy_server.user_api_key_cache") as mock_cache,
|
||||
patch("litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin"),
|
||||
|
|
@ -7091,6 +7124,8 @@ async def test_update_team_org_scoped_budget_bypasses_user_limit(
|
|||
mock_org.litellm_budget_table = mock_budget_table
|
||||
|
||||
with (
|
||||
_team_admin_may_edit("max_budget"),
|
||||
_not_org_admin(),
|
||||
patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma,
|
||||
patch("litellm.proxy.proxy_server.user_api_key_cache") as mock_cache,
|
||||
patch("litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin"),
|
||||
|
|
@ -7202,6 +7237,8 @@ async def test_update_team_org_scoped_models_bypasses_user_limit(
|
|||
mock_org.litellm_budget_table = None
|
||||
|
||||
with (
|
||||
_team_admin_may_edit("models"),
|
||||
_not_org_admin(),
|
||||
patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma,
|
||||
patch("litellm.proxy.proxy_server.user_api_key_cache") as mock_cache,
|
||||
patch("litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin"),
|
||||
|
|
@ -7304,6 +7341,8 @@ async def test_update_team_org_scoped_models_not_in_org_models():
|
|||
mock_org.litellm_budget_table = None
|
||||
|
||||
with (
|
||||
_team_admin_may_edit("models"),
|
||||
_not_org_admin(),
|
||||
patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma,
|
||||
patch("litellm.proxy.proxy_server.user_api_key_cache") as mock_cache,
|
||||
patch("litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin"),
|
||||
|
|
@ -7393,6 +7432,8 @@ async def test_update_team_org_scoped_models_with_all_proxy_models(
|
|||
mock_org.litellm_budget_table = None
|
||||
|
||||
with (
|
||||
_team_admin_may_edit("models"),
|
||||
_not_org_admin(),
|
||||
patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma,
|
||||
patch("litellm.proxy.proxy_server.user_api_key_cache") as mock_cache,
|
||||
patch("litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin"),
|
||||
|
|
@ -7502,6 +7543,7 @@ async def test_update_team_tpm_limit_not_gated_by_user_limit(
|
|||
dummy_request = MagicMock(spec=Request)
|
||||
|
||||
with (
|
||||
_team_admin_may_edit("tpm_limit"),
|
||||
patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma,
|
||||
patch("litellm.proxy.proxy_server.user_api_key_cache") as mock_cache,
|
||||
patch("litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin"),
|
||||
|
|
@ -7584,6 +7626,7 @@ async def test_update_team_rpm_limit_not_gated_by_user_limit(
|
|||
dummy_request = MagicMock(spec=Request)
|
||||
|
||||
with (
|
||||
_team_admin_may_edit("rpm_limit"),
|
||||
patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma,
|
||||
patch("litellm.proxy.proxy_server.user_api_key_cache") as mock_cache,
|
||||
patch("litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin"),
|
||||
|
|
@ -7981,6 +8024,8 @@ async def test_update_team_org_scoped_tpm_exceeds_org_limit():
|
|||
mock_org.litellm_budget_table = mock_budget_table
|
||||
|
||||
with (
|
||||
_team_admin_may_edit("tpm_limit"),
|
||||
_not_org_admin(),
|
||||
patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma,
|
||||
patch("litellm.proxy.proxy_server.user_api_key_cache") as mock_cache,
|
||||
patch("litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin"),
|
||||
|
|
@ -8067,6 +8112,8 @@ async def test_update_team_org_scoped_rpm_exceeds_org_limit():
|
|||
mock_org.litellm_budget_table = mock_budget_table
|
||||
|
||||
with (
|
||||
_team_admin_may_edit("rpm_limit"),
|
||||
_not_org_admin(),
|
||||
patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma,
|
||||
patch("litellm.proxy.proxy_server.user_api_key_cache") as mock_cache,
|
||||
patch("litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin"),
|
||||
|
|
@ -8158,6 +8205,8 @@ async def test_update_team_org_scoped_tpm_rpm_bypasses_user_limit(
|
|||
mock_org.litellm_budget_table = mock_budget_table
|
||||
|
||||
with (
|
||||
_team_admin_may_edit("tpm_limit", "rpm_limit"),
|
||||
_not_org_admin(),
|
||||
patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma,
|
||||
patch("litellm.proxy.proxy_server.user_api_key_cache") as mock_cache,
|
||||
patch("litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin"),
|
||||
|
|
@ -8286,6 +8335,7 @@ async def test_update_team_guardrails_with_org_id(
|
|||
}
|
||||
|
||||
with (
|
||||
_team_admin_may_edit("guardrails", "organization_id"),
|
||||
patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma,
|
||||
patch("litellm.proxy.proxy_server.user_api_key_cache") as mock_cache,
|
||||
patch("litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin"),
|
||||
|
|
@ -11177,8 +11227,8 @@ async def test_update_team_blocks_non_admin_passthrough_routes(mock_db_client):
|
|||
mock_db_client.db.litellm_teamtable.find_unique = AsyncMock(return_value=existing)
|
||||
|
||||
with patch(
|
||||
"litellm.proxy.management_endpoints.team_endpoints._verify_team_access",
|
||||
AsyncMock(return_value=None),
|
||||
"litellm.proxy.management_endpoints.team_endpoints._resolve_team_access",
|
||||
AsyncMock(return_value="org_admin"),
|
||||
):
|
||||
with pytest.raises(ProxyException) as exc:
|
||||
await update_team(
|
||||
|
|
@ -13246,6 +13296,7 @@ async def test_update_team_output_token_estimate_lowered_rejected_for_team_admin
|
|||
|
||||
with contextlib.ExitStack() as stack:
|
||||
_wire_update_team(stack, {_TEAM_ESTIMATE: 4000})
|
||||
stack.enter_context(_team_admin_may_edit("default_estimated_output_tokens"))
|
||||
with pytest.raises(ProxyException) as exc:
|
||||
await update_team(
|
||||
data=UpdateTeamRequest(team_id="test_team_id", default_estimated_output_tokens=1),
|
||||
|
|
@ -13277,6 +13328,7 @@ async def test_update_team_output_token_estimate_unchanged_allows_team_admin_edi
|
|||
|
||||
with contextlib.ExitStack() as stack:
|
||||
prisma = _wire_update_team(stack, {_TEAM_ESTIMATE: 4000})
|
||||
stack.enter_context(_team_admin_may_edit("team_alias"))
|
||||
await update_team(
|
||||
data=UpdateTeamRequest(
|
||||
team_id="test_team_id",
|
||||
|
|
@ -13336,6 +13388,7 @@ async def test_update_team_batch_enqueued_token_limit_raised_rejected_for_team_a
|
|||
|
||||
with contextlib.ExitStack() as stack:
|
||||
_wire_update_team(stack, {_TEAM_BATCH_LIMIT: 100000})
|
||||
stack.enter_context(_team_admin_may_edit("metadata"))
|
||||
with pytest.raises(ProxyException) as exc:
|
||||
await update_team(
|
||||
data=UpdateTeamRequest(team_id="test_team_id", metadata={_TEAM_BATCH_LIMIT: 10**12}),
|
||||
|
|
@ -14894,3 +14947,369 @@ async def test_update_team_model_max_budget_raise_blocked_for_team_admin():
|
|||
assert exc.value.code == "403"
|
||||
assert "proxy admin" in str(exc.value.message).lower()
|
||||
mock_prisma.db.litellm_teamtable.update.assert_not_awaited()
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# LIT-5722: team admins reach update_team through self_managed_routes and are
|
||||
# filtered by the team_admin_editable_team_fields setting.
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
_TEAM_ADMIN_CALLER = UserAPIKeyAuth(
|
||||
user_role=LitellmUserRoles.INTERNAL_USER, api_key="sk-team-admin", user_id="team-admin"
|
||||
)
|
||||
_PROXY_ADMIN_CALLER = UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN, api_key="sk-admin", user_id="admin")
|
||||
|
||||
|
||||
def _update_request_stub():
|
||||
from unittest.mock import Mock
|
||||
|
||||
from fastapi import Request
|
||||
|
||||
return Mock(spec=Request)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_update_team_team_admin_is_refused_before_any_write_when_no_fields_are_enabled():
|
||||
import contextlib
|
||||
|
||||
with contextlib.ExitStack() as stack:
|
||||
prisma = _wire_update_team(stack, {})
|
||||
stack.enter_context(_team_admin_may_edit())
|
||||
with pytest.raises(ProxyException) as exc:
|
||||
await update_team(
|
||||
data=UpdateTeamRequest(team_id="test_team_id", team_alias="renamed"),
|
||||
http_request=_update_request_stub(),
|
||||
user_api_key_dict=_TEAM_ADMIN_CALLER,
|
||||
)
|
||||
|
||||
assert str(exc.value.code) == "403"
|
||||
assert "cannot edit team settings" in str(exc.value.message)
|
||||
assert not prisma.db.litellm_teamtable.update.called
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_update_team_configured_but_unsupported_field_does_not_open_editing():
|
||||
"""Only fields in SUPPORTED_TEAM_ADMIN_EDITABLE_TEAM_FIELDS count, whatever general_settings says."""
|
||||
import contextlib
|
||||
|
||||
with contextlib.ExitStack() as stack:
|
||||
prisma = _wire_update_team(stack, {})
|
||||
stack.enter_context(
|
||||
patch("litellm.proxy.proxy_server.general_settings", {"team_admin_editable_team_fields": ["team_alias"]}) # test-quality-ok: update_team reads general_settings as a proxy_server module global
|
||||
)
|
||||
with pytest.raises(ProxyException) as exc:
|
||||
await update_team(
|
||||
data=UpdateTeamRequest(team_id="test_team_id", team_alias="renamed"),
|
||||
http_request=_update_request_stub(),
|
||||
user_api_key_dict=_TEAM_ADMIN_CALLER,
|
||||
)
|
||||
|
||||
assert str(exc.value.code) == "403"
|
||||
assert "cannot edit team settings" in str(exc.value.message)
|
||||
assert not prisma.db.litellm_teamtable.update.called
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_update_team_team_admin_changing_an_unpermitted_field_is_refused_by_name():
|
||||
import contextlib
|
||||
|
||||
with contextlib.ExitStack() as stack:
|
||||
prisma = _wire_update_team(stack, {})
|
||||
stack.enter_context(_team_admin_may_edit("team_alias"))
|
||||
with pytest.raises(ProxyException) as exc:
|
||||
await update_team(
|
||||
data=UpdateTeamRequest(team_id="test_team_id", team_alias="renamed", tpm_limit=10),
|
||||
http_request=_update_request_stub(),
|
||||
user_api_key_dict=_TEAM_ADMIN_CALLER,
|
||||
)
|
||||
|
||||
assert str(exc.value.code) == "403"
|
||||
assert "'tpm_limit'" in str(exc.value.message)
|
||||
assert not prisma.db.litellm_teamtable.update.called
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_update_team_team_admin_echoing_unpermitted_fields_unchanged_is_allowed(
|
||||
disable_audit_logging_for_mocked_team,
|
||||
):
|
||||
"""The dashboard resends the whole form, so only a value that differs from what is stored counts."""
|
||||
import contextlib
|
||||
|
||||
with contextlib.ExitStack() as stack:
|
||||
prisma = _wire_update_team(stack, {})
|
||||
stack.enter_context(_team_admin_may_edit("team_alias"))
|
||||
result = await update_team(
|
||||
data=UpdateTeamRequest(team_id="test_team_id", team_alias="renamed", tpm_limit=None, models=[]),
|
||||
http_request=_update_request_stub(),
|
||||
user_api_key_dict=_TEAM_ADMIN_CALLER,
|
||||
)
|
||||
|
||||
assert result["data"].team_id == "test_team_id"
|
||||
assert prisma.db.litellm_teamtable.update.called
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_update_team_team_admin_changes_tpm_limit_once_a_proxy_admin_enables_it(
|
||||
disable_audit_logging_for_mocked_team,
|
||||
):
|
||||
"""tpm_limit is the first field a proxy admin can open to team admins; every other field stays admin-only."""
|
||||
import contextlib
|
||||
|
||||
with contextlib.ExitStack() as stack:
|
||||
prisma = _wire_update_team(stack, {})
|
||||
stack.enter_context(
|
||||
patch("litellm.proxy.proxy_server.general_settings", {"team_admin_editable_team_fields": ["tpm_limit"]}) # test-quality-ok: update_team reads general_settings as a proxy_server module global
|
||||
)
|
||||
await update_team(
|
||||
data=UpdateTeamRequest(team_id="test_team_id", tpm_limit=5000),
|
||||
http_request=_update_request_stub(),
|
||||
user_api_key_dict=_TEAM_ADMIN_CALLER,
|
||||
)
|
||||
with pytest.raises(ProxyException) as refused:
|
||||
await update_team(
|
||||
data=UpdateTeamRequest(team_id="test_team_id", tpm_limit=6000, rpm_limit=10),
|
||||
http_request=_update_request_stub(),
|
||||
user_api_key_dict=_TEAM_ADMIN_CALLER,
|
||||
)
|
||||
|
||||
assert prisma.db.litellm_teamtable.update.await_count == 1
|
||||
assert prisma.db.litellm_teamtable.update.call_args.kwargs["data"]["tpm_limit"] == 5000
|
||||
assert str(refused.value.code) == "403"
|
||||
assert "'rpm_limit'" in str(refused.value.message)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_update_team_team_admin_resending_budget_settings_does_not_push_back_budget_resets(
|
||||
disable_audit_logging_for_mocked_team,
|
||||
):
|
||||
"""A resent budget_duration or budget_limits would otherwise recompute the reset timestamps from now."""
|
||||
import contextlib
|
||||
|
||||
stored_windows = [{"budget_duration": "7d", "max_budget": 5.0, "reset_at": "2026-09-20T00:00:00Z"}]
|
||||
budgeted_team = MagicMock()
|
||||
budgeted_team.metadata = {}
|
||||
budgeted_team.model_dump.return_value = {
|
||||
"team_id": "test_team_id",
|
||||
"team_alias": "test_team",
|
||||
"metadata": {},
|
||||
"budget_duration": "30d",
|
||||
"budget_limits": stored_windows,
|
||||
"members_with_roles": [{"user_id": "team-admin", "role": "admin"}],
|
||||
}
|
||||
|
||||
with contextlib.ExitStack() as stack:
|
||||
prisma = _wire_update_team(stack, {})
|
||||
prisma.db.litellm_teamtable.find_unique = AsyncMock(return_value=budgeted_team)
|
||||
stack.enter_context(_team_admin_may_edit("tpm_limit"))
|
||||
await update_team(
|
||||
data=UpdateTeamRequest(
|
||||
team_id="test_team_id", tpm_limit=5000, budget_duration="30d", budget_limits=stored_windows
|
||||
),
|
||||
http_request=_update_request_stub(),
|
||||
user_api_key_dict=_TEAM_ADMIN_CALLER,
|
||||
)
|
||||
|
||||
written = prisma.db.litellm_teamtable.update.call_args.kwargs["data"]
|
||||
assert written["tpm_limit"] == 5000
|
||||
assert not {"budget_duration", "budget_reset_at", "budget_limits"} & written.keys()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_update_team_holds_a_team_admin_to_the_org_tpm_limit(disable_audit_logging_for_mocked_team):
|
||||
"""The org ceiling lives on the org's budget row, so /team/update must load it to enforce the cap."""
|
||||
import contextlib
|
||||
|
||||
capped_org = LiteLLM_OrganizationTable(
|
||||
organization_id="capped-org",
|
||||
budget_id="capped-budget",
|
||||
created_by="admin",
|
||||
updated_by="admin",
|
||||
litellm_budget_table=LiteLLM_BudgetTable(tpm_limit=10000),
|
||||
)
|
||||
|
||||
async def org_lookup(**kwargs):
|
||||
return capped_org if kwargs.get("include_budget_table") else capped_org.model_copy(
|
||||
update={"litellm_budget_table": None}
|
||||
)
|
||||
|
||||
org_team = MagicMock()
|
||||
org_team.metadata = {}
|
||||
org_team.organization_id = "capped-org"
|
||||
org_team.model_dump.return_value = {
|
||||
"team_id": "test_team_id",
|
||||
"team_alias": "test_team",
|
||||
"organization_id": "capped-org",
|
||||
"metadata": {},
|
||||
"members_with_roles": [{"user_id": "team-admin", "role": "admin"}],
|
||||
}
|
||||
|
||||
with contextlib.ExitStack() as stack:
|
||||
prisma = _wire_update_team(stack, {})
|
||||
prisma.db.litellm_teamtable.find_unique = AsyncMock(return_value=org_team)
|
||||
stack.enter_context(_team_admin_may_edit("tpm_limit"))
|
||||
stack.enter_context(
|
||||
patch( # test-quality-ok: the org-admin lookup needs a real prisma client this file's MagicMock cannot provide
|
||||
"litellm.proxy.management_endpoints.team_endpoints._is_user_org_admin_for_team",
|
||||
AsyncMock(return_value=False),
|
||||
)
|
||||
)
|
||||
stack.enter_context(
|
||||
patch( # test-quality-ok: update_team reads orgs through this module-level import; no seam to inject
|
||||
"litellm.proxy.management_endpoints.team_endpoints.get_org_object",
|
||||
AsyncMock(side_effect=org_lookup),
|
||||
)
|
||||
)
|
||||
with pytest.raises(ProxyException) as over_cap:
|
||||
await update_team(
|
||||
data=UpdateTeamRequest(team_id="test_team_id", tpm_limit=20000),
|
||||
http_request=_update_request_stub(),
|
||||
user_api_key_dict=_TEAM_ADMIN_CALLER,
|
||||
)
|
||||
await update_team(
|
||||
data=UpdateTeamRequest(team_id="test_team_id", tpm_limit=8000),
|
||||
http_request=_update_request_stub(),
|
||||
user_api_key_dict=_TEAM_ADMIN_CALLER,
|
||||
)
|
||||
|
||||
assert str(over_cap.value.code) == "400"
|
||||
assert "exceeds organization's tpm_limit (10000)" in str(over_cap.value.message)
|
||||
assert prisma.db.litellm_teamtable.update.await_count == 1
|
||||
assert prisma.db.litellm_teamtable.update.call_args.kwargs["data"]["tpm_limit"] == 8000
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_update_team_org_admin_is_not_filtered_by_the_team_admin_field_list(
|
||||
disable_audit_logging_for_mocked_team,
|
||||
):
|
||||
"""A caller who is both org admin and roster admin keeps unrestricted edits."""
|
||||
import contextlib
|
||||
|
||||
with contextlib.ExitStack() as stack:
|
||||
prisma = _wire_update_team(stack, {})
|
||||
stack.enter_context(_team_admin_may_edit())
|
||||
stack.enter_context(
|
||||
patch( # test-quality-ok: the org-admin lookup needs a real prisma client this file's MagicMock cannot provide
|
||||
"litellm.proxy.management_endpoints.team_endpoints._is_user_org_admin_for_team",
|
||||
AsyncMock(return_value=True),
|
||||
)
|
||||
)
|
||||
result = await update_team(
|
||||
data=UpdateTeamRequest(team_id="test_team_id", team_alias="renamed"),
|
||||
http_request=_update_request_stub(),
|
||||
user_api_key_dict=_TEAM_ADMIN_CALLER,
|
||||
)
|
||||
|
||||
assert result["data"].team_id == "test_team_id"
|
||||
assert prisma.db.litellm_teamtable.update.called
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_update_team_unknown_team_is_403_for_non_proxy_admins_and_404_for_proxy_admins():
|
||||
"""Now that any authenticated caller reaches the handler, 'team not found' must not leak team ids."""
|
||||
import contextlib
|
||||
|
||||
with contextlib.ExitStack() as stack:
|
||||
prisma = _wire_update_team(stack, {})
|
||||
prisma.db.litellm_teamtable.find_unique = AsyncMock(return_value=None)
|
||||
|
||||
with pytest.raises(ProxyException) as denied:
|
||||
await update_team(
|
||||
data=UpdateTeamRequest(team_id="no-such-team", team_alias="renamed"),
|
||||
http_request=_update_request_stub(),
|
||||
user_api_key_dict=_TEAM_ADMIN_CALLER,
|
||||
)
|
||||
with pytest.raises(ProxyException) as missing:
|
||||
await update_team(
|
||||
data=UpdateTeamRequest(team_id="no-such-team", team_alias="renamed"),
|
||||
http_request=_update_request_stub(),
|
||||
user_api_key_dict=_PROXY_ADMIN_CALLER,
|
||||
)
|
||||
|
||||
assert str(denied.value.code) == "403"
|
||||
assert "do not have access to this team" in str(denied.value.message)
|
||||
assert "no-such-team" not in str(denied.value.message)
|
||||
assert str(missing.value.code) == "404"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_resolve_team_access_ranks_proxy_admin_then_org_admin_then_team_admin():
|
||||
from litellm.proxy.management_endpoints.team_endpoints import _resolve_team_access
|
||||
|
||||
team = LiteLLM_TeamTable(
|
||||
team_id="team-1",
|
||||
organization_id="org-1",
|
||||
members_with_roles=[Member(user_id="team-admin", role="admin")],
|
||||
)
|
||||
roster_admin = UserAPIKeyAuth(user_role=LitellmUserRoles.INTERNAL_USER, user_id="team-admin")
|
||||
outsider = UserAPIKeyAuth(user_role=LitellmUserRoles.INTERNAL_USER, user_id="someone-else")
|
||||
org_lookup = AsyncMock(return_value=False)
|
||||
|
||||
with patch("litellm.proxy.management_endpoints.team_endpoints._is_user_org_admin_for_team", org_lookup): # test-quality-ok: the org-admin lookup needs a real prisma client this file's MagicMock cannot provide
|
||||
assert await _resolve_team_access(team_obj=team, user_api_key_dict=_PROXY_ADMIN_CALLER) == "proxy_admin"
|
||||
assert org_lookup.await_count == 0
|
||||
assert await _resolve_team_access(team_obj=team, user_api_key_dict=roster_admin) == "team_admin"
|
||||
assert await _resolve_team_access(team_obj=team, user_api_key_dict=outsider) is None
|
||||
org_lookup.return_value = True
|
||||
assert await _resolve_team_access(team_obj=team, user_api_key_dict=roster_admin) == "org_admin"
|
||||
|
||||
|
||||
_ROSTER_ADMIN_CALLER = UserAPIKeyAuth(user_role=LitellmUserRoles.INTERNAL_USER, user_id="admin-1")
|
||||
_MEMBER_CALLER = UserAPIKeyAuth(user_role=LitellmUserRoles.INTERNAL_USER, user_id="member-1")
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"caller, org_admin, enabled_fields, expected",
|
||||
[
|
||||
pytest.param(_PROXY_ADMIN_CALLER, False, (), {"kind": "unrestricted"}, id="proxy-admin"),
|
||||
pytest.param(
|
||||
UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY, user_id="viewer"),
|
||||
False,
|
||||
("tpm_limit",),
|
||||
{"kind": "none"},
|
||||
id="proxy-admin-viewer",
|
||||
),
|
||||
pytest.param(_ROSTER_ADMIN_CALLER, True, (), {"kind": "unrestricted"}, id="org-admin-who-is-also-team-admin"),
|
||||
pytest.param(_ROSTER_ADMIN_CALLER, False, (), {"kind": "team_admin_disabled"}, id="team-admin-nothing-enabled"),
|
||||
pytest.param(
|
||||
_ROSTER_ADMIN_CALLER,
|
||||
False,
|
||||
("tpm_limit",),
|
||||
{"kind": "team_admin", "editable_fields": ["tpm_limit"]},
|
||||
id="team-admin-field-enabled",
|
||||
),
|
||||
pytest.param(_MEMBER_CALLER, False, ("tpm_limit",), {"kind": "none"}, id="plain-member"),
|
||||
],
|
||||
)
|
||||
@pytest.mark.asyncio
|
||||
async def test_team_info_reports_what_the_caller_may_edit(caller, org_admin, enabled_fields, expected):
|
||||
"""The dashboard gates its edit form on this field instead of guessing the caller's role from the org list,
|
||||
which is premium-gated and can be empty for a dual-role org admin."""
|
||||
from fastapi import Request
|
||||
|
||||
from litellm.proxy.management_endpoints import team_endpoints
|
||||
|
||||
team_row = LiteLLM_TeamTable(
|
||||
team_id="team-1",
|
||||
organization_id="org-1",
|
||||
members_with_roles=[Member(user_id="admin-1", role="admin"), Member(user_id="member-1", role="user")],
|
||||
)
|
||||
mock_prisma = MagicMock()
|
||||
mock_prisma.db.litellm_teamtable.find_unique = AsyncMock(return_value=team_row)
|
||||
mock_prisma.db.litellm_usertable.find_many = AsyncMock(return_value=[])
|
||||
mock_prisma.get_data = AsyncMock(return_value=[])
|
||||
|
||||
with (
|
||||
patch("litellm.proxy.proxy_server.prisma_client", mock_prisma), # test-quality-ok: no seam on team_info
|
||||
patch.object(team_endpoints, "get_all_team_memberships", AsyncMock(return_value=[])), # test-quality-ok: no seam on team_info
|
||||
patch.object( # test-quality-ok: the org-admin lookup needs a real prisma client this file's MagicMock cannot provide
|
||||
team_endpoints, "_is_user_org_admin_for_team", AsyncMock(return_value=org_admin)
|
||||
),
|
||||
_team_admin_may_edit(*enabled_fields),
|
||||
):
|
||||
response = await team_endpoints.team_info(
|
||||
http_request=MagicMock(spec=Request),
|
||||
team_id="team-1",
|
||||
user_api_key_dict=caller,
|
||||
)
|
||||
|
||||
assert response["team_info"].caller_edit_access.model_dump(mode="json") == expected
|
||||
|
|
|
|||
|
|
@ -1983,6 +1983,144 @@ class TestBedrockAgentRuntimePassthroughToggle:
|
|||
create_route.assert_called_once()
|
||||
|
||||
|
||||
class TestBedrockAgentRuntimePassthroughVirtualKeyLeak:
|
||||
|
||||
VKEY: Final = "sk-litellm-victim-key"
|
||||
MASTER_KEY: Final = "sk-master-1234"
|
||||
ENDPOINT: Final = "knowledgebases/KB1234567/retrieve"
|
||||
AMBIENT_AWS_ENV: Final = (
|
||||
"AWS_BEARER_TOKEN_BEDROCK",
|
||||
"AWS_SESSION_TOKEN",
|
||||
"AWS_SESSION_NAME",
|
||||
"AWS_PROFILE_NAME",
|
||||
"AWS_ROLE_NAME",
|
||||
"AWS_WEB_IDENTITY_TOKEN",
|
||||
"AWS_STS_ENDPOINT",
|
||||
"AWS_EXTERNAL_ID",
|
||||
)
|
||||
|
||||
async def _upstream_headers(self, monkeypatch, headers: list[tuple[bytes, bytes]]) -> dict:
|
||||
from litellm.proxy.pass_through_endpoints.pass_through_endpoints import HttpPassThroughEndpointHelpers
|
||||
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.master_key", self.MASTER_KEY)
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", {})
|
||||
for ambient in self.AMBIENT_AWS_ENV:
|
||||
monkeypatch.delenv(ambient, raising=False)
|
||||
monkeypatch.setenv("AWS_ACCESS_KEY_ID", "ak")
|
||||
monkeypatch.setenv("AWS_SECRET_ACCESS_KEY", "sk")
|
||||
monkeypatch.setenv("AWS_REGION_NAME", "us-east-1")
|
||||
caller: Final = UserAPIKeyAuth(api_key=self.VKEY)
|
||||
|
||||
async def receive():
|
||||
return {"type": "http.request", "body": b'{"retrievalQuery": {"text": "hi"}}', "more_body": False}
|
||||
|
||||
request: Final = Request(
|
||||
{
|
||||
"type": "http",
|
||||
"method": "POST",
|
||||
"path": f"/bedrock/{self.ENDPOINT}",
|
||||
"headers": headers,
|
||||
"query_string": b"",
|
||||
},
|
||||
receive=receive,
|
||||
)
|
||||
captured: dict = {}
|
||||
|
||||
def fake_create_pass_through_route(**kwargs):
|
||||
captured.update(kwargs)
|
||||
return AsyncMock(return_value={"status": "success"})
|
||||
|
||||
module: Final = "litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints"
|
||||
with (
|
||||
patch(f"{module}.create_request_copy", Mock()),
|
||||
patch(f"{module}.create_pass_through_route", side_effect=fake_create_pass_through_route),
|
||||
):
|
||||
await bedrock_proxy_route(
|
||||
endpoint=self.ENDPOINT,
|
||||
request=request,
|
||||
fastapi_response=Response(),
|
||||
user_api_key_dict=caller,
|
||||
)
|
||||
return HttpPassThroughEndpointHelpers.forward_headers_from_request(
|
||||
request_headers=dict(request.headers),
|
||||
headers=dict(captured["custom_headers"] or {}),
|
||||
forward_headers=captured.get("_forward_headers", False),
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _blob(upstream: dict) -> str:
|
||||
return " ".join(f"{name}:{value}" for name, value in upstream.items())
|
||||
|
||||
@staticmethod
|
||||
def _names_matching(upstream: dict, lowercase_name: str) -> list[str]:
|
||||
return [name for name in upstream if name.lower() == lowercase_name]
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
"header_name", ["x-api-key", "x-litellm-api-key", "api-key", "x-goog-api-key", "ocp-apim-subscription-key"]
|
||||
)
|
||||
async def test_virtual_key_in_a_credential_header_never_reaches_aws(self, monkeypatch, header_name: str):
|
||||
upstream: Final = await self._upstream_headers(
|
||||
monkeypatch,
|
||||
[
|
||||
(header_name.encode(), self.VKEY.encode()),
|
||||
(b"content-type", b"application/json"),
|
||||
(b"x-request-id", b"trace-1"),
|
||||
],
|
||||
)
|
||||
|
||||
assert self.VKEY not in self._blob(upstream)
|
||||
assert self._names_matching(upstream, header_name) == []
|
||||
assert upstream["x-request-id"] == "trace-1", "a benign caller header still reaches AWS"
|
||||
assert upstream["Authorization"].startswith("AWS4-HMAC-SHA256")
|
||||
assert self._names_matching(upstream, "content-type") == ["Content-Type"], "the signed header is the only one"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_credential_headers_are_dropped_by_name_even_when_they_carry_someone_elses_key(self, monkeypatch):
|
||||
other_key: Final = "sk-other-tenant-key"
|
||||
upstream: Final = await self._upstream_headers(
|
||||
monkeypatch,
|
||||
[
|
||||
(b"x-api-key", other_key.encode()),
|
||||
(b"x-litellm-api-key", other_key.encode()),
|
||||
(b"x-request-id", b"trace-3"),
|
||||
],
|
||||
)
|
||||
|
||||
assert other_key not in self._blob(upstream)
|
||||
assert self._names_matching(upstream, "x-api-key") == []
|
||||
assert self._names_matching(upstream, "x-litellm-api-key") == []
|
||||
assert upstream["x-request-id"] == "trace-3"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_virtual_key_in_authorization_bearer_is_replaced_by_the_sigv4_signature(self, monkeypatch):
|
||||
upstream: Final = await self._upstream_headers(
|
||||
monkeypatch,
|
||||
[(b"authorization", f"Bearer {self.VKEY}".encode()), (b"content-type", b"application/json")],
|
||||
)
|
||||
|
||||
assert self.VKEY not in self._blob(upstream)
|
||||
assert self._names_matching(upstream, "authorization") == ["Authorization"]
|
||||
assert upstream["Authorization"].startswith("AWS4-HMAC-SHA256")
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_authenticated_secrets_in_any_other_header_never_reach_aws(self, monkeypatch):
|
||||
upstream: Final = await self._upstream_headers(
|
||||
monkeypatch,
|
||||
[
|
||||
(b"x-api-key", self.VKEY.encode()),
|
||||
(b"x-forwarded-key", self.VKEY.encode()),
|
||||
(b"x-operator-token", self.MASTER_KEY.encode()),
|
||||
(b"x-request-id", b"trace-2"),
|
||||
],
|
||||
)
|
||||
|
||||
assert self.VKEY not in self._blob(upstream) and self.MASTER_KEY not in self._blob(upstream)
|
||||
assert self._names_matching(upstream, "x-forwarded-key") == []
|
||||
assert self._names_matching(upstream, "x-operator-token") == []
|
||||
assert upstream["x-request-id"] == "trace-2"
|
||||
|
||||
|
||||
class TestLLMPassthroughFactoryProxyRoute:
|
||||
@pytest.mark.asyncio
|
||||
async def test_llm_passthrough_factory_proxy_route_success(self):
|
||||
|
|
|
|||
|
|
@ -3938,3 +3938,58 @@ async def test_ProxyConfig__init_guardrails_in_db_skips_only_the_unloadable_row(
|
|||
|
||||
assert sorted(handler.IN_MEMORY_GUARDRAILS) == ["first", "last"]
|
||||
assert handler.reconciled_with == [{"first", "broken", "last"}]
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# add_deployment: UI settings convergence
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_add_deployment_re_reads_ui_settings_so_other_pods_converge(monkeypatch):
|
||||
"""The periodic config reload picks up a UI setting written through another pod.
|
||||
|
||||
Startup used to be the only read, so a proxy admin flipping a runtime flag reached the pod
|
||||
that served the PATCH and nowhere else until every other pod restarted.
|
||||
"""
|
||||
general_settings: Dict[str, Any] = {"allow_agents_for_team_admins": False}
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", general_settings)
|
||||
|
||||
prisma_client = MagicMock()
|
||||
prisma_client.db.litellm_config.find_many = AsyncMock(return_value=[])
|
||||
prisma_client.db.litellm_config.find_first = AsyncMock(return_value=None)
|
||||
prisma_client.db.litellm_credentialstable.find_many = AsyncMock(return_value=[])
|
||||
prisma_client.db.litellm_uisettings.find_unique = AsyncMock(
|
||||
return_value=SimpleNamespace(
|
||||
ui_settings=json.dumps({"allow_agents_for_team_admins": True, "enable_chat_ui": False})
|
||||
)
|
||||
)
|
||||
|
||||
config = ProxyConfig()
|
||||
config._should_load_db_object = MagicMock(return_value=False)
|
||||
config._init_non_llm_objects_in_db = AsyncMock()
|
||||
|
||||
await config.add_deployment(prisma_client=prisma_client, proxy_logging_obj=MagicMock())
|
||||
|
||||
prisma_client.db.litellm_uisettings.find_unique.assert_awaited_once_with(where={"id": "ui_settings"})
|
||||
assert general_settings["allow_agents_for_team_admins"] is True
|
||||
assert "enable_chat_ui" not in general_settings
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_add_deployment_syncs_ui_settings_even_when_the_model_reconcile_fails(monkeypatch):
|
||||
"""A broken model reconcile must not strand every pod on stale settings."""
|
||||
general_settings: Dict[str, Any] = {"allow_agents_for_team_admins": False}
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", general_settings)
|
||||
|
||||
prisma_client = MagicMock()
|
||||
prisma_client.db.litellm_uisettings.find_unique = AsyncMock(
|
||||
return_value=SimpleNamespace(ui_settings={"allow_agents_for_team_admins": True})
|
||||
)
|
||||
|
||||
config = ProxyConfig()
|
||||
config._should_load_db_object = MagicMock(side_effect=RuntimeError("db down"))
|
||||
|
||||
await config.add_deployment(prisma_client=prisma_client, proxy_logging_obj=MagicMock())
|
||||
|
||||
assert general_settings["allow_agents_for_team_admins"] is True
|
||||
|
|
|
|||
|
|
@ -9,7 +9,7 @@ Pins (PR2):
|
|||
|
||||
from __future__ import annotations
|
||||
|
||||
from unittest.mock import MagicMock
|
||||
from unittest.mock import AsyncMock, MagicMock
|
||||
|
||||
import pytest
|
||||
|
||||
|
|
@ -128,7 +128,6 @@ def test_v1_model_info_no_model_list_error(client, auth_as, null_router, path):
|
|||
assert "LLM Model List not loaded" in response.text
|
||||
|
||||
|
||||
|
||||
def test_get_proxy_model_info_surfaces_supports_parallel_function_calling(local_model_cost_map):
|
||||
"""``GET /v1/model/info`` enriches each deployment through ``_get_proxy_model_info``; a registry
|
||||
entry declaring parallel function calling must land in ``model_info`` instead of null."""
|
||||
|
|
@ -161,9 +160,7 @@ def test_v1_model_info_star_wildcard_filter_keeps_provider_expansion(monkeypatch
|
|||
router.get_model_list = MagicMock(return_value=[deployment])
|
||||
monkeypatch.setattr(model_checks, "get_provider_models", fake_get_provider_models)
|
||||
|
||||
expanded_deployments = proxy_server.expand_wildcard_deployments_for_model_info(
|
||||
[deployment]
|
||||
)
|
||||
expanded_deployments = proxy_server.expand_wildcard_deployments_for_model_info([deployment])
|
||||
allowed_model_names = proxy_server._get_v1_model_info_allowed_model_names(
|
||||
user_api_key_dict=UserAPIKeyAuth(
|
||||
api_key="sk-test",
|
||||
|
|
@ -308,6 +305,80 @@ def test_model_group_info_invalid_method(client, auth_as, null_router):
|
|||
assert len(response.content) > 0
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def model_group_info_router(monkeypatch):
|
||||
from litellm.types.proxy.management_endpoints.model_management_endpoints import ModelGroupInfoProxy
|
||||
|
||||
model_names = ["gpt-4", "claude-3"]
|
||||
router = MagicMock()
|
||||
router.get_model_names.return_value = model_names
|
||||
router.get_model_access_groups.return_value = {}
|
||||
router.get_model_list.return_value = []
|
||||
|
||||
def model_group_info(*, llm_router, all_models_str, model_group):
|
||||
return [ModelGroupInfoProxy(model_group=name, providers=[]) for name in all_models_str]
|
||||
|
||||
async def append_agents_to_model_group(*, model_groups, user_api_key_dict):
|
||||
return model_groups
|
||||
|
||||
monkeypatch.setattr(proxy_server, "llm_router", router)
|
||||
monkeypatch.setattr(proxy_server, "llm_model_list", [{"model_name": name} for name in model_names])
|
||||
monkeypatch.setattr(proxy_server, "user_model", None)
|
||||
monkeypatch.setattr(proxy_server, "general_settings", {})
|
||||
monkeypatch.setattr(proxy_server, "prisma_client", None)
|
||||
monkeypatch.setattr(proxy_server, "proxy_logging_obj", None)
|
||||
monkeypatch.setattr(proxy_server, "user_api_key_cache", None)
|
||||
monkeypatch.setattr(proxy_server, "_get_model_group_info", model_group_info)
|
||||
|
||||
from litellm.proxy.agent_endpoints import model_list_helpers
|
||||
|
||||
monkeypatch.setattr(
|
||||
model_list_helpers,
|
||||
"append_agents_to_model_group",
|
||||
AsyncMock(side_effect=append_agents_to_model_group),
|
||||
)
|
||||
return router
|
||||
|
||||
|
||||
@pytest.mark.parametrize("admin_role", ["proxy_admin", "proxy_admin_viewer"])
|
||||
def test_model_group_info_proxy_admin_ignores_key_model_restriction(
|
||||
client, auth_as, model_group_info_router, admin_role
|
||||
):
|
||||
from litellm.proxy._types import LitellmUserRoles
|
||||
|
||||
with auth_as(LitellmUserRoles(admin_role), models=["no-default-models"]):
|
||||
response = client.get("/model_group/info")
|
||||
|
||||
assert response.status_code == 200
|
||||
assert [model["model_group"] for model in response.json()["data"]] == ["gpt-4", "claude-3"]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("admin_role", ["proxy_admin", "proxy_admin_viewer"])
|
||||
def test_model_group_info_proxy_admin_expands_wildcard_deployments(client, auth_as, model_group_info_router, admin_role):
|
||||
from litellm.proxy._types import LitellmUserRoles
|
||||
from litellm.proxy.auth.model_checks import get_known_models_from_wildcard
|
||||
|
||||
model_group_info_router.get_model_names.return_value = ["gpt-4", "anthropic/*"]
|
||||
known_anthropic_models = get_known_models_from_wildcard(wildcard_model="anthropic/*")
|
||||
assert known_anthropic_models
|
||||
|
||||
with auth_as(LitellmUserRoles(admin_role), models=["no-default-models"]):
|
||||
response = client.get("/model_group/info")
|
||||
|
||||
assert response.status_code == 200
|
||||
assert [model["model_group"] for model in response.json()["data"]] == ["gpt-4", *known_anthropic_models]
|
||||
|
||||
|
||||
def test_model_group_info_internal_user_key_model_restriction_applies(client, auth_as, model_group_info_router):
|
||||
from litellm.proxy._types import LitellmUserRoles
|
||||
|
||||
with auth_as(LitellmUserRoles.INTERNAL_USER, models=["gpt-4"]):
|
||||
response = client.get("/model_group/info")
|
||||
|
||||
assert response.status_code == 200
|
||||
assert [model["model_group"] for model in response.json()["data"]] == ["gpt-4"]
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# GET /v2/model/info?exclude_auto_routers
|
||||
# ---------------------------------------------------------------------------
|
||||
|
|
@ -399,14 +470,10 @@ def test_v2_model_info_exclude_auto_routers_shrinks_total_count(client, auth_as,
|
|||
assert len(payload["data"]) == payload["total_count"]
|
||||
|
||||
|
||||
def test_v2_model_info_exclude_auto_routers_paginates_over_the_filtered_set(
|
||||
client, auth_as, mixed_auto_router_router
|
||||
):
|
||||
def test_v2_model_info_exclude_auto_routers_paginates_over_the_filtered_set(client, auth_as, mixed_auto_router_router):
|
||||
"""Page size applies to the filtered list, so no page silently comes back short."""
|
||||
with auth_as():
|
||||
response = client.get(
|
||||
"/v2/model/info", params={"exclude_auto_routers": "true", "page": 1, "size": 1}
|
||||
)
|
||||
response = client.get("/v2/model/info", params={"exclude_auto_routers": "true", "page": 1, "size": 1})
|
||||
payload = response.json()
|
||||
assert payload["total_count"] == 2
|
||||
assert payload["total_pages"] == 2
|
||||
|
|
|
|||
|
|
@ -3266,3 +3266,174 @@ class TestPtuCostAttributionUISetting:
|
|||
assert response.status_code == 400
|
||||
assert "enable_ptu_cost_attribution" in str(response.json()["detail"])
|
||||
assert not mock_prisma.db.litellm_uisettings.upsert.called
|
||||
|
||||
|
||||
class TestTeamAdminEditableTeamFieldsSetting:
|
||||
"""team_admin_editable_team_fields: the proxy-wide allow-list update_team applies to team admins."""
|
||||
|
||||
def _as_proxy_admin(self, monkeypatch):
|
||||
from unittest.mock import AsyncMock, MagicMock
|
||||
|
||||
from litellm.proxy._types import UserAPIKeyAuth
|
||||
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
|
||||
|
||||
app.dependency_overrides[user_api_key_auth] = lambda: UserAPIKeyAuth(
|
||||
user_id="test-user-123",
|
||||
user_role=LitellmUserRoles.PROXY_ADMIN,
|
||||
)
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.store_model_in_db", True)
|
||||
mock_prisma = MagicMock()
|
||||
mock_prisma.db.litellm_uisettings.upsert = AsyncMock()
|
||||
mock_prisma.db.litellm_uisettings.find_unique = AsyncMock(return_value=None)
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma)
|
||||
return mock_prisma
|
||||
|
||||
def test_patch_rejects_field_names_the_proxy_does_not_support(self, monkeypatch):
|
||||
mock_prisma = self._as_proxy_admin(monkeypatch)
|
||||
monkeypatch.setattr(
|
||||
"litellm.proxy.ui_crud_endpoints.proxy_setting_endpoints.SUPPORTED_TEAM_ADMIN_EDITABLE_TEAM_FIELDS",
|
||||
frozenset({"tpm_limit"}),
|
||||
)
|
||||
|
||||
try:
|
||||
response = client.patch(
|
||||
"/update/ui_settings",
|
||||
json={"team_admin_editable_team_fields": ["tpm_limit", "blocked", "organization_id"]},
|
||||
)
|
||||
finally:
|
||||
app.dependency_overrides.clear()
|
||||
|
||||
assert response.status_code == 400
|
||||
detail = response.json()["detail"]["error"]
|
||||
assert "['blocked', 'organization_id']" in detail
|
||||
assert "['tpm_limit']" in detail
|
||||
assert not mock_prisma.db.litellm_uisettings.upsert.called
|
||||
|
||||
def test_patch_rejects_a_non_list_value(self, monkeypatch):
|
||||
self._as_proxy_admin(monkeypatch)
|
||||
|
||||
try:
|
||||
response = client.patch("/update/ui_settings", json={"team_admin_editable_team_fields": "tpm_limit"})
|
||||
finally:
|
||||
app.dependency_overrides.clear()
|
||||
|
||||
assert response.status_code == 422
|
||||
|
||||
def test_patch_persists_and_syncs_the_list_to_general_settings(self, monkeypatch):
|
||||
mock_prisma = self._as_proxy_admin(monkeypatch)
|
||||
general_settings: dict = {"team_admin_editable_team_fields": []}
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", general_settings)
|
||||
|
||||
try:
|
||||
response = client.patch("/update/ui_settings", json={"team_admin_editable_team_fields": ["tpm_limit"]})
|
||||
finally:
|
||||
app.dependency_overrides.clear()
|
||||
|
||||
assert response.status_code == 200
|
||||
stored = json.loads(mock_prisma.db.litellm_uisettings.upsert.call_args.kwargs["data"]["create"]["ui_settings"])
|
||||
assert stored["team_admin_editable_team_fields"] == ["tpm_limit"]
|
||||
assert general_settings["team_admin_editable_team_fields"] == ["tpm_limit"]
|
||||
|
||||
def test_patch_with_an_empty_list_turns_team_admin_editing_off_again(self, monkeypatch):
|
||||
mock_prisma = self._as_proxy_admin(monkeypatch)
|
||||
general_settings: dict = {"team_admin_editable_team_fields": ["tpm_limit"]}
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", general_settings)
|
||||
|
||||
try:
|
||||
response = client.patch("/update/ui_settings", json={"team_admin_editable_team_fields": []})
|
||||
finally:
|
||||
app.dependency_overrides.clear()
|
||||
|
||||
assert response.status_code == 200
|
||||
stored = json.loads(mock_prisma.db.litellm_uisettings.upsert.call_args.kwargs["data"]["create"]["ui_settings"])
|
||||
assert stored["team_admin_editable_team_fields"] == []
|
||||
assert general_settings["team_admin_editable_team_fields"] == []
|
||||
|
||||
def test_get_reports_the_stored_list_and_advertises_supported_fields(self, mock_auth, monkeypatch):
|
||||
from unittest.mock import AsyncMock, MagicMock
|
||||
|
||||
mock_prisma = MagicMock()
|
||||
mock_db_record = MagicMock()
|
||||
mock_db_record.ui_settings = {"team_admin_editable_team_fields": ["tpm_limit"]}
|
||||
mock_prisma.db.litellm_uisettings.find_unique = AsyncMock(return_value=mock_db_record)
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma)
|
||||
general_settings: dict = {}
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", general_settings)
|
||||
|
||||
response = client.get("/get/ui_settings")
|
||||
|
||||
assert response.status_code == 200
|
||||
data = response.json()
|
||||
assert data["values"]["team_admin_editable_team_fields"] == ["tpm_limit"]
|
||||
assert general_settings["team_admin_editable_team_fields"] == ["tpm_limit"]
|
||||
field_schema = data["field_schema"]["properties"]["team_admin_editable_team_fields"]
|
||||
assert field_schema["type"] == "array"
|
||||
assert field_schema["items"]["type"] == "string"
|
||||
assert "tpm_limit" in field_schema["items"]["enum"]
|
||||
|
||||
|
||||
class TestSyncUiSettingsToGeneralSettings:
|
||||
"""The DB re-read each pod runs on startup and on every config reload."""
|
||||
|
||||
def _sync(self):
|
||||
from litellm.proxy.ui_crud_endpoints.proxy_setting_endpoints import (
|
||||
sync_ui_settings_to_general_settings,
|
||||
)
|
||||
|
||||
return sync_ui_settings_to_general_settings
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_applies_runtime_flags_and_leaves_other_ui_settings_alone(self, monkeypatch):
|
||||
from unittest.mock import AsyncMock, MagicMock
|
||||
|
||||
general_settings: dict = {"allow_agents_for_team_admins": False}
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", general_settings)
|
||||
mock_prisma = MagicMock()
|
||||
record = MagicMock()
|
||||
record.ui_settings = json.dumps(
|
||||
{
|
||||
"allow_agents_for_team_admins": True,
|
||||
"team_admin_editable_team_fields": ["tpm_limit"],
|
||||
"enable_chat_ui": False,
|
||||
}
|
||||
)
|
||||
mock_prisma.db.litellm_uisettings.find_unique = AsyncMock(return_value=record)
|
||||
|
||||
applied = await self._sync()(mock_prisma)
|
||||
|
||||
assert dict(applied) == {
|
||||
"allow_agents_for_team_admins": True,
|
||||
"team_admin_editable_team_fields": ["tpm_limit"],
|
||||
}
|
||||
assert general_settings["allow_agents_for_team_admins"] is True
|
||||
assert general_settings["team_admin_editable_team_fields"] == ["tpm_limit"]
|
||||
assert "enable_chat_ui" not in general_settings
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_reads_a_row_the_prisma_client_already_deserialized(self, monkeypatch):
|
||||
from unittest.mock import AsyncMock, MagicMock
|
||||
|
||||
general_settings: dict = {}
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", general_settings)
|
||||
mock_prisma = MagicMock()
|
||||
record = MagicMock()
|
||||
record.ui_settings = {"team_admin_editable_team_fields": ["rpm_limit"]}
|
||||
mock_prisma.db.litellm_uisettings.find_unique = AsyncMock(return_value=record)
|
||||
|
||||
await self._sync()(mock_prisma)
|
||||
|
||||
assert general_settings["team_admin_editable_team_fields"] == ["rpm_limit"]
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_without_a_stored_row_general_settings_is_left_untouched(self, monkeypatch):
|
||||
from unittest.mock import AsyncMock, MagicMock
|
||||
|
||||
general_settings: dict = {"allow_agents_for_team_admins": True}
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", general_settings)
|
||||
mock_prisma = MagicMock()
|
||||
mock_prisma.db.litellm_uisettings.find_unique = AsyncMock(return_value=None)
|
||||
|
||||
applied = await self._sync()(mock_prisma)
|
||||
|
||||
assert dict(applied) == {}
|
||||
assert general_settings == {"allow_agents_for_team_admins": True}
|
||||
|
|
|
|||
|
|
@ -115,13 +115,13 @@ def _respx_interceptable_httpx_client(monkeypatch):
|
|||
],
|
||||
)
|
||||
def test_resolver_opt_in_gates_openai_like_config(model_info, expected_type):
|
||||
config = _resolve_responses_api_provider_config("my-model", "custom_openai", model_info)
|
||||
config = _resolve_responses_api_provider_config("my-model", "custom_openai", model_info, None)
|
||||
assert type(config) is expected_type
|
||||
|
||||
|
||||
def test_resolver_keeps_native_provider_config():
|
||||
"""`openai/` already routes /v1/responses natively; the opt-in must not swap its config."""
|
||||
config = _resolve_responses_api_provider_config("gpt-4.1", "openai", OPT_IN)
|
||||
config = _resolve_responses_api_provider_config("gpt-4.1", "openai", OPT_IN, None)
|
||||
assert type(config) is OpenAIResponsesAPIConfig
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -17,6 +17,7 @@ import SCIMConfig from "@/components/SCIM";
|
|||
import LoggingSettings from "@/components/Settings/AdminSettings/LoggingSettings/LoggingSettings";
|
||||
import SSOSettings from "@/components/Settings/AdminSettings/SSOSettings/SSOSettings";
|
||||
import UISettings from "@/components/Settings/AdminSettings/UISettings/UISettings";
|
||||
import TeamAdminEditableFieldsSettings from "@/components/Settings/AdminSettings/UISettings/TeamAdminEditableFieldsSettings";
|
||||
import UserBannerSettings from "@/components/Settings/AdminSettings/UserBannerSettings/UserBannerSettings";
|
||||
import CyberArk from "@/components/Settings/AdminSettings/CyberArk/CyberArk";
|
||||
import HashicorpVault from "@/components/Settings/AdminSettings/HashicorpVault/HashicorpVault";
|
||||
|
|
@ -382,6 +383,7 @@ const AdminPanel: React.FC<AdminPanelProps> = ({ proxySettings }) => {
|
|||
children: (
|
||||
<div className="flex flex-col gap-4">
|
||||
<UISettings />
|
||||
<TeamAdminEditableFieldsSettings />
|
||||
<UserBannerSettings />
|
||||
</div>
|
||||
),
|
||||
|
|
|
|||
|
|
@ -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);
|
||||
});
|
||||
});
|
||||
|
|
|
|||
|
|
@ -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}
|
||||
/>
|
||||
</>
|
||||
|
|
|
|||
|
|
@ -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();
|
||||
});
|
||||
});
|
||||
|
|
|
|||
|
|
@ -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} />}
|
||||
|
|
|
|||
|
|
@ -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 : "";
|
||||
};
|
||||
|
|
@ -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());
|
||||
});
|
||||
});
|
||||
|
|
|
|||
|
|
@ -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}
|
||||
/>
|
||||
|
|
|
|||
|
|
@ -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"
|
||||
|
|
|
|||
|
|
@ -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: [] }],
|
||||
|
|
|
|||
|
|
@ -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 (
|
||||
|
|
|
|||
|
|
@ -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", () => {
|
||||
|
|
|
|||
|
|
@ -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}
|
||||
/>
|
||||
|
|
|
|||
|
|
@ -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]);
|
||||
}
|
||||
|
|
@ -0,0 +1,175 @@
|
|||
import { beforeEach, describe, expect, it, vi } from "vitest";
|
||||
|
||||
import { fireEvent, renderWithProviders, screen, waitFor } from "@/../tests/test-utils";
|
||||
import { toast } from "@/lib/toast";
|
||||
|
||||
import TeamAdminEditableFieldsSettings from "./TeamAdminEditableFieldsSettings";
|
||||
|
||||
const mockUseUISettings = vi.hoisted(() => vi.fn());
|
||||
const mockUseUpdateUISettings = vi.hoisted(() => vi.fn());
|
||||
|
||||
vi.mock("@/app/(dashboard)/hooks/useAuthorized", () => ({
|
||||
default: () => ({ accessToken: "test-token" }),
|
||||
}));
|
||||
|
||||
vi.mock("@/app/(dashboard)/hooks/uiSettings/useUISettings", () => ({
|
||||
useUISettings: mockUseUISettings,
|
||||
}));
|
||||
|
||||
vi.mock("@/app/(dashboard)/hooks/uiSettings/useUpdateUISettings", () => ({
|
||||
useUpdateUISettings: mockUseUpdateUISettings,
|
||||
}));
|
||||
|
||||
const TPM_LABEL = "Tokens per minute Limit (TPM)";
|
||||
|
||||
const mockSettings = (supported: readonly string[], enabled: readonly string[]) =>
|
||||
mockUseUISettings.mockReturnValue({
|
||||
isLoading: false,
|
||||
data: {
|
||||
field_schema: {
|
||||
properties: {
|
||||
team_admin_editable_team_fields: {
|
||||
description: "Fields a team admin may change",
|
||||
items: { type: "string", enum: supported },
|
||||
},
|
||||
},
|
||||
},
|
||||
values: { team_admin_editable_team_fields: enabled },
|
||||
},
|
||||
});
|
||||
|
||||
const mockSave = ({
|
||||
isPending = false,
|
||||
outcome = "success",
|
||||
}: {
|
||||
isPending?: boolean;
|
||||
outcome?: "success" | "error";
|
||||
}) => {
|
||||
const mutate = vi.fn((_settings: unknown, options: { onSuccess: () => void; onError: (error: Error) => void }) =>
|
||||
outcome === "success" ? options.onSuccess() : options.onError(new Error("save failed")),
|
||||
);
|
||||
mockUseUpdateUISettings.mockReturnValue({ mutate, isPending });
|
||||
return mutate;
|
||||
};
|
||||
|
||||
const saveButton = () => screen.getByRole("button", { name: "Save" });
|
||||
|
||||
describe("TeamAdminEditableFieldsSettings", () => {
|
||||
beforeEach(() => {
|
||||
vi.clearAllMocks();
|
||||
});
|
||||
|
||||
it("explains that nothing can be enabled when the proxy supports no fields", () => {
|
||||
mockSettings([], []);
|
||||
mockSave({});
|
||||
|
||||
renderWithProviders(<TeamAdminEditableFieldsSettings />);
|
||||
|
||||
expect(screen.getByText("Team admins cannot edit team settings")).toBeInTheDocument();
|
||||
expect(screen.getByText(/does not support enabling any team settings fields/)).toBeInTheDocument();
|
||||
expect(screen.queryByRole("checkbox")).not.toBeInTheDocument();
|
||||
expect(screen.queryByRole("button", { name: "Save" })).not.toBeInTheDocument();
|
||||
});
|
||||
|
||||
it("renders one checkbox per supported field, checked for the saved ones, with Save disabled until something changes", () => {
|
||||
mockSettings(["max_budget", "tpm_limit"], ["tpm_limit"]);
|
||||
mockSave({});
|
||||
|
||||
renderWithProviders(<TeamAdminEditableFieldsSettings />);
|
||||
|
||||
expect(screen.getByText("Team admin editable fields")).toBeInTheDocument();
|
||||
expect(screen.getByText("1 field enabled")).toBeInTheDocument();
|
||||
expect(screen.getByText("Fields a team admin may change")).toBeInTheDocument();
|
||||
expect(screen.getByRole("checkbox", { name: "max_budget" })).not.toBeChecked();
|
||||
expect(screen.getByRole("checkbox", { name: TPM_LABEL })).toBeChecked();
|
||||
expect(saveButton()).toBeDisabled();
|
||||
});
|
||||
|
||||
it("only saves a ticked field once Save is clicked", async () => {
|
||||
mockSettings(["max_budget", "tpm_limit"], ["tpm_limit"]);
|
||||
const mutate = mockSave({});
|
||||
|
||||
renderWithProviders(<TeamAdminEditableFieldsSettings />);
|
||||
fireEvent.click(screen.getByRole("checkbox", { name: "max_budget" }));
|
||||
|
||||
expect(screen.getByRole("checkbox", { name: "max_budget" })).toBeChecked();
|
||||
expect(mutate).not.toHaveBeenCalled();
|
||||
|
||||
fireEvent.click(saveButton());
|
||||
|
||||
await waitFor(() => expect(toast.success).toHaveBeenCalledWith("Team admin editable fields updated successfully"));
|
||||
expect(mutate).toHaveBeenCalledWith(
|
||||
{ team_admin_editable_team_fields: ["max_budget", "tpm_limit"] },
|
||||
expect.anything(),
|
||||
);
|
||||
expect(saveButton()).toBeDisabled();
|
||||
});
|
||||
|
||||
it("saves the list without an unticked field", async () => {
|
||||
mockSettings(["max_budget", "tpm_limit"], ["max_budget", "tpm_limit"]);
|
||||
const mutate = mockSave({});
|
||||
|
||||
renderWithProviders(<TeamAdminEditableFieldsSettings />);
|
||||
fireEvent.click(screen.getByRole("checkbox", { name: TPM_LABEL }));
|
||||
fireEvent.click(saveButton());
|
||||
|
||||
await waitFor(() => expect(mutate).toHaveBeenCalledTimes(1));
|
||||
expect(mutate).toHaveBeenCalledWith({ team_admin_editable_team_fields: ["max_budget"] }, expect.anything());
|
||||
});
|
||||
|
||||
it("disables Save again when the draft is ticked back to the saved list", () => {
|
||||
mockSettings(["tpm_limit"], []);
|
||||
mockSave({});
|
||||
|
||||
renderWithProviders(<TeamAdminEditableFieldsSettings />);
|
||||
fireEvent.click(screen.getByRole("checkbox", { name: TPM_LABEL }));
|
||||
|
||||
expect(saveButton()).toBeEnabled();
|
||||
|
||||
fireEvent.click(screen.getByRole("checkbox", { name: TPM_LABEL }));
|
||||
|
||||
expect(screen.getByRole("checkbox", { name: TPM_LABEL })).not.toBeChecked();
|
||||
expect(saveButton()).toBeDisabled();
|
||||
});
|
||||
|
||||
it("treats a saved list in another order, or with fields this proxy dropped, as the same selection", () => {
|
||||
mockSettings(["max_budget", "tpm_limit"], ["tpm_limit", "retired_field", "max_budget"]);
|
||||
mockSave({});
|
||||
|
||||
renderWithProviders(<TeamAdminEditableFieldsSettings />);
|
||||
|
||||
expect(screen.getByText("2 fields enabled")).toBeInTheDocument();
|
||||
|
||||
fireEvent.click(screen.getByRole("checkbox", { name: TPM_LABEL }));
|
||||
fireEvent.click(screen.getByRole("checkbox", { name: TPM_LABEL }));
|
||||
|
||||
expect(saveButton()).toBeDisabled();
|
||||
});
|
||||
|
||||
it("keeps the draft and shows the error when the save fails", async () => {
|
||||
mockSettings(["tpm_limit"], []);
|
||||
const mutate = mockSave({ outcome: "error" });
|
||||
|
||||
renderWithProviders(<TeamAdminEditableFieldsSettings />);
|
||||
fireEvent.click(screen.getByRole("checkbox", { name: TPM_LABEL }));
|
||||
fireEvent.click(saveButton());
|
||||
|
||||
await waitFor(() => expect(toast.fromError).toHaveBeenCalledTimes(1));
|
||||
expect(mutate).toHaveBeenCalledTimes(1);
|
||||
expect(toast.success).not.toHaveBeenCalled();
|
||||
expect(screen.getByRole("checkbox", { name: TPM_LABEL })).toBeChecked();
|
||||
expect(saveButton()).toBeEnabled();
|
||||
});
|
||||
|
||||
it("blocks ticking and saving while a save is in flight", () => {
|
||||
mockSettings(["tpm_limit"], []);
|
||||
const mutate = mockSave({ isPending: true });
|
||||
|
||||
renderWithProviders(<TeamAdminEditableFieldsSettings />);
|
||||
fireEvent.click(screen.getByRole("checkbox", { name: TPM_LABEL }));
|
||||
|
||||
expect(screen.getByRole("checkbox", { name: TPM_LABEL })).not.toBeChecked();
|
||||
expect(screen.getByRole("button", { name: "Saving..." })).toBeDisabled();
|
||||
expect(mutate).not.toHaveBeenCalled();
|
||||
});
|
||||
});
|
||||
|
|
@ -0,0 +1,138 @@
|
|||
"use client";
|
||||
|
||||
import { Controller } from "react-hook-form";
|
||||
import { z } from "zod/v4";
|
||||
|
||||
import { useUISettings } from "@/app/(dashboard)/hooks/uiSettings/useUISettings";
|
||||
import { useUpdateUISettings } from "@/app/(dashboard)/hooks/uiSettings/useUpdateUISettings";
|
||||
import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized";
|
||||
import {
|
||||
parseSupportedTeamAdminEditableFields,
|
||||
parseTeamAdminEditableFields,
|
||||
teamAdminFieldLabel,
|
||||
} from "@/components/team/teamAdminEditAccess";
|
||||
import { Badge } from "@/components/ui/badge";
|
||||
import { Button } from "@/components/ui/button";
|
||||
import { Card, CardContent, CardDescription, CardHeader, CardTitle } from "@/components/ui/card";
|
||||
import { Checkbox } from "@/components/ui/checkbox";
|
||||
import { Skeleton } from "@/components/ui/skeleton";
|
||||
import { useZodForm } from "@/lib/forms/useZodForm";
|
||||
import { toast } from "@/lib/toast";
|
||||
|
||||
const editableFieldsSchema = z.object({ team_admin_editable_team_fields: z.array(z.string()) });
|
||||
|
||||
type SaveEditableFields = ReturnType<typeof useUpdateUISettings>["mutate"];
|
||||
|
||||
export default function TeamAdminEditableFieldsSettings() {
|
||||
const { accessToken } = useAuthorized();
|
||||
const { data, isLoading } = useUISettings();
|
||||
const { mutate: saveSettings, isPending } = useUpdateUISettings(accessToken);
|
||||
const supportedFields = parseSupportedTeamAdminEditableFields(data?.field_schema);
|
||||
const savedFields = parseTeamAdminEditableFields(data?.values);
|
||||
const enabledFields = supportedFields.filter((field) => savedFields.includes(field));
|
||||
|
||||
return (
|
||||
<Card>
|
||||
<CardHeader>
|
||||
<div className="flex items-center gap-2">
|
||||
<CardTitle>Team admin editable fields</CardTitle>
|
||||
<Badge variant={enabledFields.length > 0 ? "secondary" : "outline"}>
|
||||
{enabledFields.length > 0
|
||||
? `${enabledFields.length} field${enabledFields.length !== 1 ? "s" : ""} enabled`
|
||||
: "Team admins cannot edit team settings"}
|
||||
</Badge>
|
||||
</div>
|
||||
<CardDescription>
|
||||
{data?.field_schema?.properties?.team_admin_editable_team_fields?.description ??
|
||||
"Team settings fields a team admin may change on the teams they administer."}
|
||||
</CardDescription>
|
||||
</CardHeader>
|
||||
<CardContent>
|
||||
{isLoading ? (
|
||||
<Skeleton className="h-16 w-full" />
|
||||
) : (
|
||||
<TeamAdminEditableFieldsForm
|
||||
key={enabledFields.join(",")}
|
||||
enabledFields={enabledFields}
|
||||
supportedFields={supportedFields}
|
||||
isPending={isPending}
|
||||
saveSettings={saveSettings}
|
||||
/>
|
||||
)}
|
||||
</CardContent>
|
||||
</Card>
|
||||
);
|
||||
}
|
||||
|
||||
interface TeamAdminEditableFieldsFormProps {
|
||||
enabledFields: readonly string[];
|
||||
supportedFields: readonly string[];
|
||||
isPending: boolean;
|
||||
saveSettings: SaveEditableFields;
|
||||
}
|
||||
|
||||
function TeamAdminEditableFieldsForm({
|
||||
enabledFields,
|
||||
supportedFields,
|
||||
isPending,
|
||||
saveSettings,
|
||||
}: TeamAdminEditableFieldsFormProps) {
|
||||
const form = useZodForm(editableFieldsSchema, {
|
||||
defaultValues: { team_admin_editable_team_fields: [...enabledFields] },
|
||||
});
|
||||
const submit = form.handleSubmit((values) =>
|
||||
saveSettings(values, {
|
||||
onSuccess: () => {
|
||||
form.reset(values);
|
||||
toast.success("Team admin editable fields updated successfully");
|
||||
},
|
||||
onError: (error) => {
|
||||
toast.fromError(error);
|
||||
},
|
||||
}),
|
||||
);
|
||||
|
||||
if (supportedFields.length === 0) {
|
||||
return (
|
||||
<p className="text-sm italic text-muted-foreground">
|
||||
This proxy version does not support enabling any team settings fields for team admins yet.
|
||||
</p>
|
||||
);
|
||||
}
|
||||
|
||||
return (
|
||||
<form onSubmit={(event) => void submit(event)} className="space-y-4">
|
||||
<Controller
|
||||
control={form.control}
|
||||
name="team_admin_editable_team_fields"
|
||||
render={({ field }) => (
|
||||
<div className="space-y-2">
|
||||
{supportedFields.map((name) => {
|
||||
const checkboxId = `team-admin-editable-${name}`;
|
||||
return (
|
||||
<label key={name} htmlFor={checkboxId} className="flex cursor-pointer items-center gap-2">
|
||||
<Checkbox
|
||||
id={checkboxId}
|
||||
checked={field.value.includes(name)}
|
||||
disabled={isPending}
|
||||
onCheckedChange={(checked) =>
|
||||
field.onChange(
|
||||
supportedFields.filter((item) => (item === name ? checked : field.value.includes(item))),
|
||||
)
|
||||
}
|
||||
/>
|
||||
<span className="text-sm text-foreground">{teamAdminFieldLabel(name)}</span>
|
||||
</label>
|
||||
);
|
||||
})}
|
||||
</div>
|
||||
)}
|
||||
/>
|
||||
<div className="flex justify-end">
|
||||
<Button type="submit" disabled={isPending || !form.formState.isDirty}>
|
||||
{isPending ? "Saving..." : "Save"}
|
||||
</Button>
|
||||
</div>
|
||||
</form>
|
||||
);
|
||||
}
|
||||
|
|
@ -145,7 +145,7 @@ function UserBannerSettingsForm({ persisted, isLoading, isPending, saveBanner }:
|
|||
</div>
|
||||
)}
|
||||
|
||||
<div>
|
||||
<div className="flex justify-end">
|
||||
<Button onClick={handleSave} disabled={isPending || messageMissing}>
|
||||
{isPending ? "Saving..." : "Save banner"}
|
||||
</Button>
|
||||
|
|
|
|||
|
|
@ -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";
|
||||
|
|
@ -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();
|
||||
});
|
||||
});
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -0,0 +1,85 @@
|
|||
import { describe, expect, it, vi } from "vitest";
|
||||
import userEvent from "@testing-library/user-event";
|
||||
|
||||
import { fireEvent, renderWithProviders, screen, waitFor } from "@/../tests/test-utils";
|
||||
|
||||
import TeamAdminSettingsForm from "./TeamAdminSettingsForm";
|
||||
|
||||
const renderForm = (editableFields: ReadonlySet<string>, overrides: { isSaving?: boolean } = {}) => {
|
||||
const onSave = vi.fn().mockResolvedValue(undefined);
|
||||
const onCancel = vi.fn();
|
||||
renderWithProviders(
|
||||
<TeamAdminSettingsForm
|
||||
initialValues={{ tpm_limit: 1000 }}
|
||||
editableFields={editableFields}
|
||||
isSaving={overrides.isSaving ?? false}
|
||||
onCancel={onCancel}
|
||||
onSave={onSave}
|
||||
/>,
|
||||
);
|
||||
return { onSave, onCancel };
|
||||
};
|
||||
|
||||
describe("TeamAdminSettingsForm", () => {
|
||||
it("shows the team's current TPM limit when the proxy lets team admins edit it", () => {
|
||||
renderForm(new Set(["tpm_limit"]));
|
||||
|
||||
expect(screen.getByLabelText("Tokens per minute Limit (TPM)")).toHaveValue(1000);
|
||||
});
|
||||
|
||||
it("hides the TPM limit when the proxy has not enabled it for team admins", () => {
|
||||
renderForm(new Set(["max_budget"]));
|
||||
|
||||
expect(screen.queryByLabelText("Tokens per minute Limit (TPM)")).not.toBeInTheDocument();
|
||||
});
|
||||
|
||||
it("saves the new TPM limit and nothing else", async () => {
|
||||
const user = userEvent.setup();
|
||||
const { onSave } = renderForm(new Set(["tpm_limit"]));
|
||||
|
||||
fireEvent.change(screen.getByLabelText("Tokens per minute Limit (TPM)"), { target: { value: "5000" } });
|
||||
await user.click(screen.getByRole("button", { name: /save changes/i }));
|
||||
|
||||
await waitFor(() => expect(onSave).toHaveBeenCalledWith({ tpm_limit: 5000 }));
|
||||
});
|
||||
|
||||
it("saves a cleared TPM limit as no limit", async () => {
|
||||
const user = userEvent.setup();
|
||||
const { onSave } = renderForm(new Set(["tpm_limit"]));
|
||||
|
||||
fireEvent.change(screen.getByLabelText("Tokens per minute Limit (TPM)"), { target: { value: "" } });
|
||||
await user.click(screen.getByRole("button", { name: /save changes/i }));
|
||||
|
||||
await waitFor(() => expect(onSave).toHaveBeenCalledWith({ tpm_limit: null }));
|
||||
});
|
||||
|
||||
it("keeps Save disabled until the TPM limit differs from the team's", () => {
|
||||
renderForm(new Set(["tpm_limit"]));
|
||||
const tpmInput = screen.getByLabelText("Tokens per minute Limit (TPM)");
|
||||
const save = screen.getByRole("button", { name: /save changes/i });
|
||||
|
||||
expect(save).toBeDisabled();
|
||||
fireEvent.change(tpmInput, { target: { value: "5000" } });
|
||||
expect(save).toBeEnabled();
|
||||
fireEvent.change(tpmInput, { target: { value: "1000" } });
|
||||
expect(save).toBeDisabled();
|
||||
});
|
||||
|
||||
it("closes without saving on cancel", async () => {
|
||||
const user = userEvent.setup();
|
||||
const { onSave, onCancel } = renderForm(new Set(["tpm_limit"]));
|
||||
|
||||
await user.click(screen.getByRole("button", { name: "Cancel" }));
|
||||
|
||||
expect(onCancel).toHaveBeenCalledTimes(1);
|
||||
expect(onSave).not.toHaveBeenCalled();
|
||||
});
|
||||
|
||||
it("locks both buttons while a save is in flight", () => {
|
||||
renderForm(new Set(["tpm_limit"]), { isSaving: true });
|
||||
fireEvent.change(screen.getByLabelText("Tokens per minute Limit (TPM)"), { target: { value: "5000" } });
|
||||
|
||||
expect(screen.getByRole("button", { name: "Cancel" })).toBeDisabled();
|
||||
expect(screen.getByRole("button", { name: /save changes/i })).toBeDisabled();
|
||||
});
|
||||
});
|
||||
|
|
@ -0,0 +1,69 @@
|
|||
"use client";
|
||||
|
||||
import { Save } from "lucide-react";
|
||||
import { useWatch } from "react-hook-form";
|
||||
import { z } from "zod/v4";
|
||||
|
||||
import { FormField } from "@/components/shared/form/FormField";
|
||||
import { Button } from "@/components/ui/button";
|
||||
import { FieldGroup } from "@/components/ui/field";
|
||||
import { UiLoadingSpinner } from "@/components/ui/ui-loading-spinner";
|
||||
import { useZodForm } from "@/lib/forms/useZodForm";
|
||||
|
||||
import NumericalInput from "../shared/numerical_input";
|
||||
import {
|
||||
teamAdminFieldLabel,
|
||||
teamAdminSettingsChanges,
|
||||
type TeamAdminSettingsChanges,
|
||||
type TeamAdminSettingsValues,
|
||||
} from "./teamAdminEditAccess";
|
||||
|
||||
const teamAdminSettingsSchema = z.object({
|
||||
tpm_limit: z.union([z.string(), z.number()]).nullish(),
|
||||
});
|
||||
|
||||
interface TeamAdminSettingsFormProps {
|
||||
initialValues: TeamAdminSettingsValues;
|
||||
editableFields: ReadonlySet<string>;
|
||||
isSaving: boolean;
|
||||
onCancel: () => void;
|
||||
onSave: (changes: TeamAdminSettingsChanges) => Promise<void>;
|
||||
}
|
||||
|
||||
export default function TeamAdminSettingsForm({
|
||||
initialValues,
|
||||
editableFields,
|
||||
isSaving,
|
||||
onCancel,
|
||||
onSave,
|
||||
}: TeamAdminSettingsFormProps) {
|
||||
const form = useZodForm(teamAdminSettingsSchema, { defaultValues: initialValues });
|
||||
const draft = useWatch({ control: form.control });
|
||||
const hasChanges = Object.keys(teamAdminSettingsChanges(draft, initialValues, editableFields)).length > 0;
|
||||
const submit = form.handleSubmit((values) => onSave(teamAdminSettingsChanges(values, initialValues, editableFields)));
|
||||
|
||||
return (
|
||||
<form onSubmit={(event) => void submit(event)}>
|
||||
<FieldGroup>
|
||||
<p className="text-sm text-muted-foreground">
|
||||
A proxy admin chose which settings team admins can change. Ask a proxy admin to change anything else.
|
||||
</p>
|
||||
{editableFields.has("tpm_limit") && (
|
||||
<FormField control={form.control} name="tpm_limit" label={teamAdminFieldLabel("tpm_limit")}>
|
||||
{({ ref, value, ...field }) => <NumericalInput {...field} ref={ref} value={value ?? ""} step={1} />}
|
||||
</FormField>
|
||||
)}
|
||||
</FieldGroup>
|
||||
|
||||
<div className="mt-6 flex items-center justify-end gap-2">
|
||||
<Button type="button" variant="outline" onClick={onCancel} disabled={isSaving}>
|
||||
Cancel
|
||||
</Button>
|
||||
<Button type="submit" disabled={isSaving || !hasChanges}>
|
||||
{isSaving ? <UiLoadingSpinner className="size-4" /> : <Save className="size-4" />}
|
||||
Save Changes
|
||||
</Button>
|
||||
</div>
|
||||
</form>
|
||||
);
|
||||
}
|
||||
|
|
@ -69,6 +69,10 @@ vi.mock("@/app/(dashboard)/hooks/teams/useTeamMetadataSchema", () => ({
|
|||
useTeamMetadataSchema: vi.fn(() => ({ data: [], isLoading: false })),
|
||||
}));
|
||||
|
||||
vi.mock("@/app/(dashboard)/hooks/uiSettings/useUISettings", () => ({
|
||||
useUISettings: vi.fn(),
|
||||
}));
|
||||
|
||||
vi.mock("@/app/(dashboard)/hooks/models/useModels", () => ({
|
||||
useAllProxyModels: vi.fn(),
|
||||
}));
|
||||
|
|
@ -228,6 +232,7 @@ import { useCurrentUser } from "@/app/(dashboard)/hooks/users/useCurrentUser";
|
|||
import { useMCPServers } from "@/app/(dashboard)/hooks/mcpServers/useMCPServers";
|
||||
import { useMCPToolsets } from "@/app/(dashboard)/hooks/mcpServers/useMCPToolsets";
|
||||
import { useAccessGroups } from "@/app/(dashboard)/hooks/accessGroups/useAccessGroups";
|
||||
import { useUISettings } from "@/app/(dashboard)/hooks/uiSettings/useUISettings";
|
||||
|
||||
const mockUseAllProxyModels = vi.mocked(useAllProxyModels);
|
||||
const mockUseKeys = vi.mocked(useKeys);
|
||||
|
|
@ -237,6 +242,7 @@ const mockUseCurrentUser = vi.mocked(useCurrentUser);
|
|||
const mockUseMCPServers = vi.mocked(useMCPServers);
|
||||
const mockUseMCPToolsets = vi.mocked(useMCPToolsets);
|
||||
const mockUseAccessGroups = vi.mocked(useAccessGroups);
|
||||
const mockUseUISettings = vi.mocked(useUISettings);
|
||||
|
||||
const createMockTeamData = (overrides = {}) => ({
|
||||
team_id: "123",
|
||||
|
|
@ -305,6 +311,10 @@ const seedDefaultMocks = () => {
|
|||
isLoading: false,
|
||||
isError: false,
|
||||
} as any);
|
||||
mockUseUISettings.mockReturnValue({
|
||||
data: { values: {} },
|
||||
isLoading: false,
|
||||
} as any);
|
||||
mockUseKeys.mockReturnValue({
|
||||
data: { keys: [], total_count: 0, current_page: 1, total_pages: 1 },
|
||||
isPending: false,
|
||||
|
|
@ -656,19 +666,9 @@ describe("TeamInfoView", () => {
|
|||
});
|
||||
});
|
||||
|
||||
it("shows edit tabs when the fetched team data marks the session user as team admin, even without the is_team_admin prop", async () => {
|
||||
it("shows edit tabs when the proxy reports the session user may edit, even without the is_team_admin prop", async () => {
|
||||
vi.mocked(networking.teamInfoCall).mockResolvedValue(
|
||||
createMockTeamData({
|
||||
members_with_roles: [
|
||||
{
|
||||
user_id: "user-1",
|
||||
user_email: "admin@test.com",
|
||||
role: "admin",
|
||||
spend: 0,
|
||||
budget_id: "budget1",
|
||||
},
|
||||
],
|
||||
}),
|
||||
createMockTeamData({ caller_edit_access: { kind: "team_admin_disabled" } }),
|
||||
);
|
||||
|
||||
renderWithProviders(<TeamInfoView {...defaultProps} is_team_admin={false} is_proxy_admin={false} />);
|
||||
|
|
@ -1863,6 +1863,92 @@ describe("TeamInfoView", () => {
|
|||
});
|
||||
});
|
||||
});
|
||||
|
||||
describe("team admin edit access", () => {
|
||||
const teamAdminProps = { ...defaultProps, is_proxy_admin: false, is_team_admin: true };
|
||||
|
||||
beforeEach(() => {
|
||||
authState.userRole = "Internal User";
|
||||
});
|
||||
|
||||
it("tells a team admin to ask a proxy admin when the proxy reports no team field is enabled for them", async () => {
|
||||
const user = userEvent.setup({ delay: null });
|
||||
vi.mocked(networking.teamInfoCall).mockResolvedValue(
|
||||
createMockTeamData({ caller_edit_access: { kind: "team_admin_disabled" } }),
|
||||
);
|
||||
|
||||
renderWithProviders(<TeamInfoView {...teamAdminProps} />);
|
||||
|
||||
await user.click(await screen.findByRole("tab", { name: "Settings" }));
|
||||
await user.click(await screen.findByRole("button", { name: /edit settings/i }));
|
||||
|
||||
expect(toast.error).toHaveBeenCalledWith("Team admins cannot edit team settings on this proxy", {
|
||||
description: "Ask a proxy admin to enable fields under Settings > UI > Team admin editable fields.",
|
||||
});
|
||||
expect(screen.queryByLabelText("Team Name")).not.toBeInTheDocument();
|
||||
expect(screen.getByRole("button", { name: /edit settings/i })).toBeInTheDocument();
|
||||
});
|
||||
|
||||
it("gives a team admin only the fields the proxy enabled and sends only those on save", async () => {
|
||||
const user = userEvent.setup({ delay: null });
|
||||
vi.mocked(networking.teamInfoCall).mockResolvedValue(
|
||||
createMockTeamData({
|
||||
tpm_limit: 1000,
|
||||
caller_edit_access: { kind: "team_admin", editable_fields: ["tpm_limit"] },
|
||||
}),
|
||||
);
|
||||
vi.mocked(networking.teamUpdateCall).mockResolvedValue({ data: {}, team_id: "123" } as any);
|
||||
|
||||
renderWithProviders(<TeamInfoView {...teamAdminProps} />);
|
||||
|
||||
await user.click(await screen.findByRole("tab", { name: "Settings" }));
|
||||
await user.click(await screen.findByRole("button", { name: /edit settings/i }));
|
||||
|
||||
const tpmInput = await screen.findByLabelText("Tokens per minute Limit (TPM)");
|
||||
expect(tpmInput).toHaveValue(1000);
|
||||
expect(screen.queryByLabelText("Team Name")).not.toBeInTheDocument();
|
||||
expect(screen.queryByLabelText("Requests per minute Limit (RPM)")).not.toBeInTheDocument();
|
||||
|
||||
fireEvent.change(tpmInput, { target: { value: "5000" } });
|
||||
await user.click(screen.getByRole("button", { name: /save changes/i }));
|
||||
|
||||
await waitFor(() => expect(networking.teamUpdateCall).toHaveBeenCalledTimes(1));
|
||||
expect(vi.mocked(networking.teamUpdateCall).mock.calls[0][1]).toStrictEqual({ team_id: "123", tpm_limit: 5000 });
|
||||
expect(toast.success).toHaveBeenCalledWith("Team settings updated successfully");
|
||||
expect(toast.error).not.toHaveBeenCalled();
|
||||
});
|
||||
|
||||
it("opens the form when the proxy reports unrestricted access although the props only mark a team admin", async () => {
|
||||
const user = userEvent.setup({ delay: null });
|
||||
vi.mocked(networking.teamInfoCall).mockResolvedValue(
|
||||
createMockTeamData({ caller_edit_access: { kind: "unrestricted" } }),
|
||||
);
|
||||
|
||||
renderWithProviders(<TeamInfoView {...teamAdminProps} />);
|
||||
|
||||
await user.click(await screen.findByRole("tab", { name: "Settings" }));
|
||||
await user.click(await screen.findByRole("button", { name: /edit settings/i }));
|
||||
|
||||
expect(await screen.findByLabelText("Team Name")).toBeInTheDocument();
|
||||
expect(toast.error).not.toHaveBeenCalled();
|
||||
});
|
||||
|
||||
it("never gates a proxy admin on the team admin field list", async () => {
|
||||
authState.userRole = "Admin";
|
||||
const user = userEvent.setup({ delay: null });
|
||||
vi.mocked(networking.teamInfoCall).mockResolvedValue(
|
||||
createMockTeamData({ caller_edit_access: { kind: "unrestricted" } }),
|
||||
);
|
||||
|
||||
renderWithProviders(<TeamInfoView {...defaultProps} />);
|
||||
|
||||
await user.click(await screen.findByRole("tab", { name: "Settings" }));
|
||||
await user.click(await screen.findByRole("button", { name: /edit settings/i }));
|
||||
|
||||
expect(await screen.findByLabelText("Team Name")).toBeInTheDocument();
|
||||
expect(toast.error).not.toHaveBeenCalled();
|
||||
});
|
||||
});
|
||||
});
|
||||
|
||||
describe("TeamInfoView - which team member fields reach the update payload depends on the open sections", () => {
|
||||
|
|
|
|||
|
|
@ -48,6 +48,14 @@ import React, { useEffect, useMemo, useState } from "react";
|
|||
import { useFieldArray } from "react-hook-form";
|
||||
import { z } from "zod/v4";
|
||||
import GuardrailsSelect from "./GuardrailsSelect";
|
||||
import {
|
||||
type CallerEditAccess,
|
||||
parseTeamEditAccess,
|
||||
TEAM_ADMIN_EDITING_DISABLED_DESCRIPTION,
|
||||
TEAM_ADMIN_EDITING_DISABLED_TITLE,
|
||||
type TeamAdminSettingsChanges,
|
||||
} from "./teamAdminEditAccess";
|
||||
import TeamAdminSettingsForm from "./TeamAdminSettingsForm";
|
||||
import { copyToClipboard as utilCopyToClipboard } from "../../utils/dataUtils";
|
||||
import AccessGroupSelector from "../common_components/AccessGroupSelector";
|
||||
import BudgetDurationDropdown, { NEVER_RESETS_BUDGET_DURATION } from "../common_components/budget_duration_dropdown";
|
||||
|
|
@ -297,6 +305,7 @@ export interface TeamData {
|
|||
guardrails?: string[];
|
||||
policies?: string[];
|
||||
object_permission?: ObjectPermission | null;
|
||||
caller_edit_access?: CallerEditAccess;
|
||||
team_member_budget_table: {
|
||||
max_budget: number;
|
||||
budget_duration: string | null;
|
||||
|
|
@ -315,7 +324,6 @@ export interface TeamInfoProps {
|
|||
accessToken: string | null;
|
||||
is_team_admin: boolean;
|
||||
is_proxy_admin: boolean;
|
||||
is_org_admin?: boolean;
|
||||
userModels: string[];
|
||||
editTeam: boolean;
|
||||
premiumUser?: boolean;
|
||||
|
|
@ -531,7 +539,6 @@ const TeamInfoView: React.FC<TeamInfoProps> = ({
|
|||
accessToken,
|
||||
is_team_admin,
|
||||
is_proxy_admin,
|
||||
is_org_admin = false,
|
||||
userModels,
|
||||
editTeam,
|
||||
premiumUser = false,
|
||||
|
|
@ -575,7 +582,7 @@ const TeamInfoView: React.FC<TeamInfoProps> = ({
|
|||
const [teamModelMaxBudget, setTeamModelMaxBudget] = useState<ModelMaxBudget>({});
|
||||
const routerSettingsRef = React.useRef<RouterSettingsAccordionRef>(null);
|
||||
const [organization, setOrganization] = useState<Organization | null>(null);
|
||||
const { userRole, userId } = useAuthorized();
|
||||
const { userRole } = useAuthorized();
|
||||
const { data: allMcpServers = [], isError: mcpServersFailed, isLoading: mcpServersLoading } = useMCPServers();
|
||||
const { data: allMcpToolsets = [], isError: mcpToolsetsFailed, isLoading: mcpToolsetsLoading } = useMCPToolsets();
|
||||
const { data: allAccessGroups = [], isError: accessGroupsFailed, isLoading: accessGroupsLoading } = useAccessGroups();
|
||||
|
|
@ -585,14 +592,6 @@ const TeamInfoView: React.FC<TeamInfoProps> = ({
|
|||
const { data: teamMetadataSchemaFields = [], isLoading: isTeamMetadataSchemaLoading } = useTeamMetadataSchema();
|
||||
const queryClient = useQueryClient();
|
||||
|
||||
// Check if user is org admin for this team's organization
|
||||
const isOrgAdminForTeam = useMemo(() => {
|
||||
const teamOrgId = teamData?.team_info?.organization_id;
|
||||
if (!teamOrgId || !userId) return false;
|
||||
const org = userOrganizations.find((o) => o.organization_id === teamOrgId);
|
||||
return org?.members?.some((m: any) => m.user_id === userId && m.user_role === "org_admin") ?? false;
|
||||
}, [teamData, userOrganizations, userId]);
|
||||
|
||||
// Models currently selected in the team edit form, used to scope the per-model
|
||||
// rate limit dropdown to models this team actually has access to.
|
||||
const watchedModels = form.watch("models");
|
||||
|
|
@ -616,15 +615,8 @@ const TeamInfoView: React.FC<TeamInfoProps> = ({
|
|||
return unfurlWildcardModelsInList(selected, userModels);
|
||||
}, [watchedModels, teamData, userModels]);
|
||||
|
||||
const isTeamAdminFromTeamData = useMemo(
|
||||
() =>
|
||||
teamData?.team_info?.members_with_roles?.some(
|
||||
(member) => member.user_id != null && member.user_id === userId && member.role === "admin",
|
||||
) ?? false,
|
||||
[teamData, userId],
|
||||
);
|
||||
|
||||
const canEditTeam = is_team_admin || is_proxy_admin || is_org_admin || isOrgAdminForTeam || isTeamAdminFromTeamData;
|
||||
const teamEditAccess = useMemo(() => parseTeamEditAccess(teamData?.team_info?.caller_edit_access), [teamData]);
|
||||
const canEditTeam = is_team_admin || is_proxy_admin || teamEditAccess.kind !== "none";
|
||||
const visibleTabs = useMemo(() => getTeamInfoVisibleTabs(canEditTeam), [canEditTeam]);
|
||||
const defaultTabKey = useMemo(() => getTeamInfoDefaultTab(editTeam, canEditTeam), [editTeam, canEditTeam]);
|
||||
const { onTabChange, hasVisited } = useVisitedTabs(defaultTabKey);
|
||||
|
|
@ -644,6 +636,15 @@ const TeamInfoView: React.FC<TeamInfoProps> = ({
|
|||
setIsEditing(true);
|
||||
};
|
||||
|
||||
const openSettingsEditor = (modelAliases: Record<string, string>) => {
|
||||
if (teamEditAccess.kind === "team_admin_disabled") {
|
||||
toast.error(TEAM_ADMIN_EDITING_DISABLED_TITLE, { description: TEAM_ADMIN_EDITING_DISABLED_DESCRIPTION });
|
||||
return;
|
||||
}
|
||||
setTeamModelAliases(modelAliases);
|
||||
startEditing();
|
||||
};
|
||||
|
||||
const applyKillSwitchToGuardrails = (checked: boolean) => {
|
||||
const current = form.getValues("guardrails") ?? [];
|
||||
const nonGlobals = current.filter((name) => !globalGuardrailNames.has(name));
|
||||
|
|
@ -863,6 +864,27 @@ const TeamInfoView: React.FC<TeamInfoProps> = ({
|
|||
setMemberToDelete(null);
|
||||
};
|
||||
|
||||
const persistTeamUpdate = async (token: string, updateData: Record<string, unknown>) => {
|
||||
await teamUpdateCall(token, updateData);
|
||||
queryClient.invalidateQueries({ queryKey: organizationKeys.all });
|
||||
|
||||
toast.success("Team settings updated successfully");
|
||||
setIsEditing(false);
|
||||
fetchTeamInfo();
|
||||
};
|
||||
|
||||
const saveTeamAdminSettings = async (changes: TeamAdminSettingsChanges) => {
|
||||
if (!accessToken) return;
|
||||
setIsTeamSaving(true);
|
||||
try {
|
||||
await persistTeamUpdate(accessToken, { team_id: teamId, ...changes });
|
||||
} catch (error) {
|
||||
console.error("Error updating team:", error);
|
||||
} finally {
|
||||
setIsTeamSaving(false);
|
||||
}
|
||||
};
|
||||
|
||||
const handleTeamUpdate = async (values: any) => {
|
||||
try {
|
||||
if (!accessToken) return;
|
||||
|
|
@ -1113,12 +1135,7 @@ const TeamInfoView: React.FC<TeamInfoProps> = ({
|
|||
}
|
||||
}
|
||||
|
||||
await teamUpdateCall(accessToken, updateData);
|
||||
queryClient.invalidateQueries({ queryKey: organizationKeys.all });
|
||||
|
||||
toast.success("Team settings updated successfully");
|
||||
setIsEditing(false);
|
||||
fetchTeamInfo();
|
||||
await persistTeamUpdate(accessToken, updateData);
|
||||
} catch (error) {
|
||||
console.error("Error updating team:", error);
|
||||
} finally {
|
||||
|
|
@ -1136,6 +1153,17 @@ const TeamInfoView: React.FC<TeamInfoProps> = ({
|
|||
|
||||
const { team_info: info } = teamData;
|
||||
|
||||
const teamAdminSettingsEditor =
|
||||
teamEditAccess.kind === "team_admin" ? (
|
||||
<TeamAdminSettingsForm
|
||||
initialValues={{ tpm_limit: info.tpm_limit }}
|
||||
editableFields={teamEditAccess.editableFields}
|
||||
isSaving={isTeamSaving}
|
||||
onCancel={() => setIsEditing(false)}
|
||||
onSave={saveTeamAdminSettings}
|
||||
/>
|
||||
) : null;
|
||||
|
||||
const inheritedMcpServers = computeInheritedGrants(
|
||||
info.access_group_mcp_server_ids,
|
||||
info.access_group_details,
|
||||
|
|
@ -1340,10 +1368,7 @@ const TeamInfoView: React.FC<TeamInfoProps> = ({
|
|||
{canEditTeam && !isEditing && (
|
||||
<Button
|
||||
variant="outline"
|
||||
onClick={() => {
|
||||
setTeamModelAliases(info.litellm_model_table?.model_aliases ?? {});
|
||||
startEditing();
|
||||
}}
|
||||
onClick={() => openSettingsEditor(info.litellm_model_table?.model_aliases ?? {})}
|
||||
>
|
||||
<Pencil />
|
||||
Edit Settings
|
||||
|
|
@ -1351,8 +1376,8 @@ const TeamInfoView: React.FC<TeamInfoProps> = ({
|
|||
)}
|
||||
</div>
|
||||
|
||||
{isEditing && isGuardrailsLoading ? (
|
||||
<div className="p-4">Loading...</div>
|
||||
{isEditing && (teamAdminSettingsEditor !== null || isGuardrailsLoading) ? (
|
||||
teamAdminSettingsEditor ?? <div className="p-4">Loading...</div>
|
||||
) : isEditing ? (
|
||||
<TooltipProvider>
|
||||
<form onSubmit={(event) => void form.handleSubmit(onTeamUpdateSubmit)(event)}>
|
||||
|
|
|
|||
|
|
@ -0,0 +1,117 @@
|
|||
import { describe, expect, it } from "vitest";
|
||||
|
||||
import {
|
||||
parseSupportedTeamAdminEditableFields,
|
||||
parseTeamAdminEditableFields,
|
||||
parseTeamEditAccess,
|
||||
teamAdminFieldLabel,
|
||||
teamAdminSettingsChanges,
|
||||
} from "./teamAdminEditAccess";
|
||||
|
||||
describe("teamAdminFieldLabel", () => {
|
||||
it("names tpm_limit the way the team settings form does", () => {
|
||||
expect(teamAdminFieldLabel("tpm_limit")).toBe("Tokens per minute Limit (TPM)");
|
||||
});
|
||||
|
||||
it("falls back to the raw field name for a field the dashboard has no label for", () => {
|
||||
expect(teamAdminFieldLabel("max_budget")).toBe("max_budget");
|
||||
});
|
||||
});
|
||||
|
||||
describe("teamAdminSettingsChanges", () => {
|
||||
const tpmEnabled = new Set(["tpm_limit"]);
|
||||
const stored = { tpm_limit: 1000 };
|
||||
|
||||
it.each([
|
||||
["a typed number string", "5000", 5000],
|
||||
["a number", 1200, 1200],
|
||||
["zero", "0", 0],
|
||||
["an emptied input", "", null],
|
||||
["whitespace", " ", null],
|
||||
["no limit", null, null],
|
||||
["an unset value", undefined, null],
|
||||
])("sends tpm_limit changed to %s", (_label, tpm_limit, expected) => {
|
||||
expect(teamAdminSettingsChanges({ tpm_limit }, stored, tpmEnabled)).toStrictEqual({ tpm_limit: expected });
|
||||
});
|
||||
|
||||
it.each([
|
||||
["the stored number", 1000, { tpm_limit: 1000 }],
|
||||
["the stored number typed back in", "1000", { tpm_limit: 1000 }],
|
||||
["an emptied input over no stored limit", "", { tpm_limit: null }],
|
||||
["an unset value over no stored limit", undefined, { tpm_limit: null }],
|
||||
])("sends nothing for %s", (_label, tpm_limit, initialValues) => {
|
||||
expect(teamAdminSettingsChanges({ tpm_limit }, initialValues, tpmEnabled)).toStrictEqual({});
|
||||
});
|
||||
|
||||
it("leaves tpm_limit out when the proxy did not enable it for team admins", () => {
|
||||
expect(teamAdminSettingsChanges({ tpm_limit: "5000" }, stored, new Set(["max_budget"]))).toStrictEqual({});
|
||||
});
|
||||
});
|
||||
|
||||
describe("parseTeamAdminEditableFields", () => {
|
||||
it("returns the configured list", () => {
|
||||
expect(parseTeamAdminEditableFields({ team_admin_editable_team_fields: ["tpm_limit", "rpm_limit"] })).toEqual([
|
||||
"tpm_limit",
|
||||
"rpm_limit",
|
||||
]);
|
||||
});
|
||||
|
||||
it.each([
|
||||
["no values yet", undefined],
|
||||
["setting missing", {}],
|
||||
["setting is null", { team_admin_editable_team_fields: null }],
|
||||
["setting is a string", { team_admin_editable_team_fields: "tpm_limit" }],
|
||||
["list holds a non-string", { team_admin_editable_team_fields: ["tpm_limit", 7] }],
|
||||
])("fails closed to an empty list when %s", (_label, values) => {
|
||||
expect(parseTeamAdminEditableFields(values)).toEqual([]);
|
||||
});
|
||||
});
|
||||
|
||||
describe("parseSupportedTeamAdminEditableFields", () => {
|
||||
it("reads the enum the proxy advertises on the setting's items schema", () => {
|
||||
const schema = {
|
||||
properties: {
|
||||
team_admin_editable_team_fields: {
|
||||
type: "array",
|
||||
items: { type: "string", enum: ["max_budget", "tpm_limit"] },
|
||||
},
|
||||
},
|
||||
};
|
||||
expect(parseSupportedTeamAdminEditableFields(schema)).toEqual(["max_budget", "tpm_limit"]);
|
||||
});
|
||||
|
||||
it.each([
|
||||
["schema not loaded", undefined],
|
||||
["property absent", { properties: {} }],
|
||||
["items has no enum", { properties: { team_admin_editable_team_fields: { items: { type: "string" } } } }],
|
||||
["enum is not a string list", { properties: { team_admin_editable_team_fields: { items: { enum: [1] } } } }],
|
||||
])("returns no supported fields when %s", (_label, schema) => {
|
||||
expect(parseSupportedTeamAdminEditableFields(schema)).toEqual([]);
|
||||
});
|
||||
});
|
||||
|
||||
describe("parseTeamEditAccess", () => {
|
||||
it.each([
|
||||
["unrestricted", { kind: "unrestricted" }],
|
||||
["team_admin_disabled", { kind: "team_admin_disabled" }],
|
||||
["none", { kind: "none" }],
|
||||
])("passes the proxy's %s verdict through", (_kind, verdict) => {
|
||||
expect(parseTeamEditAccess(verdict)).toEqual(verdict);
|
||||
});
|
||||
|
||||
it("hands a team admin the fields the proxy enabled", () => {
|
||||
expect(parseTeamEditAccess({ kind: "team_admin", editable_fields: ["tpm_limit"] })).toEqual({
|
||||
kind: "team_admin",
|
||||
editableFields: new Set(["tpm_limit"]),
|
||||
});
|
||||
});
|
||||
|
||||
it.each([
|
||||
["the proxy sent nothing", undefined],
|
||||
["the kind is unknown", { kind: "owner" }],
|
||||
["a team admin verdict lacks its field list", { kind: "team_admin" }],
|
||||
["the field list holds a non-string", { kind: "team_admin", editable_fields: [7] }],
|
||||
])("fails closed to no access when %s", (_label, value) => {
|
||||
expect(parseTeamEditAccess(value)).toEqual({ kind: "none" });
|
||||
});
|
||||
});
|
||||
|
|
@ -0,0 +1,78 @@
|
|||
import { z } from "zod/v4";
|
||||
|
||||
export const TEAM_ADMIN_EDITABLE_TEAM_FIELDS_SETTING = "team_admin_editable_team_fields";
|
||||
|
||||
export const TEAM_ADMIN_EDITING_DISABLED_TITLE = "Team admins cannot edit team settings on this proxy";
|
||||
export const TEAM_ADMIN_EDITING_DISABLED_DESCRIPTION =
|
||||
"Ask a proxy admin to enable fields under Settings > UI > Team admin editable fields.";
|
||||
|
||||
const callerEditAccessSchema = z.discriminatedUnion("kind", [
|
||||
z.object({ kind: z.literal("unrestricted") }),
|
||||
z.object({ kind: z.literal("team_admin"), editable_fields: z.array(z.string()) }),
|
||||
z.object({ kind: z.literal("team_admin_disabled") }),
|
||||
z.object({ kind: z.literal("none") }),
|
||||
]);
|
||||
|
||||
export type CallerEditAccess = z.infer<typeof callerEditAccessSchema>;
|
||||
|
||||
export type TeamEditAccess =
|
||||
| { readonly kind: "unrestricted" }
|
||||
| { readonly kind: "team_admin"; readonly editableFields: ReadonlySet<string> }
|
||||
| { readonly kind: "team_admin_disabled" }
|
||||
| { readonly kind: "none" };
|
||||
|
||||
const fieldListSchema = z.array(z.string()).catch([]);
|
||||
|
||||
export const parseTeamAdminEditableFields = (uiSettingsValues: unknown): readonly string[] => {
|
||||
const values = z.record(z.string(), z.unknown()).catch({}).parse(uiSettingsValues);
|
||||
return fieldListSchema.parse(values[TEAM_ADMIN_EDITABLE_TEAM_FIELDS_SETTING]);
|
||||
};
|
||||
|
||||
export const parseSupportedTeamAdminEditableFields = (uiSettingsFieldSchema: unknown): readonly string[] => {
|
||||
const property = z
|
||||
.object({ properties: z.object({ [TEAM_ADMIN_EDITABLE_TEAM_FIELDS_SETTING]: z.object({ items: z.unknown() }) }) })
|
||||
.safeParse(uiSettingsFieldSchema);
|
||||
if (!property.success) return [];
|
||||
const items = z
|
||||
.object({ enum: z.unknown() })
|
||||
.safeParse(property.data.properties[TEAM_ADMIN_EDITABLE_TEAM_FIELDS_SETTING].items);
|
||||
return items.success ? fieldListSchema.parse(items.data.enum) : [];
|
||||
};
|
||||
|
||||
const TEAM_ADMIN_FIELD_LABELS: ReadonlyMap<string, string> = new Map([["tpm_limit", "Tokens per minute Limit (TPM)"]]);
|
||||
|
||||
export const teamAdminFieldLabel = (field: string): string => TEAM_ADMIN_FIELD_LABELS.get(field) ?? field;
|
||||
|
||||
export interface TeamAdminSettingsValues {
|
||||
readonly tpm_limit?: string | number | null;
|
||||
}
|
||||
|
||||
export interface TeamAdminSettingsChanges {
|
||||
readonly tpm_limit?: number | null;
|
||||
}
|
||||
|
||||
const numberOrNull = (value: string | number | null | undefined): number | null => {
|
||||
if (value === null || value === undefined || String(value).trim() === "") return null;
|
||||
const parsed = Number(value);
|
||||
return Number.isNaN(parsed) ? null : parsed;
|
||||
};
|
||||
|
||||
export const teamAdminSettingsChanges = (
|
||||
values: TeamAdminSettingsValues,
|
||||
initialValues: TeamAdminSettingsValues,
|
||||
editableFields: ReadonlySet<string>,
|
||||
): TeamAdminSettingsChanges => {
|
||||
const tpmLimit = numberOrNull(values.tpm_limit);
|
||||
return editableFields.has("tpm_limit") && tpmLimit !== numberOrNull(initialValues.tpm_limit)
|
||||
? { tpm_limit: tpmLimit }
|
||||
: {};
|
||||
};
|
||||
|
||||
export const parseTeamEditAccess = (callerEditAccess: unknown): TeamEditAccess => {
|
||||
const parsed = callerEditAccessSchema.safeParse(callerEditAccess);
|
||||
if (!parsed.success) return { kind: "none" };
|
||||
if (parsed.data.kind === "team_admin") {
|
||||
return { kind: "team_admin", editableFields: new Set(parsed.data.editable_fields) };
|
||||
}
|
||||
return parsed.data;
|
||||
};
|
||||
Loading…
Add table
Reference in a new issue