mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-08 22:21:35 +00:00
Merge branch 'main' into litellm_mcp_oauth2_m2m_ui
This commit is contained in:
Ishaan Jaffer
commit
89167b0cc7
11 changed files with 564 additions and 143 deletions
|
|
@ -189,4 +189,4 @@ curl http://localhost:4000/mcp-rest/tools/call \
|
|||
| `client_id` | Yes | OAuth2 client ID. Supports `os.environ/VAR_NAME` |
|
||||
| `client_secret` | Yes | OAuth2 client secret. Supports `os.environ/VAR_NAME` |
|
||||
| `token_url` | Yes | Token endpoint URL |
|
||||
| `scopes` | No | List of scopes to request |
|
||||
| `scopes` | No | List of scopes to request |
|
||||
|
|
|
|||
|
|
@ -29,7 +29,10 @@ from litellm.proxy._types import (
|
|||
UserAPIKeyAuth,
|
||||
)
|
||||
from litellm.types.integrations.prometheus import *
|
||||
from litellm.types.integrations.prometheus import _sanitize_prometheus_label_name
|
||||
from litellm.types.integrations.prometheus import (
|
||||
_sanitize_prometheus_label_name,
|
||||
_sanitize_prometheus_label_value,
|
||||
)
|
||||
from litellm.types.utils import StandardLoggingPayload
|
||||
|
||||
if TYPE_CHECKING:
|
||||
|
|
@ -1276,11 +1279,17 @@ class PrometheusLogger(CustomLogger):
|
|||
)
|
||||
|
||||
self.litellm_remaining_api_key_requests_for_model.labels(
|
||||
user_api_key, user_api_key_alias, model_group, model_id
|
||||
_sanitize_prometheus_label_value(user_api_key),
|
||||
_sanitize_prometheus_label_value(user_api_key_alias),
|
||||
_sanitize_prometheus_label_value(model_group),
|
||||
_sanitize_prometheus_label_value(model_id),
|
||||
).set(remaining_requests)
|
||||
|
||||
self.litellm_remaining_api_key_tokens_for_model.labels(
|
||||
user_api_key, user_api_key_alias, model_group, model_id
|
||||
_sanitize_prometheus_label_value(user_api_key),
|
||||
_sanitize_prometheus_label_value(user_api_key_alias),
|
||||
_sanitize_prometheus_label_value(model_group),
|
||||
_sanitize_prometheus_label_value(model_id),
|
||||
).set(remaining_tokens)
|
||||
|
||||
def _set_latency_metrics(
|
||||
|
|
@ -1401,14 +1410,14 @@ class PrometheusLogger(CustomLogger):
|
|||
|
||||
try:
|
||||
self.litellm_llm_api_failed_requests_metric.labels(
|
||||
end_user_id,
|
||||
user_api_key,
|
||||
user_api_key_alias,
|
||||
model,
|
||||
user_api_team,
|
||||
user_api_team_alias,
|
||||
user_id,
|
||||
standard_logging_payload.get("model_id", ""),
|
||||
_sanitize_prometheus_label_value(end_user_id),
|
||||
_sanitize_prometheus_label_value(user_api_key),
|
||||
_sanitize_prometheus_label_value(user_api_key_alias),
|
||||
_sanitize_prometheus_label_value(model),
|
||||
_sanitize_prometheus_label_value(user_api_team),
|
||||
_sanitize_prometheus_label_value(user_api_team_alias),
|
||||
_sanitize_prometheus_label_value(user_id),
|
||||
_sanitize_prometheus_label_value(standard_logging_payload.get("model_id", "")),
|
||||
).inc()
|
||||
self.set_llm_deployment_failure_metrics(kwargs)
|
||||
except Exception as e:
|
||||
|
|
@ -2354,7 +2363,11 @@ class PrometheusLogger(CustomLogger):
|
|||
increment metric when litellm.Router / load balancing logic places a deployment in cool down
|
||||
"""
|
||||
self.litellm_deployment_cooled_down.labels(
|
||||
litellm_model_name, model_id, api_base, api_provider, exception_status
|
||||
_sanitize_prometheus_label_value(litellm_model_name),
|
||||
_sanitize_prometheus_label_value(model_id),
|
||||
_sanitize_prometheus_label_value(api_base),
|
||||
_sanitize_prometheus_label_value(api_provider),
|
||||
_sanitize_prometheus_label_value(exception_status),
|
||||
).inc()
|
||||
|
||||
def increment_callback_logging_failure(
|
||||
|
|
@ -3074,9 +3087,10 @@ def prometheus_label_factory(
|
|||
# Extract dictionary from Pydantic object
|
||||
enum_dict = enum_values.model_dump()
|
||||
|
||||
# Filter supported labels
|
||||
# Filter supported labels and sanitize values to prevent breaking
|
||||
# the Prometheus text format (e.g. U+2028 Line Separator in label values)
|
||||
filtered_labels = {
|
||||
label: value
|
||||
label: _sanitize_prometheus_label_value(value)
|
||||
for label, value in enum_dict.items()
|
||||
if label in supported_enum_labels
|
||||
}
|
||||
|
|
@ -3094,14 +3108,14 @@ def prometheus_label_factory(
|
|||
# check sanitized key
|
||||
sanitized_key = _sanitize_prometheus_label_name(key)
|
||||
if sanitized_key in supported_enum_labels:
|
||||
filtered_labels[sanitized_key] = value
|
||||
filtered_labels[sanitized_key] = _sanitize_prometheus_label_value(value)
|
||||
|
||||
# Add custom tags if configured
|
||||
if enum_values.tags is not None:
|
||||
custom_tag_labels = get_custom_labels_from_tags(enum_values.tags)
|
||||
for key, value in custom_tag_labels.items():
|
||||
if key in supported_enum_labels:
|
||||
filtered_labels[key] = value
|
||||
filtered_labels[key] = _sanitize_prometheus_label_value(value)
|
||||
|
||||
for label in supported_enum_labels:
|
||||
if label not in filtered_labels:
|
||||
|
|
|
|||
|
|
@ -8,6 +8,8 @@ with ``client_id``, ``client_secret``, and ``token_url``.
|
|||
import asyncio
|
||||
from typing import TYPE_CHECKING, Dict, Optional, Tuple, Union
|
||||
|
||||
import httpx
|
||||
|
||||
from litellm._logging import verbose_logger
|
||||
from litellm.caching.in_memory_cache import InMemoryCache
|
||||
from litellm.constants import (
|
||||
|
|
@ -39,9 +41,7 @@ class MCPOAuth2TokenCache(InMemoryCache):
|
|||
self._locks: Dict[str, asyncio.Lock] = {}
|
||||
|
||||
def _get_lock(self, server_id: str) -> asyncio.Lock:
|
||||
if server_id not in self._locks:
|
||||
self._locks[server_id] = asyncio.Lock()
|
||||
return self._locks[server_id]
|
||||
return self._locks.setdefault(server_id, asyncio.Lock())
|
||||
|
||||
async def async_get_token(self, server: "MCPServer") -> Optional[str]:
|
||||
"""Return a valid access token, fetching or refreshing as needed.
|
||||
|
|
@ -76,8 +76,13 @@ class MCPOAuth2TokenCache(InMemoryCache):
|
|||
"""
|
||||
client = get_async_httpx_client(llm_provider=httpxSpecialProvider.MCP)
|
||||
|
||||
assert server.client_id is not None, "client_id must be set"
|
||||
assert server.client_secret is not None, "client_secret must be set"
|
||||
if not server.client_id or not server.client_secret or not server.token_url:
|
||||
raise ValueError(
|
||||
f"MCP server '{server.server_id}' missing required OAuth2 fields: "
|
||||
f"client_id={bool(server.client_id)}, "
|
||||
f"client_secret={bool(server.client_secret)}, "
|
||||
f"token_url={bool(server.token_url)}"
|
||||
)
|
||||
|
||||
data: Dict[str, str] = {
|
||||
"grant_type": "client_credentials",
|
||||
|
|
@ -88,23 +93,41 @@ class MCPOAuth2TokenCache(InMemoryCache):
|
|||
data["scope"] = " ".join(server.scopes)
|
||||
|
||||
verbose_logger.debug(
|
||||
"Fetching OAuth2 client_credentials token for MCP server %s from %s",
|
||||
"Fetching OAuth2 client_credentials token for MCP server %s",
|
||||
server.server_id,
|
||||
server.token_url,
|
||||
)
|
||||
|
||||
response = await client.post(server.token_url, data=data) # type: ignore[arg-type]
|
||||
response.raise_for_status()
|
||||
try:
|
||||
response = await client.post(server.token_url, data=data)
|
||||
response.raise_for_status()
|
||||
except httpx.HTTPStatusError as exc:
|
||||
raise ValueError(
|
||||
f"OAuth2 token request for MCP server '{server.server_id}' "
|
||||
f"failed with status {exc.response.status_code}"
|
||||
) from exc
|
||||
|
||||
body = response.json()
|
||||
|
||||
if not isinstance(body, dict):
|
||||
raise ValueError(
|
||||
f"OAuth2 token response for MCP server '{server.server_id}' "
|
||||
f"returned non-object JSON (got {type(body).__name__})"
|
||||
)
|
||||
|
||||
access_token = body.get("access_token")
|
||||
if not access_token:
|
||||
raise ValueError(
|
||||
f"OAuth2 token response for MCP server '{server.server_id}' "
|
||||
f"missing 'access_token': {body}"
|
||||
f"missing 'access_token'"
|
||||
)
|
||||
|
||||
expires_in = int(body.get("expires_in", 3600))
|
||||
# Safely parse expires_in — providers may return null or non-numeric values
|
||||
raw_expires_in = body.get("expires_in")
|
||||
try:
|
||||
expires_in = int(raw_expires_in) if raw_expires_in is not None else MCP_OAUTH2_TOKEN_CACHE_DEFAULT_TTL
|
||||
except (TypeError, ValueError):
|
||||
expires_in = MCP_OAUTH2_TOKEN_CACHE_DEFAULT_TTL
|
||||
|
||||
ttl = max(expires_in - MCP_OAUTH2_TOKEN_EXPIRY_BUFFER_SECONDS, MCP_OAUTH2_TOKEN_CACHE_MIN_TTL)
|
||||
|
||||
verbose_logger.info(
|
||||
|
|
|
|||
|
|
@ -1,7 +1,7 @@
|
|||
import re
|
||||
from dataclasses import dataclass
|
||||
from enum import Enum
|
||||
from typing import Dict, List, Literal, Optional, Tuple
|
||||
from typing import Any, Dict, List, Literal, Optional, Tuple
|
||||
|
||||
from pydantic import BaseModel, Field
|
||||
from typing_extensions import Annotated
|
||||
|
|
@ -41,6 +41,38 @@ def _sanitize_prometheus_label_name(label: str) -> str:
|
|||
return sanitized
|
||||
|
||||
|
||||
def _sanitize_prometheus_label_value(value: Optional[Any]) -> Optional[str]:
|
||||
"""
|
||||
Sanitize a label value for Prometheus text format compatibility.
|
||||
|
||||
Removes or replaces characters that break the Prometheus exposition format:
|
||||
- U+2028 (Line Separator) and U+2029 (Paragraph Separator) are removed
|
||||
- Carriage returns are removed
|
||||
- Newlines are replaced with spaces
|
||||
- Backslashes and double quotes are escaped per Prometheus spec
|
||||
"""
|
||||
if value is None:
|
||||
return None
|
||||
|
||||
# Coerce non-string values (int, bool, etc.) to str before sanitizing
|
||||
if not isinstance(value, str):
|
||||
value = str(value)
|
||||
|
||||
# Remove Unicode line/paragraph separators that break text format
|
||||
value = value.replace("\u2028", "").replace("\u2029", "")
|
||||
|
||||
# Remove carriage returns
|
||||
value = value.replace("\r", "")
|
||||
|
||||
# Replace newlines with spaces
|
||||
value = value.replace("\n", " ")
|
||||
|
||||
# Escape backslashes and double quotes per Prometheus exposition format
|
||||
value = value.replace("\\", "\\\\").replace('"', '\\"')
|
||||
|
||||
return value
|
||||
|
||||
|
||||
@dataclass
|
||||
class MetricValidationError:
|
||||
"""Error for invalid metric name"""
|
||||
|
|
|
|||
|
|
@ -292,6 +292,87 @@ def test_prometheus_metrics_use_normalized_routes():
|
|||
print("✅ Prometheus metrics use normalized routes in labels")
|
||||
|
||||
|
||||
def test_prometheus_label_value_sanitization():
|
||||
"""
|
||||
Test that Prometheus label values are sanitized to prevent breaking
|
||||
the Prometheus text format.
|
||||
|
||||
Issue: Unicode Line Separator (U+2028) in label values (e.g. from a
|
||||
malformed model name) breaks the Prometheus exposition format, causing
|
||||
scrapers like Datadog to fail parsing the entire /metrics endpoint.
|
||||
"""
|
||||
from litellm.integrations.prometheus import (
|
||||
PrometheusLogger,
|
||||
UserAPIKeyLabelValues,
|
||||
prometheus_label_factory,
|
||||
)
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
prometheus_logger = MagicMock()
|
||||
prometheus_logger.get_labels_for_metric = (
|
||||
PrometheusLogger.get_labels_for_metric.__get__(prometheus_logger)
|
||||
)
|
||||
|
||||
# Simulate a model name with U+2028 (Unicode Line Separator) appended
|
||||
# and an api_key_alias with newlines and quotes
|
||||
enum_values = UserAPIKeyLabelValues(
|
||||
requested_model="claude-haiku-4-5-20251001\u2028",
|
||||
api_key_alias='My Key "test"\nwith newline',
|
||||
route="/v1/chat/completions",
|
||||
status_code="400",
|
||||
)
|
||||
|
||||
labels = prometheus_label_factory(
|
||||
supported_enum_labels=prometheus_logger.get_labels_for_metric(
|
||||
metric_name="litellm_proxy_total_requests_metric"
|
||||
),
|
||||
enum_values=enum_values,
|
||||
)
|
||||
|
||||
# U+2028 must be stripped
|
||||
assert "\u2028" not in labels["requested_model"], (
|
||||
f"U+2028 should be removed from label value, got: {repr(labels['requested_model'])}"
|
||||
)
|
||||
assert labels["requested_model"] == "claude-haiku-4-5-20251001"
|
||||
|
||||
# Newlines must be replaced with spaces, quotes must be escaped
|
||||
assert "\n" not in labels["api_key_alias"]
|
||||
assert labels["api_key_alias"] == 'My Key \\"test\\" with newline'
|
||||
|
||||
print("✅ Prometheus label values are properly sanitized")
|
||||
|
||||
|
||||
def test_prometheus_label_value_sanitization_unicode_paragraph_separator():
|
||||
"""Test that U+2029 (Paragraph Separator) is also stripped."""
|
||||
from litellm.types.integrations.prometheus import _sanitize_prometheus_label_value
|
||||
|
||||
result = _sanitize_prometheus_label_value("model\u2029name")
|
||||
assert result == "modelname"
|
||||
assert "\u2029" not in result
|
||||
|
||||
print("✅ U+2029 Paragraph Separator is stripped")
|
||||
|
||||
|
||||
def test_prometheus_label_value_sanitization_none():
|
||||
"""Test that None values pass through unchanged."""
|
||||
from litellm.types.integrations.prometheus import _sanitize_prometheus_label_value
|
||||
|
||||
assert _sanitize_prometheus_label_value(None) is None
|
||||
|
||||
print("✅ None values pass through unchanged")
|
||||
|
||||
|
||||
def test_prometheus_label_value_sanitization_non_string_types():
|
||||
"""Test that non-string values (int, bool, etc.) are coerced to str."""
|
||||
from litellm.types.integrations.prometheus import _sanitize_prometheus_label_value
|
||||
|
||||
assert _sanitize_prometheus_label_value(200) == "200"
|
||||
assert _sanitize_prometheus_label_value(True) == "True"
|
||||
assert _sanitize_prometheus_label_value(3.14) == "3.14"
|
||||
|
||||
print("✅ Non-string values are coerced to str")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
test_user_email_in_required_metrics()
|
||||
test_user_email_label_exists()
|
||||
|
|
@ -301,4 +382,8 @@ if __name__ == "__main__":
|
|||
test_route_normalization_preserves_static_routes()
|
||||
test_route_normalization_other_dynamic_apis()
|
||||
test_prometheus_metrics_use_normalized_routes()
|
||||
test_prometheus_label_value_sanitization()
|
||||
test_prometheus_label_value_sanitization_unicode_paragraph_separator()
|
||||
test_prometheus_label_value_sanitization_none()
|
||||
test_prometheus_label_value_sanitization_non_string_types()
|
||||
print("\n✅ All prometheus label tests passed!")
|
||||
|
|
|
|||
|
|
@ -5,10 +5,9 @@ Covers the critical path: resolve_mcp_auth(), token caching, auth priority,
|
|||
fallback to static token, and the skip-condition property.
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import time
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
|
||||
from litellm.proxy._experimental.mcp_server.oauth2_token_cache import (
|
||||
|
|
@ -120,3 +119,39 @@ def test_needs_user_oauth_token_property():
|
|||
|
||||
# Non-OAuth2 → never needs user OAuth token
|
||||
assert _server(auth_type=MCPAuth.bearer_token).needs_user_oauth_token is False
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_http_error_raises_value_error():
|
||||
"""HTTP errors from the token endpoint are wrapped in a clear ValueError."""
|
||||
server = _server()
|
||||
mock_response = MagicMock()
|
||||
mock_response.status_code = 401
|
||||
mock_response.raise_for_status.side_effect = httpx.HTTPStatusError(
|
||||
"Unauthorized", request=MagicMock(), response=mock_response,
|
||||
)
|
||||
mock_client = AsyncMock()
|
||||
mock_client.post.return_value = mock_response
|
||||
|
||||
with patch(
|
||||
"litellm.proxy._experimental.mcp_server.oauth2_token_cache.get_async_httpx_client",
|
||||
return_value=mock_client,
|
||||
), pytest.raises(ValueError, match="failed with status 401"):
|
||||
await resolve_mcp_auth(server)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_non_dict_response_raises_value_error():
|
||||
"""A non-dict JSON response raises a clear ValueError."""
|
||||
server = _server()
|
||||
resp = MagicMock()
|
||||
resp.json.return_value = ["not", "a", "dict"]
|
||||
resp.raise_for_status = MagicMock()
|
||||
mock_client = AsyncMock()
|
||||
mock_client.post.return_value = resp
|
||||
|
||||
with patch(
|
||||
"litellm.proxy._experimental.mcp_server.oauth2_token_cache.get_async_httpx_client",
|
||||
return_value=mock_client,
|
||||
), pytest.raises(ValueError, match="non-object JSON"):
|
||||
await resolve_mcp_auth(server)
|
||||
|
|
|
|||
264
ui/litellm-dashboard/src/components/CreateUserButton.test.tsx
Normal file
264
ui/litellm-dashboard/src/components/CreateUserButton.test.tsx
Normal file
|
|
@ -0,0 +1,264 @@
|
|||
import { render, screen, waitFor, within } from "@testing-library/react";
|
||||
import userEvent from "@testing-library/user-event";
|
||||
import { QueryClient, QueryClientProvider } from "@tanstack/react-query";
|
||||
import { describe, it, expect, vi, beforeEach } from "vitest";
|
||||
import { CreateUserButton } from "./CreateUserButton";
|
||||
import * as networking from "./networking";
|
||||
import NotificationsManager from "./molecules/notifications_manager";
|
||||
|
||||
vi.mock("./networking", () => ({
|
||||
userCreateCall: vi.fn(),
|
||||
modelAvailableCall: vi.fn().mockResolvedValue({ data: [] }),
|
||||
invitationCreateCall: vi.fn(),
|
||||
getProxyUISettings: vi.fn().mockResolvedValue({
|
||||
PROXY_BASE_URL: null,
|
||||
PROXY_LOGOUT_URL: null,
|
||||
DEFAULT_TEAM_DISABLED: false,
|
||||
SSO_ENABLED: false,
|
||||
}),
|
||||
getProxyBaseUrl: vi.fn().mockReturnValue("http://localhost"),
|
||||
}));
|
||||
|
||||
vi.mock("./bulk_create_users_button", () => ({
|
||||
default: () => <div data-testid="bulk-create-users">Bulk Create Users</div>,
|
||||
}));
|
||||
|
||||
const mockUserCreateCall = vi.mocked(networking.userCreateCall);
|
||||
const mockInvitationCreateCall = vi.mocked(networking.invitationCreateCall);
|
||||
const mockGetProxyUISettings = vi.mocked(networking.getProxyUISettings);
|
||||
const mockNotificationsManager = vi.mocked(NotificationsManager);
|
||||
|
||||
const createQueryClient = () =>
|
||||
new QueryClient({
|
||||
defaultOptions: { queries: { retry: false, gcTime: 0 } },
|
||||
});
|
||||
|
||||
const defaultProps = {
|
||||
userID: "123",
|
||||
accessToken: "token",
|
||||
teams: [],
|
||||
possibleUIRoles: null as Record<string, Record<string, string>> | null,
|
||||
};
|
||||
|
||||
function renderWithProviders(ui: React.ReactElement) {
|
||||
const qc = createQueryClient();
|
||||
return render(<QueryClientProvider client={qc}>{ui}</QueryClientProvider>);
|
||||
}
|
||||
|
||||
describe("CreateUserButton", { timeout: 20000 }, () => {
|
||||
beforeEach(() => {
|
||||
vi.clearAllMocks();
|
||||
mockGetProxyUISettings.mockResolvedValue({
|
||||
PROXY_BASE_URL: null,
|
||||
PROXY_LOGOUT_URL: null,
|
||||
DEFAULT_TEAM_DISABLED: false,
|
||||
SSO_ENABLED: false,
|
||||
});
|
||||
});
|
||||
|
||||
it("should render the create user form when embedded", () => {
|
||||
renderWithProviders(
|
||||
<CreateUserButton {...defaultProps} isEmbedded />,
|
||||
);
|
||||
expect(screen.getByRole("button", { name: /create user/i })).toBeInTheDocument();
|
||||
});
|
||||
|
||||
it("should render the invite user button when not embedded", async () => {
|
||||
renderWithProviders(<CreateUserButton {...defaultProps} />);
|
||||
await waitFor(() => {
|
||||
expect(screen.getByRole("button", { name: /\+ invite user/i })).toBeInTheDocument();
|
||||
});
|
||||
});
|
||||
|
||||
it("should open the invite modal when invite user button is clicked", async () => {
|
||||
const user = userEvent.setup();
|
||||
renderWithProviders(<CreateUserButton {...defaultProps} />);
|
||||
await waitFor(() => {
|
||||
expect(screen.getByRole("button", { name: /\+ invite user/i })).toBeInTheDocument();
|
||||
});
|
||||
await user.click(screen.getByRole("button", { name: /\+ invite user/i }));
|
||||
const dialog = screen.getByRole("dialog", { name: /invite user/i });
|
||||
expect(dialog).toBeInTheDocument();
|
||||
expect(within(dialog).getByRole("button", { name: /invite user/i })).toBeInTheDocument();
|
||||
});
|
||||
|
||||
it("should display email invitations info message in embedded mode", () => {
|
||||
renderWithProviders(<CreateUserButton {...defaultProps} isEmbedded />);
|
||||
expect(screen.getByText("Email invitations")).toBeInTheDocument();
|
||||
});
|
||||
|
||||
it("should display user role options when possibleUIRoles is provided", async () => {
|
||||
const possibleUIRoles = {
|
||||
proxy_admin: { ui_label: "Admin", description: "Full access" },
|
||||
proxy_user: { ui_label: "User", description: "Limited access" },
|
||||
};
|
||||
renderWithProviders(
|
||||
<CreateUserButton {...defaultProps} possibleUIRoles={possibleUIRoles} isEmbedded />,
|
||||
);
|
||||
await userEvent.click(screen.getByRole("combobox", { name: /user role/i }));
|
||||
expect(screen.getByText("Admin")).toBeInTheDocument();
|
||||
expect(screen.getByText("User")).toBeInTheDocument();
|
||||
});
|
||||
|
||||
it("should call userCreateCall when form is submitted in embedded mode", async () => {
|
||||
const user = userEvent.setup();
|
||||
mockUserCreateCall.mockResolvedValue({ data: { user_id: "new-user-123" } });
|
||||
mockInvitationCreateCall.mockResolvedValue({
|
||||
id: "inv-1",
|
||||
user_id: "new-user-123",
|
||||
has_user_setup_sso: false,
|
||||
} as any);
|
||||
|
||||
renderWithProviders(
|
||||
<CreateUserButton {...defaultProps} possibleUIRoles={{ proxy_user: { ui_label: "User", description: "" } }} isEmbedded />,
|
||||
);
|
||||
|
||||
await user.type(screen.getByLabelText(/user email/i), "test@example.com");
|
||||
await user.click(screen.getByRole("combobox", { name: /user role/i }));
|
||||
await user.click(screen.getByText("User"));
|
||||
await user.click(screen.getByRole("button", { name: /create user/i }));
|
||||
|
||||
await waitFor(() => {
|
||||
expect(mockUserCreateCall).toHaveBeenCalledWith("token", null, expect.objectContaining({
|
||||
user_email: "test@example.com",
|
||||
user_role: "proxy_user",
|
||||
}));
|
||||
});
|
||||
});
|
||||
|
||||
it("should call onUserCreated callback when user is created in embedded mode", async () => {
|
||||
const user = userEvent.setup();
|
||||
const onUserCreated = vi.fn();
|
||||
mockUserCreateCall.mockResolvedValue({ data: { user_id: "new-user-456" } });
|
||||
|
||||
renderWithProviders(
|
||||
<CreateUserButton {...defaultProps} onUserCreated={onUserCreated} possibleUIRoles={{ proxy_user: { ui_label: "User", description: "" } }} isEmbedded />,
|
||||
);
|
||||
|
||||
await user.type(screen.getByLabelText(/user email/i), "embedded@example.com");
|
||||
await user.click(screen.getByRole("combobox", { name: /user role/i }));
|
||||
await user.click(screen.getByText("User"));
|
||||
await user.click(screen.getByRole("button", { name: /create user/i }));
|
||||
|
||||
await waitFor(() => {
|
||||
expect(onUserCreated).toHaveBeenCalledWith("new-user-456");
|
||||
});
|
||||
});
|
||||
|
||||
it("should show success notification when user is created successfully in standalone mode", async () => {
|
||||
const user = userEvent.setup();
|
||||
mockUserCreateCall.mockResolvedValue({ data: { user_id: "new-user-789" } });
|
||||
mockInvitationCreateCall.mockResolvedValue({
|
||||
id: "inv-2",
|
||||
user_id: "new-user-789",
|
||||
has_user_setup_sso: false,
|
||||
} as any);
|
||||
|
||||
renderWithProviders(
|
||||
<CreateUserButton {...defaultProps} possibleUIRoles={{ proxy_user: { ui_label: "User", description: "" } }} />,
|
||||
);
|
||||
|
||||
await waitFor(() => {
|
||||
expect(screen.getByRole("button", { name: /\+ invite user/i })).toBeInTheDocument();
|
||||
});
|
||||
await user.click(screen.getByRole("button", { name: /\+ invite user/i }));
|
||||
|
||||
const dialog = screen.getByRole("dialog", { name: /invite user/i });
|
||||
await user.type(within(dialog).getByLabelText(/user email/i), "standalone@example.com");
|
||||
await user.click(within(dialog).getByRole("combobox", { name: /global proxy role/i }));
|
||||
await user.click(screen.getByText("User"));
|
||||
await user.click(within(dialog).getByRole("button", { name: /invite user/i }));
|
||||
|
||||
await waitFor(() => {
|
||||
expect(mockNotificationsManager.success).toHaveBeenCalledWith("API user Created");
|
||||
});
|
||||
});
|
||||
|
||||
it("should show error notification when user creation fails", async () => {
|
||||
const user = userEvent.setup();
|
||||
mockUserCreateCall.mockRejectedValue({ response: { data: { detail: "Email already exists" } } });
|
||||
|
||||
renderWithProviders(
|
||||
<CreateUserButton {...defaultProps} possibleUIRoles={{ proxy_user: { ui_label: "User", description: "" } }} isEmbedded />,
|
||||
);
|
||||
|
||||
await user.type(screen.getByLabelText(/user email/i), "duplicate@example.com");
|
||||
await user.click(screen.getByRole("combobox", { name: /user role/i }));
|
||||
await user.click(screen.getByText("User"));
|
||||
await user.click(screen.getByRole("button", { name: /create user/i }));
|
||||
|
||||
await waitFor(() => {
|
||||
expect(mockNotificationsManager.fromBackend).toHaveBeenCalledWith("Email already exists");
|
||||
});
|
||||
});
|
||||
|
||||
it("should show info notification when making API call", async () => {
|
||||
const user = userEvent.setup();
|
||||
mockUserCreateCall.mockResolvedValue({ data: { user_id: "new-user" } });
|
||||
mockInvitationCreateCall.mockResolvedValue({
|
||||
id: "inv-3",
|
||||
user_id: "new-user",
|
||||
has_user_setup_sso: false,
|
||||
} as any);
|
||||
|
||||
renderWithProviders(
|
||||
<CreateUserButton {...defaultProps} possibleUIRoles={{ proxy_user: { ui_label: "User", description: "" } }} isEmbedded />,
|
||||
);
|
||||
|
||||
await user.type(screen.getByLabelText(/user email/i), "info@example.com");
|
||||
await user.click(screen.getByRole("combobox", { name: /user role/i }));
|
||||
await user.click(screen.getByText("User"));
|
||||
await user.click(screen.getByRole("button", { name: /create user/i }));
|
||||
|
||||
await waitFor(() => {
|
||||
expect(mockNotificationsManager.info).toHaveBeenCalledWith("Making API Call");
|
||||
});
|
||||
});
|
||||
|
||||
it("should close modal when cancel is clicked in standalone mode", async () => {
|
||||
const user = userEvent.setup();
|
||||
renderWithProviders(<CreateUserButton {...defaultProps} />);
|
||||
|
||||
await waitFor(() => {
|
||||
expect(screen.getByRole("button", { name: /\+ invite user/i })).toBeInTheDocument();
|
||||
});
|
||||
await user.click(screen.getByRole("button", { name: /\+ invite user/i }));
|
||||
expect(screen.getByRole("dialog", { name: /invite user/i })).toBeInTheDocument();
|
||||
|
||||
const dialog = screen.getByRole("dialog", { name: /invite user/i });
|
||||
await user.click(within(dialog).getByRole("button", { name: /close/i }));
|
||||
expect(screen.queryByRole("dialog")).not.toBeInTheDocument();
|
||||
});
|
||||
|
||||
it("should show onboarding modal when user is created and SSO is disabled", async () => {
|
||||
const user = userEvent.setup();
|
||||
mockUserCreateCall.mockResolvedValue({ data: { user_id: "sso-user" } });
|
||||
mockInvitationCreateCall.mockResolvedValue({
|
||||
id: "inv-sso",
|
||||
user_id: "sso-user",
|
||||
has_user_setup_sso: false,
|
||||
} as any);
|
||||
|
||||
renderWithProviders(
|
||||
<CreateUserButton {...defaultProps} possibleUIRoles={{ proxy_user: { ui_label: "User", description: "" } }} />,
|
||||
);
|
||||
|
||||
await waitFor(() => {
|
||||
expect(screen.getByRole("button", { name: /\+ invite user/i })).toBeInTheDocument();
|
||||
});
|
||||
await user.click(screen.getByRole("button", { name: /\+ invite user/i }));
|
||||
|
||||
const dialog = screen.getByRole("dialog", { name: /invite user/i });
|
||||
await user.type(within(dialog).getByLabelText(/user email/i), "sso@example.com");
|
||||
await user.click(within(dialog).getByRole("combobox", { name: /global proxy role/i }));
|
||||
await user.click(screen.getByText("User"));
|
||||
await user.click(within(dialog).getByRole("button", { name: /invite user/i }));
|
||||
|
||||
await waitFor(() => {
|
||||
expect(mockInvitationCreateCall).toHaveBeenCalledWith("token", "sso-user");
|
||||
});
|
||||
await waitFor(() => {
|
||||
expect(mockNotificationsManager.success).toHaveBeenCalledWith("API user Created");
|
||||
});
|
||||
});
|
||||
});
|
||||
|
|
@ -1,33 +1,29 @@
|
|||
import React, { useState, useEffect } from "react";
|
||||
import { Button, Modal, Form, Input, Select, Select as Select2 } from "antd";
|
||||
import {
|
||||
Button as Button2,
|
||||
Text,
|
||||
TextInput,
|
||||
SelectItem,
|
||||
Accordion,
|
||||
AccordionHeader,
|
||||
AccordionBody,
|
||||
Title,
|
||||
} from "@tremor/react";
|
||||
import OnboardingModal from "./onboarding_link";
|
||||
import { InvitationLink } from "./onboarding_link";
|
||||
import {
|
||||
userCreateCall,
|
||||
modelAvailableCall,
|
||||
invitationCreateCall,
|
||||
getProxyUISettings,
|
||||
getProxyBaseUrl,
|
||||
} from "./networking";
|
||||
import BulkCreateUsers from "./bulk_create_users_button";
|
||||
const { Option } = Select;
|
||||
import { Tooltip } from "antd";
|
||||
import { InfoCircleOutlined } from "@ant-design/icons";
|
||||
import { getModelDisplayName } from "./key_team_helpers/fetch_available_models_team_key";
|
||||
import { InfoCircleOutlined, UserAddOutlined } from "@ant-design/icons";
|
||||
import { useQueryClient } from "@tanstack/react-query";
|
||||
import NotificationsManager from "./molecules/notifications_manager";
|
||||
import {
|
||||
Accordion,
|
||||
AccordionBody,
|
||||
AccordionHeader,
|
||||
Button as Button2,
|
||||
SelectItem,
|
||||
TextInput,
|
||||
} from "@tremor/react";
|
||||
import { Alert, Button, Form, Input, Modal, Select, Select as Select2, Space, Tooltip, Typography } from "antd";
|
||||
import React, { useEffect, useState } from "react";
|
||||
import BulkCreateUsers from "./bulk_create_users_button";
|
||||
import TeamDropdown from "./common_components/team_dropdown";
|
||||
|
||||
import { getModelDisplayName } from "./key_team_helpers/fetch_available_models_team_key";
|
||||
import NotificationsManager from "./molecules/notifications_manager";
|
||||
import {
|
||||
getProxyBaseUrl,
|
||||
getProxyUISettings,
|
||||
invitationCreateCall,
|
||||
modelAvailableCall,
|
||||
userCreateCall,
|
||||
} from "./networking";
|
||||
import OnboardingModal, { InvitationLink } from "./onboarding_link";
|
||||
const { Option } = Select;
|
||||
const { Text, Link, Title } = Typography;
|
||||
// Helper function to generate UUID compatible across all environments
|
||||
const generateUUID = (): string => {
|
||||
if (typeof crypto !== "undefined" && crypto.randomUUID) {
|
||||
|
|
@ -58,14 +54,8 @@ interface UISettings {
|
|||
SSO_ENABLED: boolean;
|
||||
}
|
||||
|
||||
const Createuser: React.FC<CreateuserProps> = ({
|
||||
userID,
|
||||
accessToken,
|
||||
teams,
|
||||
possibleUIRoles,
|
||||
onUserCreated,
|
||||
isEmbedded = false,
|
||||
}) => {
|
||||
export const CreateUserButton: React.FC<CreateuserProps> = ({
|
||||
userID, accessToken, teams, possibleUIRoles, onUserCreated, isEmbedded = false }) => {
|
||||
const queryClient = useQueryClient();
|
||||
const [uiSettings, setUISettings] = useState<UISettings | null>(null);
|
||||
const [form] = Form.useForm();
|
||||
|
|
@ -75,28 +65,18 @@ const Createuser: React.FC<CreateuserProps> = ({
|
|||
const [isInvitationLinkModalVisible, setIsInvitationLinkModalVisible] = useState(false);
|
||||
const [invitationLinkData, setInvitationLinkData] = useState<InvitationLink | null>(null);
|
||||
const [baseUrl, setBaseUrl] = useState<string | null>(null);
|
||||
// get all models
|
||||
useEffect(() => {
|
||||
const fetchData = async () => {
|
||||
try {
|
||||
const userRole = "any"; // You may need to get the user role dynamically
|
||||
const userRole = "any";
|
||||
const modelDataResponse = await modelAvailableCall(accessToken, userID, userRole);
|
||||
// Assuming modelDataResponse.data contains an array of model objects with a 'model_name' property
|
||||
const availableModels = [];
|
||||
for (let i = 0; i < modelDataResponse.data.length; i++) {
|
||||
const model = modelDataResponse.data[i];
|
||||
availableModels.push(model.id);
|
||||
}
|
||||
console.log("Model data response:", modelDataResponse.data);
|
||||
console.log("Available models:", availableModels);
|
||||
|
||||
// Assuming modelDataResponse.data contains an array of model names
|
||||
setUserModels(availableModels);
|
||||
|
||||
// get ui settings
|
||||
const uiSettingsResponse = await getProxyUISettings(accessToken);
|
||||
console.log("uiSettingsResponse:", uiSettingsResponse);
|
||||
|
||||
setUISettings(uiSettingsResponse);
|
||||
} catch (error) {
|
||||
console.error("Error fetching model data:", error);
|
||||
|
|
@ -104,9 +84,8 @@ const Createuser: React.FC<CreateuserProps> = ({
|
|||
};
|
||||
|
||||
setBaseUrl(getProxyBaseUrl());
|
||||
|
||||
fetchData(); // Call the function to fetch model data when the component mounts
|
||||
}, []); // Empty dependency array to run only once
|
||||
fetchData();
|
||||
}, []);
|
||||
|
||||
const handleOk = () => {
|
||||
setIsModalVisible(false);
|
||||
|
|
@ -126,25 +105,19 @@ const Createuser: React.FC<CreateuserProps> = ({
|
|||
setIsModalVisible(true);
|
||||
}
|
||||
if ((!formValues.models || formValues.models.length === 0) && formValues.user_role !== "proxy_admin") {
|
||||
console.log("formValues.user_role", formValues.user_role);
|
||||
// If models is empty or undefined, set it to "no-default-models"
|
||||
formValues.models = ["no-default-models"];
|
||||
}
|
||||
console.log("formValues in create user:", formValues);
|
||||
const response = await userCreateCall(accessToken, null, formValues);
|
||||
await queryClient.invalidateQueries({ queryKey: ["userList"] });
|
||||
console.log("user create Response:", response);
|
||||
setApiuser(true);
|
||||
const user_id = response.data?.user_id || response.user_id;
|
||||
|
||||
// Call the callback if provided (for embedded mode)
|
||||
if (onUserCreated && isEmbedded) {
|
||||
onUserCreated(user_id);
|
||||
form.resetFields();
|
||||
return; // Skip the invitation flow when embedded
|
||||
return;
|
||||
}
|
||||
|
||||
// only do invite link flow if sso is not enabled
|
||||
if (!uiSettings?.SSO_ENABLED) {
|
||||
invitationCreateCall(accessToken, user_id).then((data) => {
|
||||
data.has_user_setup_sso = false;
|
||||
|
|
@ -184,6 +157,21 @@ const Createuser: React.FC<CreateuserProps> = ({
|
|||
if (isEmbedded) {
|
||||
return (
|
||||
<Form form={form} onFinish={handleCreate} labelCol={{ span: 8 }} wrapperCol={{ span: 16 }} labelAlign="left">
|
||||
<Alert
|
||||
message="Email invitations"
|
||||
description={
|
||||
<>
|
||||
New users receive an email invite only when an email integration (SMTP, Resend, or SendGrid) is configured.
|
||||
{" "}
|
||||
<Link href="https://docs.litellm.ai/docs/proxy/email" target="_blank">
|
||||
Learn how to set up email notifications
|
||||
</Link>
|
||||
</>
|
||||
}
|
||||
type="info"
|
||||
showIcon
|
||||
className="mb-4"
|
||||
/>
|
||||
<Form.Item label="User Email" name="user_email">
|
||||
<TextInput placeholder="" />
|
||||
</Form.Item>
|
||||
|
|
@ -194,9 +182,9 @@ const Createuser: React.FC<CreateuserProps> = ({
|
|||
<SelectItem key={role} value={role} title={ui_label}>
|
||||
<div className="flex">
|
||||
{ui_label}{" "}
|
||||
<p className="ml-2" style={{ color: "gray", fontSize: "12px" }}>
|
||||
<Text className="ml-2" style={{ color: "gray", fontSize: "12px" }}>
|
||||
{description}
|
||||
</p>
|
||||
</Text>
|
||||
</div>
|
||||
</SelectItem>
|
||||
))}
|
||||
|
|
@ -234,16 +222,33 @@ const Createuser: React.FC<CreateuserProps> = ({
|
|||
onOk={handleOk}
|
||||
onCancel={handleCancel}
|
||||
>
|
||||
<Text className="mb-1">Create a User who can own keys</Text>
|
||||
<Space direction="vertical" size="middle">
|
||||
<Text className="mb-1">Create a User who can own keys</Text>
|
||||
<Alert
|
||||
message="Email invitations"
|
||||
description={
|
||||
<>
|
||||
New users receive an email invite only when an email integration (SMTP, Resend, or SendGrid) is configured.
|
||||
{" "}
|
||||
<Link href="https://docs.litellm.ai/docs/proxy/email" target="_blank">
|
||||
Learn how to set up email notifications
|
||||
</Link>
|
||||
</>
|
||||
}
|
||||
type="info"
|
||||
showIcon
|
||||
className="mb-4"
|
||||
/>
|
||||
</Space>
|
||||
<Form form={form} onFinish={handleCreate} labelCol={{ span: 8 }} wrapperCol={{ span: 16 }} labelAlign="left">
|
||||
<Form.Item label="User Email" name="user_email">
|
||||
<TextInput placeholder="" />
|
||||
<Input />
|
||||
</Form.Item>
|
||||
<Form.Item
|
||||
label={
|
||||
<span>
|
||||
Global Proxy Role{" "}
|
||||
<Tooltip title="This is the role that the user will globally on the proxy. This role is independent of any team/org specific roles.">
|
||||
<Tooltip title="This role is independent of any team/org specific roles. Configure Team / Organization Admins in the Settings">
|
||||
<InfoCircleOutlined />
|
||||
</Tooltip>
|
||||
</span>
|
||||
|
|
@ -254,12 +259,12 @@ const Createuser: React.FC<CreateuserProps> = ({
|
|||
{possibleUIRoles &&
|
||||
Object.entries(possibleUIRoles).map(([role, { ui_label, description }]) => (
|
||||
<SelectItem key={role} value={role} title={ui_label}>
|
||||
<div className="flex">
|
||||
{ui_label}{" "}
|
||||
<p className="ml-2" style={{ color: "gray", fontSize: "12px" }}>
|
||||
{description}
|
||||
</p>
|
||||
</div>
|
||||
<Text>
|
||||
{ui_label}
|
||||
</Text>
|
||||
<Text type="secondary">
|
||||
{" - "}{description}
|
||||
</Text>
|
||||
</SelectItem>
|
||||
))}
|
||||
</Select2>
|
||||
|
|
@ -279,7 +284,7 @@ const Createuser: React.FC<CreateuserProps> = ({
|
|||
</Form.Item>
|
||||
<Accordion>
|
||||
<AccordionHeader>
|
||||
<Title>Personal Key Creation</Title>
|
||||
<Text strong>Personal Key Creation</Text>
|
||||
</AccordionHeader>
|
||||
<AccordionBody>
|
||||
<Form.Item
|
||||
|
|
@ -312,7 +317,7 @@ const Createuser: React.FC<CreateuserProps> = ({
|
|||
</AccordionBody>
|
||||
</Accordion>
|
||||
<div style={{ textAlign: "right", marginTop: "10px" }}>
|
||||
<Button htmlType="submit">Create User</Button>
|
||||
<Button type="primary" icon={<UserAddOutlined />} htmlType="submit">Invite User</Button>
|
||||
</div>
|
||||
</Form>
|
||||
</Modal>
|
||||
|
|
@ -326,6 +331,4 @@ const Createuser: React.FC<CreateuserProps> = ({
|
|||
)}
|
||||
</div>
|
||||
);
|
||||
};
|
||||
|
||||
export default Createuser;
|
||||
};
|
||||
|
|
@ -1,35 +0,0 @@
|
|||
import React from "react";
|
||||
import { render } from "@testing-library/react";
|
||||
import { describe, it, expect, vi } from "vitest";
|
||||
import { QueryClient, QueryClientProvider } from "@tanstack/react-query";
|
||||
import Createuser from "./create_user_button";
|
||||
|
||||
vi.mock("./networking", () => ({
|
||||
userCreateCall: vi.fn(),
|
||||
modelAvailableCall: vi.fn().mockResolvedValue({ data: [] }),
|
||||
invitationCreateCall: vi.fn(),
|
||||
getProxyUISettings: vi.fn().mockResolvedValue({
|
||||
PROXY_BASE_URL: null,
|
||||
PROXY_LOGOUT_URL: null,
|
||||
DEFAULT_TEAM_DISABLED: false,
|
||||
SSO_ENABLED: false,
|
||||
}),
|
||||
getProxyBaseUrl: vi.fn().mockReturnValue("http://localhost"),
|
||||
}));
|
||||
|
||||
const createQueryClient = () =>
|
||||
new QueryClient({
|
||||
defaultOptions: { queries: { retry: false, gcTime: 0 } },
|
||||
});
|
||||
|
||||
describe("Create User Button", () => {
|
||||
it("should render the create user button", () => {
|
||||
const qc = createQueryClient();
|
||||
const { getByText } = render(
|
||||
<QueryClientProvider client={qc}>
|
||||
<Createuser userID="123" accessToken="123" teams={[]} possibleUIRoles={{}} isEmbedded />
|
||||
</QueryClientProvider>,
|
||||
);
|
||||
expect(getByText("Create User")).toBeInTheDocument();
|
||||
});
|
||||
});
|
||||
|
|
@ -21,7 +21,7 @@ import PremiumLoggingSettings from "../common_components/PremiumLoggingSettings"
|
|||
import RateLimitTypeFormItem from "../common_components/RateLimitTypeFormItem";
|
||||
import RouterSettingsAccordion, { RouterSettingsAccordionValue } from "../common_components/RouterSettingsAccordion";
|
||||
import TeamDropdown from "../common_components/team_dropdown";
|
||||
import Createuser from "../create_user_button";
|
||||
import { CreateUserButton } from "../CreateUserButton";
|
||||
import { getModelDisplayName } from "../key_team_helpers/fetch_available_models_team_key";
|
||||
import { Team } from "../key_team_helpers/key_list";
|
||||
import MCPServerSelector from "../mcp_server_management/MCPServerSelector";
|
||||
|
|
@ -1347,7 +1347,7 @@ const CreateKey: React.FC<CreateKeyProps> = ({ team, teams, data, addKey }) => {
|
|||
footer={null}
|
||||
width={800}
|
||||
>
|
||||
<Createuser
|
||||
<CreateUserButton
|
||||
userID={userID}
|
||||
accessToken={accessToken}
|
||||
teams={teams}
|
||||
|
|
|
|||
|
|
@ -3,7 +3,7 @@ import React, { useEffect, useState } from "react";
|
|||
|
||||
import { Button } from "@tremor/react";
|
||||
import BulkEditUserModal from "./BulkEditUsers";
|
||||
import CreateUser from "./create_user_button";
|
||||
import { CreateUserButton } from "./CreateUserButton";
|
||||
import EditUserModal from "./edit_user";
|
||||
import {
|
||||
getPossibleUserRoles,
|
||||
|
|
@ -301,7 +301,7 @@ const ViewUserDashboard: React.FC<ViewUserDashboardProps> = ({ accessToken, toke
|
|||
</>
|
||||
) : userID && accessToken ? (
|
||||
<>
|
||||
<CreateUser userID={userID} accessToken={accessToken} teams={teams} possibleUIRoles={possibleUIRoles} />
|
||||
<CreateUserButton userID={userID} accessToken={accessToken} teams={teams} possibleUIRoles={possibleUIRoles} />
|
||||
|
||||
<Button
|
||||
onClick={handleToggleSelectionMode}
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue