From eeedf98d4004b44f11ec0069fe32661165f79611 Mon Sep 17 00:00:00 2001 From: shrey-berri Date: Fri, 23 Oct 2026 16:04:09 -0700 Subject: [PATCH 1/7] fix(azure_ai): drop tool_reference parts from tool messages on chat (#48032) --- litellm/llms/azure_ai/chat/transformation.py | 3 ++- .../chat/test_azure_ai_transformation.py | 26 +++++++++++++++++++ 2 files changed, 28 insertions(+), 1 deletion(-) diff --git a/litellm/llms/azure_ai/chat/transformation.py b/litellm/llms/azure_ai/chat/transformation.py index 0fa54cccbc9..83d00a621b4 100644 --- a/litellm/llms/azure_ai/chat/transformation.py +++ b/litellm/llms/azure_ai/chat/transformation.py @@ -12,6 +12,7 @@ from litellm._logging import verbose_logger from litellm.litellm_core_utils.prompt_templates.common_utils import ( audio_or_image_in_message_content, convert_content_list_to_str, + drop_tool_reference_parts_from_tool_messages, filter_value_from_dict, ) from litellm.llms.azure.common_utils import BaseAzureLLM @@ -250,7 +251,7 @@ class AzureAIStudioConfig(OpenAIConfig): Operates on a deep copy so the caller's messages keep their thinking blocks and provider metadata, which a fallback to another provider still needs. """ - stripped_messages: Final = copy.deepcopy(messages) + stripped_messages: Final = copy.deepcopy(drop_tool_reference_parts_from_tool_messages(messages)) for message in stripped_messages: message_dict = cast(dict, message) # cast-ok: TypedDict is a runtime dict stripped on our copy for field in NON_OPENAI_SPEC_MESSAGE_FIELDS: diff --git a/tests/unit/llms/azure_ai/chat/test_azure_ai_transformation.py b/tests/unit/llms/azure_ai/chat/test_azure_ai_transformation.py index 5186541a62b..ccac6ffc309 100644 --- a/tests/unit/llms/azure_ai/chat/test_azure_ai_transformation.py +++ b/tests/unit/llms/azure_ai/chat/test_azure_ai_transformation.py @@ -879,6 +879,32 @@ def test_azure_ai_stripping_does_not_mutate_caller_messages(): } +def test_azure_ai_reference_only_tool_result_is_sent_as_empty_text(): + messages: Final = [ + {"role": "user", "content": "load the WebFetch tool"}, + { + "role": "assistant", + "content": None, + "tool_calls": [{"id": "call_1", "type": "function", "function": {"name": "ToolSearch", "arguments": "{}"}}], + }, + { + "role": "tool", + "tool_call_id": "call_1", + "content": [{"type": "tool_reference", "tool_name": "WebFetch"}], + }, + ] + + request: Final = AzureAIStudioConfig().transform_request( + model="fw-glm-5.2", + messages=messages, + optional_params={}, + litellm_params={}, + headers={}, + ) + + assert request["messages"][2] == {"role": "tool", "tool_call_id": "call_1", "content": ""} + + @pytest.mark.parametrize( "api_base, expected_url", [ From 91287b49f403094a754c2a3b979001d1d0fad9d1 Mon Sep 17 00:00:00 2001 From: moe-berri Date: Fri, 23 Oct 2026 16:12:51 -0700 Subject: [PATCH 2/7] feat(ui): sort every users table data column (#48105) * feat(ui): sort every users table data column * fix(ui): order unlimited user budgets as infinity and verify org scope * fix(ui): cover user sort queries and document pagination bound --- .../internal_user_endpoints.py | 50 ++++- .../management_helpers/user_list_sorting.py | 77 ++++++++ .../management/test_user_list_sorting.py | 179 ++++++++++++++++++ .../test_user_list_sorting.py | 71 +++++++ ...st.tsx => view_users.integration.test.tsx} | 15 +- ...st.tsx => UsersTable.integration.test.tsx} | 19 +- .../view_users/UsersTableColumns.tsx | 33 ++-- 7 files changed, 415 insertions(+), 29 deletions(-) create mode 100644 litellm/proxy/management_helpers/user_list_sorting.py create mode 100644 tests/integration/management/test_user_list_sorting.py create mode 100644 tests/unit/proxy/management_helpers/test_user_list_sorting.py rename ui/litellm-dashboard/src/app/(dashboard)/users/_components/{view_users.test.tsx => view_users.integration.test.tsx} (96%) rename ui/litellm-dashboard/src/app/(dashboard)/users/_components/view_users/{UsersTable.test.tsx => UsersTable.integration.test.tsx} (96%) diff --git a/litellm/proxy/management_endpoints/internal_user_endpoints.py b/litellm/proxy/management_endpoints/internal_user_endpoints.py index b563dbc5c4c..d6bf9220551 100644 --- a/litellm/proxy/management_endpoints/internal_user_endpoints.py +++ b/litellm/proxy/management_endpoints/internal_user_endpoints.py @@ -2093,7 +2093,9 @@ async def get_user_key_counts( return result -def _validate_sort_params(sort_by: str | None, sort_order: str) -> dict[str, str] | None: +def _validate_sort_params( + sort_by: str | None, sort_order: str +) -> Mapping[str, object] | Sequence[Mapping[str, object]] | None: order_by: Final[dict[str, str]] = {} if sort_by is None: @@ -2106,6 +2108,11 @@ def _validate_sort_params(sort_by: str | None, sort_order: str) -> dict[str, str "spend", "user_alias", "user_role", + "max_budget", + "updated_at", + "sso_user_id", + "status", + "key_count", ] if sort_by not in valid_columns: raise HTTPException( @@ -2120,6 +2127,11 @@ def _validate_sort_params(sort_by: str | None, sort_order: str) -> dict[str, str detail={"error": "Invalid sort order. Must be 'asc' or 'desc'"}, ) + if sort_by in {"status", "key_count"}: + return None + if sort_by in {"max_budget", "updated_at", "sso_user_id"}: + nulls: Final = "first" if sort_by == "max_budget" and sort_order.lower() == "desc" else "last" + return ({sort_by: {"sort": sort_order.lower(), "nulls": nulls}}, {"user_id": "asc"}) order_by[sort_by] = sort_order.lower() return order_by @@ -2338,16 +2350,40 @@ async def get_users( # Build order_by conditions - order_by: Final[dict[str, str] | None] = ( + order_by: Final = ( _validate_sort_params(sort_by, sort_order) if sort_by is not None and isinstance(sort_by, str) else None ) - users: Final[Sequence[prisma_models.LiteLLM_UserTable]] = await UserRepository(prisma_client).table.find_many( - where=where, - skip=skip, - take=page_size, - order=(order_by if order_by else {"created_at": "desc"}), # Default to created_at desc if no sort specified + from litellm.proxy.management_helpers.user_list_sorting import sorted_user_ids + + computed_sort: Final = sort_by in ("status", "key_count") + sorted_ids: Final = ( + await sorted_user_ids( + prisma_client, + sort_by="status" if sort_by == "status" else "key_count", + descending=sort_order.lower() == "desc", + skip=skip, + take=page_size, + role=role, + user_ids=user_ids, + user_email=user_email, + team=team, + sso_user_ids=sso_user_ids, + organization_ids=organization_ids, + search=search, + ) + if computed_sort + else () ) + page_users: Final = await UserRepository(prisma_client).table.find_many( + # bounded-ok: sorted_ids comes from LIMIT page_size, capped at 100 by the route + where={"AND": (where, {"user_id": {"in": sorted_ids}})} if computed_sort else where, + skip=0 if computed_sort else skip, + take=page_size, + order=order_by if order_by else {"created_at": "desc"}, + ) + positions: Final = {user_id: index for index, user_id in enumerate(sorted_ids)} + users: Final = sorted(page_users, key=lambda user: positions[user.user_id]) if computed_sort else page_users # Get total count of user rows total_count: Final[int] = await UserRepository(prisma_client).table.count(where=where) diff --git a/litellm/proxy/management_helpers/user_list_sorting.py b/litellm/proxy/management_helpers/user_list_sorting.py new file mode 100644 index 00000000000..d3d2e323cb0 --- /dev/null +++ b/litellm/proxy/management_helpers/user_list_sorting.py @@ -0,0 +1,77 @@ +from collections.abc import Sequence +from typing import Final, Literal + +from pydantic import BaseModel, ConfigDict, TypeAdapter + +from litellm.constants import UI_SESSION_TOKEN_TEAM_ID +from litellm.proxy.utils import PrismaClient + + +class _UserId(BaseModel): + model_config = ConfigDict(frozen=True) + user_id: str + + +_USER_IDS: Final = TypeAdapter(tuple[_UserId, ...]) + + +async def sorted_user_ids( + prisma_client: PrismaClient, + *, + sort_by: Literal["status", "key_count"], + descending: bool, + skip: int, + take: int, + role: str | None, + user_ids: str | None, + user_email: str | None, + team: str | None, + sso_user_ids: str | None, + organization_ids: str | None, + search: str | None, +) -> Sequence[str]: + ids: Final = tuple(value.strip() for value in (user_ids or "").split(",") if value.strip()) + sso_ids: Final = tuple(value.strip() for value in (sso_user_ids or "").split(",") if value.strip()) + org_ids: Final = tuple(value.strip() for value in (organization_ids or "").split(",") if value.strip()) + expression: Final = ( + "CASE WHEN u.metadata->'scim_active' = 'false'::jsonb THEN 1 ELSE 0 END" + if sort_by == "status" + else '(SELECT COUNT(*) FROM "LiteLLM_VerificationToken" k ' + "WHERE k.user_id = u.user_id AND k.team_id IS DISTINCT FROM $11::text)" + ) + direction: Final = "DESC" if descending else "ASC" + sql: Final = f""" + SELECT u.user_id FROM "LiteLLM_UserTable" u + WHERE ($1::text IS NULL OR u.user_role = $1) + AND ($2::text[] IS NULL + OR CASE WHEN cardinality($2::text[]) = 1 + THEN u.user_id ILIKE '%' || ($2::text[])[1] || '%' + ELSE u.user_id = ANY($2::text[]) END) + AND ($3::text IS NULL OR u.user_email ILIKE '%' || $3 || '%') + AND ($4::text IS NULL OR $4 = ANY(u.teams)) + AND ($5::boolean OR u.sso_user_id = ANY($6::text[])) + AND (cardinality($7::text[]) = 0 OR EXISTS ( + SELECT 1 FROM "LiteLLM_OrganizationMembership" m + WHERE m.user_id = u.user_id AND m.organization_id = ANY($7::text[]))) + AND ($8::text IS NULL OR u.user_id ILIKE '%' || $8 || '%' + OR u.user_email ILIKE '%' || $8 || '%') + ORDER BY {expression} {direction}, u.user_id ASC + LIMIT $9 OFFSET $10 + """ + rows: Final = _USER_IDS.validate_python( + await prisma_client.db.query_raw( + sql, + role or None, + ids if user_ids else None, + user_email, + team, + sso_user_ids is None, + sso_ids, + org_ids, + search or None, + take, + skip, + *((UI_SESSION_TOKEN_TEAM_ID,) if sort_by == "key_count" else ()), + ) + ) + return tuple(row.user_id for row in rows) diff --git a/tests/integration/management/test_user_list_sorting.py b/tests/integration/management/test_user_list_sorting.py new file mode 100644 index 00000000000..f5578140e1c --- /dev/null +++ b/tests/integration/management/test_user_list_sorting.py @@ -0,0 +1,179 @@ +from itertools import chain +from typing import Final +from uuid import uuid4 + +import pytest + +from litellm.constants import UI_SESSION_TOKEN_TEAM_ID +from tests.integration._support.client import Gateway, Scenario +from tests.integration._support.database import write_rows + + +@pytest.mark.parametrize("field", ("user_alias", "max_budget", "updated_at", "sso_user_id", "status", "key_count")) +@pytest.mark.parametrize("direction", ("asc", "desc")) +def test_user_list_sorts_all_pages(field: str, direction: str, gateway: Gateway, scenario: Scenario) -> None: + prefix: Final = f"sort-{uuid4().hex}" + ids: Final = tuple( + scenario.user( + user_id=f"{prefix}-{index}", + user_role="internal_user", + user_alias=f"Alias {value if index != 2 else 4:03}", + user_email=f"{prefix}-{index}@example.com", + max_budget=value, + metadata={"scim_active": value != 3}, + ) + for index, value in enumerate((20, 3, 3, 100, 0)) + ) + for user_id, value in zip(ids, (20, 3, 3, 100, 0), strict=True): + write_rows( + 'UPDATE "LiteLLM_UserTable" SET sso_user_id = %s, ' + "updated_at = timestamp '2030-01-01' + %s::int * interval '1 day' WHERE user_id = %s", + (f"{prefix}-{value:03}", str(value), user_id), + ) if value != 3 or user_id == ids[1] else write_rows( + "UPDATE \"LiteLLM_UserTable\" SET sso_user_id = %s, updated_at = timestamp '2030-01-04' WHERE user_id = %s", + (f"{prefix}-004", user_id), + ) + for index in range(value if value < 4 else 1): + scenario.key(user_id=user_id, key_alias=f"{user_id}-key-{index}") + write_rows( + 'UPDATE "LiteLLM_UserTable" SET max_budget = NULL, updated_at = NULL, sso_user_id = NULL WHERE user_id = %s', + (ids[-1],), + ) + scenario.key(user_id=ids[-1]) + write_rows( + 'UPDATE "LiteLLM_VerificationToken" SET team_id = %s WHERE user_id = %s', (UI_SESSION_TOKEN_TEAM_ID, ids[-1]) + ) + values: Final = ( + (0, 1, 1, 0, 0) + if field == "status" + else (1, 3, 3, 1, 0) + if field == "key_count" + else (20, 3, 4, 100, None) + if field == "sso_user_id" + else (20, 3, 4, 100, 0) + if field == "user_alias" + else (20, 3, 3, 100, None) + ) + expected: Final = tuple( + user_id + for _, user_id in sorted( + ((value, user_id) for value, user_id in zip(values, ids, strict=True)), + key=lambda item: ( + (item[0] is not None if field == "max_budget" and direction == "desc" else item[0] is None), + (item[0] or 0) * (1 if direction == "asc" else -1), + item[1], + ), + ) + ) + pages: Final = tuple( + gateway.request( + "GET", + "/user/list", + params={ + "search": prefix, + "sort_by": field, + "sort_order": direction, + "page_size": "2", + "page": str(page), + }, + ) + for page in (1, 2, 3) + ) + assert all(page.status_code == 200 for page in pages), [page.text for page in pages] + assert all(page.json()["total"] == len(ids) for page in pages) + assert tuple(user["user_id"] for user in chain.from_iterable(page.json()["users"] for page in pages)) == expected + + +@pytest.mark.parametrize("field", ("status", "key_count")) +def test_computed_user_sorts_preserve_filters(field: str, gateway: Gateway, scenario: Scenario) -> None: + team: Final = scenario.team() + member: Final = scenario.member(team) + other: Final = scenario.user(user_role="internal_user") + for filters in ({"team": team}, {"user_ids": member}, {"user_ids": f"{member},{other}"}, {"role": "internal_user"}): + baseline: Final = gateway.get("/user/list", {"user_ids": f"{member},{other}", **filters, "sort_by": "user_id"}) + response: Final = gateway.get("/user/list", {"user_ids": f"{member},{other}", **filters, "sort_by": field}) + assert response["total"] == baseline["total"] + assert sorted(user["user_id"] for user in response["users"]) == sorted( + user["user_id"] for user in baseline["users"] + ) + ordinary_key: Final = scenario.key(user_id=other) + denied: Final = gateway.request("GET", "/user/list", key=ordinary_key, params={"sort_by": field}) + assert denied.status_code == 403, denied.text + + +@pytest.mark.parametrize("field", ("status", "key_count")) +@pytest.mark.parametrize("direction", ("asc", "desc")) +def test_computed_user_sorts_preserve_org_admin_scope( + field: str, direction: str, gateway: Gateway, scenario: Scenario +) -> None: + prefix: Final = f"sort-org-users-{uuid4().hex}" + owner: Final = scenario.user(user_role="internal_user") + members: Final = tuple( + scenario.user(user_id=f"{prefix}-{index}", user_role="internal_user", metadata={"scim_active": index != 1}) + for index in range(4) + ) + budget: Final = scenario.budget(max_budget=10) + orgs: Final = (f"{prefix}-allowed", f"{prefix}-outside") + try: + for org in orgs: + write_rows( + 'INSERT INTO "LiteLLM_OrganizationTable" ' + "(organization_id, organization_alias, budget_id, created_by, updated_by, updated_at) " + "VALUES (%s, %s, %s, 'sorting-test', 'sorting-test', NOW())", + (org, org, budget), + ) + write_rows( + 'INSERT INTO "LiteLLM_OrganizationMembership" (user_id, organization_id, user_role) ' + "VALUES (%s, %s, 'org_admin')", + (owner, orgs[0]), + ) + for index, member in enumerate(members): + write_rows( + 'INSERT INTO "LiteLLM_OrganizationMembership" (user_id, organization_id, user_role) ' + "VALUES (%s, %s, 'internal_user')", + (member, orgs[0] if index < 3 else orgs[1]), + ) + for _ in range(index): + scenario.key(user_id=member) + caller: Final = scenario.key(user_id=owner) + values: Final = (0, 1, 0) if field == "status" else (0, 1, 2) + expected: Final = tuple( + member + for value, member in sorted( + zip(values, members[:3], strict=True), + key=lambda item: (item[0] * (-1 if direction == "desc" else 1), item[1]), + ) + ) + pages: Final = tuple( + gateway.request( + "GET", + "/user/list", + key=caller, + params={ + "search": prefix, + "sort_by": field, + "sort_order": direction, + "page_size": "2", + "page": str(page), + }, + ) + for page in (1, 2) + ) + assert all(page.status_code == 200 for page in pages), [page.text for page in pages] + assert all(page.json()["total"] == 3 for page in pages) + returned: Final = tuple(row["user_id"] for row in chain.from_iterable(page.json()["users"] for page in pages)) + assert returned == expected + denied: Final = gateway.request( + "GET", + "/user/list", + key=caller, + params={ + "organization_ids": orgs[1], + "sort_by": field, + }, + ) + assert denied.status_code == 403, denied.text + finally: + for org in orgs: + write_rows('DELETE FROM "LiteLLM_OrganizationMembership" WHERE organization_id = %s', (org,)) + write_rows('DELETE FROM "LiteLLM_OrganizationTable" WHERE organization_id = %s', (org,)) diff --git a/tests/unit/proxy/management_helpers/test_user_list_sorting.py b/tests/unit/proxy/management_helpers/test_user_list_sorting.py new file mode 100644 index 00000000000..263c1691b23 --- /dev/null +++ b/tests/unit/proxy/management_helpers/test_user_list_sorting.py @@ -0,0 +1,71 @@ +from typing import Final, Literal +from unittest.mock import AsyncMock, MagicMock + +import pytest + +from litellm.constants import UI_SESSION_TOKEN_TEAM_ID +from litellm.proxy.management_helpers.user_list_sorting import sorted_user_ids + + +@pytest.mark.parametrize("sort_by", ("status", "key_count")) +@pytest.mark.parametrize("descending", (False, True)) +@pytest.mark.asyncio +async def test_user_sort_binds_scope_search_and_page_and_excludes_ui_sessions( + sort_by: Literal["status", "key_count"], descending: bool +) -> None: + prisma: Final = MagicMock() + prisma.db.query_raw = AsyncMock(return_value=[]) + search: Final = "x' OR TRUE --" + result: Final = await sorted_user_ids( + prisma, + sort_by=sort_by, + descending=descending, + skip=25, + take=25, + role="internal_user", + user_ids=" a, b ", + user_email=search, + team="team-1", + sso_user_ids="sso-1, sso-2", + organization_ids="org-1, org-2", + search=search, + ) + assert result == () + sql, *bindings = prisma.db.query_raw.await_args.args + assert search not in sql + assert bindings == [ + "internal_user", + ("a", "b"), + search, + "team-1", + False, + ("sso-1", "sso-2"), + ("org-1", "org-2"), + search, + 25, + 25, + ] + ([UI_SESSION_TOKEN_TEAM_ID] if sort_by == "key_count" else []) + assert ("DESC" if descending else "ASC") + ", u.user_id ASC" in sql + assert ("k.team_id IS DISTINCT FROM $11::text" in sql) == (sort_by == "key_count") + + +@pytest.mark.asyncio +async def test_missing_user_filters_preserve_unfiltered_query() -> None: + prisma: Final = MagicMock() + prisma.db.query_raw = AsyncMock(return_value=[]) + result: Final = await sorted_user_ids( + prisma, + sort_by="status", + descending=False, + skip=0, + take=25, + role=None, + user_ids=None, + user_email=None, + team=None, + sso_user_ids=None, + organization_ids=None, + search=None, + ) + assert result == () + assert prisma.db.query_raw.await_args.args[1:] == (None, None, None, None, True, (), (), None, 25, 0) diff --git a/ui/litellm-dashboard/src/app/(dashboard)/users/_components/view_users.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/users/_components/view_users.integration.test.tsx similarity index 96% rename from ui/litellm-dashboard/src/app/(dashboard)/users/_components/view_users.test.tsx rename to ui/litellm-dashboard/src/app/(dashboard)/users/_components/view_users.integration.test.tsx index 13958a26adf..d728f6473f5 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/users/_components/view_users.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/users/_components/view_users.integration.test.tsx @@ -307,7 +307,14 @@ describe("ViewUserDashboard", () => { expect(sortOrder).toBe("desc"); }); - it("sends the clicked column as sort_by and resets to the first page", async () => { + it.each([ + ["user_alias", "asc"], + ["max_budget", "desc"], + ["updated_at", "asc"], + ["status", "asc"], + ["key_count", "desc"], + ["sso_user_id", "asc"], + ])("sorts %s through the API and resets pagination", async (field, direction) => { const user = userEvent.setup(); renderDashboard(); @@ -315,12 +322,12 @@ describe("ViewUserDashboard", () => { expect(screen.getByText("test@example.com")).toBeInTheDocument(); }); - await user.click(screen.getByTestId("sort-header-user_email")); + await user.click(screen.getByTestId(`sort-header-${field}`)); await waitFor(() => { const latest = userListCall.mock.calls[userListCall.mock.calls.length - 1]; - expect(latest[8]).toBe("user_email"); - expect(latest[9]).toBe("asc"); + expect(latest[8]).toBe(field); + expect(latest[9]).toBe(direction); expect(latest[2]).toBe(1); }); }); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/users/_components/view_users/UsersTable.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/users/_components/view_users/UsersTable.integration.test.tsx similarity index 96% rename from ui/litellm-dashboard/src/app/(dashboard)/users/_components/view_users/UsersTable.test.tsx rename to ui/litellm-dashboard/src/app/(dashboard)/users/_components/view_users/UsersTable.integration.test.tsx index b9ed37a2152..85d9ef386a3 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/users/_components/view_users/UsersTable.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/users/_components/view_users/UsersTable.integration.test.tsx @@ -132,10 +132,7 @@ describe("UsersTable", () => { expect(screen.getByText("$98.85")).toBeInTheDocument(); }); - // Sorting is server-side and the backend only accepts these five keys, so a sort - // control on any other column would send an invalid sort_by. Assert the exact set: - // a missing control and an extra one both have to fail. - it("exposes a sort control for exactly the five server-sortable columns", () => { + it("exposes sorting for every data column", () => { render(); const sortableIds = screen @@ -143,7 +140,19 @@ describe("UsersTable", () => { .map((node) => (node.getAttribute("data-testid") ?? "").replace("sort-header-", "")) .sort(); - expect(sortableIds).toEqual(["created_at", "spend", "user_email", "user_id", "user_role"]); + expect(sortableIds).toEqual([ + "created_at", + "key_count", + "max_budget", + "spend", + "sso_user_id", + "status", + "updated_at", + "user_alias", + "user_email", + "user_id", + "user_role", + ]); }); it("reports the clicked column to the server sorting handler", async () => { diff --git a/ui/litellm-dashboard/src/app/(dashboard)/users/_components/view_users/UsersTableColumns.tsx b/ui/litellm-dashboard/src/app/(dashboard)/users/_components/view_users/UsersTableColumns.tsx index ad0cdbbd9c1..f2593f6dda8 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/users/_components/view_users/UsersTableColumns.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/users/_components/view_users/UsersTableColumns.tsx @@ -130,10 +130,12 @@ export const getUsersTableColumns = ({ }, { id: "status", + accessorFn: (user) => (isScimInactive(user) ? "Inactive" : "Active"), meta: { title: "Status", skeleton: "badge" }, - header: "Status", + header: ({ column }) => , + sortDescFirst: false, size: 110, - enableSorting: false, + enableSorting: true, cell: ({ row }) => { if (isScimInactive(row.original)) { return ( @@ -165,9 +167,10 @@ export const getUsersTableColumns = ({ id: "user_alias", accessorKey: "user_alias", meta: { title: "User Alias" }, - header: "User Alias", + header: ({ column }) => , + sortDescFirst: false, size: 150, - enableSorting: false, + enableSorting: true, cell: ({ row }) => ( {row.original.user_alias || "-"} @@ -187,26 +190,28 @@ export const getUsersTableColumns = ({ id: "max_budget", accessorKey: "max_budget", meta: { title: "Budget (USD)", numeric: true }, - header: "Budget (USD)", + header: ({ column }) => , + sortDescFirst: true, size: 130, - enableSorting: false, + enableSorting: true, cell: ({ row }) => , }, { id: "sso_user_id", accessorKey: "sso_user_id", meta: { title: "SSO ID" }, - header: () => ( + header: ({ column }) => ( - SSO ID + } /> ), + sortDescFirst: false, size: 160, - enableSorting: false, + enableSorting: true, cell: ({ row }) => ( {row.original.sso_user_id ?? "-"} @@ -217,9 +222,10 @@ export const getUsersTableColumns = ({ id: "key_count", accessorKey: "key_count", meta: { title: "Virtual Keys", skeleton: "badge" }, - header: "Virtual Keys", + header: ({ column }) => , + sortDescFirst: true, size: 120, - enableSorting: false, + enableSorting: true, cell: ({ row }) => { const keyCount = row.original.key_count; if (keyCount > 0) { @@ -255,9 +261,10 @@ export const getUsersTableColumns = ({ id: "updated_at", accessorKey: "updated_at", meta: { title: "Updated At" }, - header: "Updated At", + header: ({ column }) => , + sortDescFirst: false, size: 130, - enableSorting: false, + enableSorting: true, cell: ({ row }) => , }, { From 683442eb61aa9d470cb01f5f7fb2c75fff9971be Mon Sep 17 00:00:00 2001 From: moe-berri Date: Fri, 23 Oct 2026 16:21:36 -0700 Subject: [PATCH 3/7] feat(ui): sort every teams table data column (#48109) * feat(ui): sort every teams table data column * fix(ui): order unlimited team limits as infinity and verify org scope * fix(ui): validate team sort filters without recursion --- .../management_endpoints/team_endpoints.py | 73 +++++- .../management_helpers/team_list_sorting.py | 109 ++++++++ .../management/test_team_list_sorting.py | 244 ++++++++++++++++++ .../test_team_list_sorting.py | 85 ++++++ ...st.tsx => TeamsTable.integration.test.tsx} | 22 +- .../components/TeamsPage/teamTableColumns.tsx | 55 +++- 6 files changed, 556 insertions(+), 32 deletions(-) create mode 100644 litellm/proxy/management_helpers/team_list_sorting.py create mode 100644 tests/integration/management/test_team_list_sorting.py create mode 100644 tests/unit/proxy/management_helpers/test_team_list_sorting.py rename ui/litellm-dashboard/src/components/TeamsPage/{TeamsTable.test.tsx => TeamsTable.integration.test.tsx} (92%) diff --git a/litellm/proxy/management_endpoints/team_endpoints.py b/litellm/proxy/management_endpoints/team_endpoints.py index d9763f72741..c7adae2ce8d 100644 --- a/litellm/proxy/management_endpoints/team_endpoints.py +++ b/litellm/proxy/management_endpoints/team_endpoints.py @@ -6109,13 +6109,52 @@ async def list_team_v2( "total_pages": 0, } - # Build order_by conditions - valid_sort_columns: Final = ["team_id", "team_alias", "created_at"] - order_by = None - if sort_by and sort_by in valid_sort_columns: - if sort_order.lower() not in ["asc", "desc"]: - sort_order = "asc" - order_by = {sort_by: sort_order.lower()} + from litellm.proxy.management_helpers.team_list_sorting import ( + COMPUTED_TEAM_SORTS, + UnsupportedTeamFilter, + sorted_team_ids, + ) + + valid_sort_columns: Final = ( + "team_id", + "team_alias", + "created_at", + "spend", + "max_budget", + "tpm_limit", + "rpm_limit", + "updated_at", + ) + sort_direction: Final = "desc" if str(sort_order).lower() == "desc" else "asc" + nullable_sort: Final = sort_by in ("max_budget", "tpm_limit", "rpm_limit") + order_by: Final = ( + ( + {sort_by: {"sort": sort_direction, "nulls": "first" if sort_direction == "desc" else "last"}}, + {"team_id": "asc"}, + ) + if nullable_sort + else ({sort_by: sort_direction}, {"team_id": "asc"}) + if sort_by in ("spend", "updated_at") + else {sort_by: sort_direction} + if sort_by in valid_sort_columns + else None + ) + computed_sort: Final = not use_deleted_table and sort_by in COMPUTED_TEAM_SORTS + sorted_ids: Final = ( + await sorted_team_ids( + prisma_client, + where=where_conditions, + sort_by=sort_by or "", + descending=sort_direction == "desc", + skip=skip, + take=page_size, + ) + if computed_sort + else () + ) + + if isinstance(sorted_ids, UnsupportedTeamFilter): + raise HTTPException(status_code=500, detail="Unsupported team sorting filter") # Get teams with pagination if use_deleted_table: @@ -6130,26 +6169,32 @@ async def list_team_v2( total_count = await _deleted_team_db(prisma_client).count(where=where_conditions) else: teams = await _team_db(prisma_client).find_many( - where=where_conditions, - skip=skip, + # bounded-ok: sorted_ids comes from LIMIT page_size, capped at 100 by the route + where={"AND": (where_conditions, {"team_id": {"in": sorted_ids}})} if computed_sort else where_conditions, + skip=0 if computed_sort else skip, take=page_size, - order=order_by if order_by else {"created_at": "desc"}, # Default sort + order=order_by if order_by else {"created_at": "desc"}, include=_INCLUDE_MODEL_TABLE, ) # Get total count for pagination total_count = await _team_db(prisma_client).count(where=where_conditions) + positions: Final = {team_id: index for index, team_id in enumerate(sorted_ids)} + sorted_teams: Final = sorted(teams, key=lambda team: positions[team.team_id]) if computed_sort else teams + # Calculate total pages total_pages: Final = -(-total_count // page_size) # Ceiling division # Aggregate virtual-key counts per team for the current page. The deleted # table does not carry keys_count, so it is skipped. - keys_count_by_team: dict[str, int] = {} - if not use_deleted_table: - keys_count_by_team = await _get_keys_count_by_team(prisma_client, teams) + keys_count_by_team: Final = ( + await _get_keys_count_by_team(prisma_client, sorted_teams) if not use_deleted_table else {} + ) # Convert Prisma models to response models with members_count and keys_count - team_list = _convert_teams_to_response_models(teams, use_deleted_table, keys_count_by_team=keys_count_by_team) + team_list: Final = _convert_teams_to_response_models( + sorted_teams, use_deleted_table, keys_count_by_team=keys_count_by_team + ) # Resolve resources inherited from access groups (single batch query) if not use_deleted_table: diff --git a/litellm/proxy/management_helpers/team_list_sorting.py b/litellm/proxy/management_helpers/team_list_sorting.py new file mode 100644 index 00000000000..caddd8afc08 --- /dev/null +++ b/litellm/proxy/management_helpers/team_list_sorting.py @@ -0,0 +1,109 @@ +from collections.abc import Mapping, Sequence +from dataclasses import dataclass +from itertools import chain +from types import MappingProxyType +from typing import Final + +from pydantic import BaseModel, ConfigDict, JsonValue, TypeAdapter + +from litellm.proxy.utils import PrismaClient + +COMPUTED_TEAM_SORTS: Final = MappingProxyType( + { + "organization_alias": "COALESCE((SELECT NULLIF(o.organization_alias, '') " + 'FROM "LiteLLM_OrganizationTable" o WHERE o.organization_id = t.organization_id), ' + "NULLIF(t.organization_id, ''))", + "members": "CASE WHEN jsonb_typeof(t.members_with_roles) = 'array' " + "THEN jsonb_array_length(t.members_with_roles) ELSE 0 END", + "models": "cardinality(t.models)", + "keys_count": '(SELECT COUNT(*) FROM "LiteLLM_VerificationToken" k WHERE k.team_id = t.team_id)', + } +) +_WHERE: Final = TypeAdapter(dict[str, JsonValue]) + + +@dataclass(frozen=True, slots=True) +class UnsupportedTeamFilter: + field: str + + +def _field_sql(field: str, value: JsonValue, index: int) -> tuple[str, tuple[object, ...]] | UnsupportedTeamFilter: + if field not in ("team_id", "team_alias", "organization_id"): + return UnsupportedTeamFilter(field) + column: Final = f't."{field}"' + match value: + case str(text): + return f"{column} = ${index}", (text,) + case {"in": list(values)}: + return f"{column} = ANY(${index}::text[])", (values,) + case {"contains": str(text), "mode": "insensitive"}: + return f"{column} ILIKE ${index}", (f"%{text}%",) + case {"startsWith": str(text)}: + return f"{column} LIKE ${index}", (f"{text}%",) + case _: + return UnsupportedTeamFilter(field) + + +def _group_tokens( + field: str, value: Sequence[JsonValue] | Mapping[str, JsonValue] +) -> tuple[str | tuple[str, JsonValue], ...]: + children: Final = tuple(value.items()) if isinstance(value, Mapping) else tuple(("AND", item) for item in value) + if not children: + return ("FALSE" if field == "OR" else "TRUE",) + return ( + "(", + *chain.from_iterable(((f" {field} " if index else ""), child) for index, child in enumerate(children)), + ")", + ) + + +def _where_sql(where: Mapping[str, JsonValue]) -> tuple[str, tuple[object, ...]] | UnsupportedTeamFilter: + pending: Final[list[str | tuple[str, JsonValue]]] = [("AND", dict(where))] # mutable-ok: iterative tree traversal + fragments: Final[list[str]] = [] # mutable-ok: append-only SQL output avoids quadratic tuple copying + params: Final[list[object]] = [] # mutable-ok: append-only bindings follow traversal order + while pending: + match pending.pop(): + case str(fragment): + fragments.append(fragment) + case (field, value) if field in ("AND", "OR") and isinstance(value, (list, dict)): + pending.extend(reversed(_group_tokens(field, value))) + case (field, value): + match _field_sql(field, value, len(params) + 1): + case UnsupportedTeamFilter() as failure: + return failure + case (fragment, values): + fragments.append(fragment) + params.extend(values) + return "".join(fragments), tuple(params) + + +class _TeamId(BaseModel): + model_config = ConfigDict(frozen=True) + team_id: str + + +_TEAM_IDS: Final = TypeAdapter(tuple[_TeamId, ...]) + + +async def sorted_team_ids( + prisma_client: PrismaClient, + *, + where: Mapping[str, object], + sort_by: str, + descending: bool, + skip: int, + take: int, +) -> Sequence[str] | UnsupportedTeamFilter: + result: Final = _where_sql(_WHERE.validate_python(where)) + if isinstance(result, UnsupportedTeamFilter): + return result + predicate, params = result + expression: Final = COMPUTED_TEAM_SORTS[sort_by] + direction: Final = "DESC" if descending else "ASC" + sql: Final = ( + f'SELECT t.team_id FROM "LiteLLM_TeamTable" t WHERE {predicate} ' + f"ORDER BY {expression} {direction} NULLS LAST, t.team_id ASC " + f"LIMIT ${len(params) + 1} OFFSET ${len(params) + 2}" + ) + rows: Final = _TEAM_IDS.validate_python(await prisma_client.db.query_raw(sql, *params, take, skip)) + return tuple(row.team_id for row in rows) diff --git a/tests/integration/management/test_team_list_sorting.py b/tests/integration/management/test_team_list_sorting.py new file mode 100644 index 00000000000..c09e51f0dfe --- /dev/null +++ b/tests/integration/management/test_team_list_sorting.py @@ -0,0 +1,244 @@ +import json +from itertools import chain +from typing import Final +from uuid import uuid4 + +import pytest + +from tests.integration._support.client import Gateway, Scenario +from tests.integration._support.database import write_rows + + +@pytest.mark.parametrize( + "field", ("spend", "max_budget", "tpm_limit", "rpm_limit", "updated_at", "members", "models", "keys_count") +) +@pytest.mark.parametrize("direction", ("asc", "desc")) +def test_team_list_sorts_across_pages_and_respects_membership( + field: str, + direction: str, + gateway: Gateway, + scenario: Scenario, +) -> None: + prefix: Final = f"sort-{uuid4().hex}" + owner: Final = scenario.user(user_role="internal_user") + caller: Final = scenario.key(user_id=owner) + values: Final = (20, 3, 3, 100, 0) + teams: Final = tuple( + scenario.team( + team_id=f"{prefix}-{index}", + team_alias=f"{prefix}-{index}", + members_with_roles=[{"user_id": owner, "role": "admin"}], + ) + for index in range(len(values)) + ) + outsider: Final = scenario.team(team_alias=f"{prefix}-outside") + for team_id, value in zip(teams, values, strict=True): + write_rows( + 'UPDATE "LiteLLM_TeamTable" SET spend = %s::float, max_budget = %s::float, ' + "tpm_limit = %s::bigint, rpm_limit = %s::bigint, " + "updated_at = timestamp '2030-01-01' + %s::int * interval '1 day', " + "members_with_roles = %s::jsonb, models = ARRAY(SELECT generate_series(1, %s::int)::text) " + "WHERE team_id = %s", + ( + str(value), + str(value), + str(value), + str(value), + str(value), + json.dumps( + [{"user_id": owner if i == 0 else f"{prefix}-member-{i}", "role": "admin"} for i in range(value)] + ), + str(value), + team_id, + ), + ) + write_rows( + 'INSERT INTO "LiteLLM_VerificationToken" (token, team_id) ' + "SELECT %s || '-key-' || n::text, %s FROM generate_series(1, %s::int) n", + (team_id, team_id, str(value)), + ) + write_rows( + 'UPDATE "LiteLLM_TeamTable" SET max_budget = NULL, tpm_limit = NULL, rpm_limit = NULL WHERE team_id = %s', + (teams[-1],), + ) + try: + expected: Final = tuple( + team_id + for value, team_id in sorted( + zip(values, teams, strict=True), + key=lambda item: ( + (item[0] != 0 if direction == "desc" else item[0] == 0) + if field in ("max_budget", "tpm_limit", "rpm_limit") + else False, + item[0] * (1 if direction == "asc" else -1), + item[1], + ), + ) + ) + pages: Final = tuple( + gateway.request( + "GET", + "/v2/team/list", + key=caller, + params={ + "search": prefix, + "user_id": owner, + "sort_by": field, + "sort_order": direction, + "page_size": "2", + "page": str(page), + }, + ) + for page in (1, 2, 3) + ) + assert all(page.status_code == 200 for page in pages), [page.text for page in pages] + returned: Final = tuple(row["team_id"] for row in chain.from_iterable(page.json()["teams"] for page in pages)) + assert returned == expected + assert outsider not in returned + assert all(page.json()["total"] == len(teams) for page in pages) + finally: + for team_id in teams: + write_rows('DELETE FROM "LiteLLM_VerificationToken" WHERE team_id = %s', (team_id,)) + + +@pytest.mark.parametrize("direction", ("asc", "desc")) +def test_team_list_sorts_organization_names_and_missing_values( + direction: str, gateway: Gateway, scenario: Scenario +) -> None: + prefix: Final = f"sort-org-{uuid4().hex}" + budget: Final = scenario.budget(max_budget=10) + organizations: Final = ((f"{prefix}-z", "alpha"), (f"{prefix}-a", "zulu"), (f"{prefix}-m", "")) + teams: Final = tuple(scenario.team(team_id=f"{prefix}-{i}", team_alias=f"{prefix}-{i}") for i in range(4)) + try: + for (organization_id, alias), team_id in zip(organizations, teams): + write_rows( + 'INSERT INTO "LiteLLM_OrganizationTable" ' + "(organization_id, organization_alias, budget_id, created_by, updated_by, updated_at) " + "VALUES (%s, %s, %s, 'sorting-test', 'sorting-test', NOW())", + (organization_id, alias, budget), + ) + write_rows( + 'UPDATE "LiteLLM_TeamTable" SET organization_id = %s WHERE team_id = %s', + (organization_id, team_id), + ) + ascending: Final = (teams[0], teams[2], teams[1]) + ordered: Final = ascending if direction == "asc" else tuple(reversed(ascending)) + pages: Final = tuple( + gateway.request( + "GET", + "/v2/team/list", + params={ + "search": prefix, + "sort_by": "organization_alias", + "sort_order": direction, + "page_size": "2", + "page": str(page), + }, + ) + for page in (1, 2) + ) + assert all(page.status_code == 200 for page in pages), [page.text for page in pages] + assert tuple(row["team_id"] for row in chain.from_iterable(page.json()["teams"] for page in pages)) == (*ordered, teams[-1]) + finally: + for team_id in teams: + write_rows('UPDATE "LiteLLM_TeamTable" SET organization_id = NULL WHERE team_id = %s', (team_id,)) + for organization_id, _ in organizations: + write_rows('DELETE FROM "LiteLLM_OrganizationTable" WHERE organization_id = %s', (organization_id,)) + + +@pytest.mark.parametrize("field", ("organization_alias", "members", "models", "keys_count")) +@pytest.mark.parametrize("direction", ("asc", "desc")) +def test_computed_team_sorts_preserve_org_and_membership_union( + field: str, direction: str, gateway: Gateway, scenario: Scenario +) -> None: + prefix: Final = f"sort-org-scope-{uuid4().hex}" + owner: Final = scenario.user(user_role="internal_user") + budget: Final = scenario.budget(max_budget=10) + orgs: Final = (f"{prefix}-allowed", f"{prefix}-outside") + teams: Final = tuple( + scenario.team( + team_id=f"{prefix}-{i}", + team_alias=f"{prefix}-{i}", + members_with_roles=[{"user_id": owner, "role": "admin"}] if i == 2 else [], + ) + for i in range(4) + ) + try: + for org, alias in zip(orgs, ("alpha", "zulu"), strict=True): + write_rows( + 'INSERT INTO "LiteLLM_OrganizationTable" ' + "(organization_id, organization_alias, budget_id, created_by, updated_by, updated_at) " + "VALUES (%s, %s, %s, 'sorting-test', 'sorting-test', NOW())", + (org, alias, budget), + ) + write_rows( + 'INSERT INTO "LiteLLM_OrganizationMembership" (user_id, organization_id, user_role) ' + "VALUES (%s, %s, 'org_admin')", + (owner, orgs[0]), + ) + for index, team in enumerate(teams): + write_rows( + 'UPDATE "LiteLLM_TeamTable" SET organization_id = %s, members_with_roles = %s::jsonb, ' + "models = ARRAY(SELECT generate_series(1, %s::int)::text) WHERE team_id = %s", + ( + orgs[0] if index < 2 else orgs[1], + json.dumps( + [ + {"user_id": owner if i == 0 else f"{prefix}-member-{i}", "role": "admin"} + for i in range(index + 1) + ] + ), + str(index + 1), + team, + ), + ) + write_rows( + 'INSERT INTO "LiteLLM_VerificationToken" (token, team_id) ' + "SELECT %s || '-key-' || n::text, %s FROM generate_series(1, %s::int) n", + (team, team, str(index + 1)), + ) + caller: Final = scenario.key(user_id=owner) + values: Final = (0, 0, 1) if field == "organization_alias" else (1, 2, 3) + expected: Final = tuple( + team + for value, team in sorted( + zip(values, teams[:3], strict=True), + key=lambda item: (item[0] * (-1 if direction == "desc" else 1), item[1]), + ) + ) + pages: Final = tuple( + gateway.request( + "GET", + "/v2/team/list", + key=caller, + params={ + "user_id": owner, + "search": prefix, + "sort_by": field, + "sort_order": direction, + "page_size": "2", + "page": str(page), + }, + ) + for page in (1, 2) + ) + assert all(page.status_code == 200 for page in pages), [page.text for page in pages] + assert all(page.json()["total"] == 3 for page in pages) + assert tuple(row["team_id"] for row in chain.from_iterable(page.json()["teams"] for page in pages)) == expected + denied: Final = gateway.request( + "GET", + "/v2/team/list", + key=caller, + params={ + "organization_id": orgs[1], + "sort_by": field, + }, + ) + assert denied.status_code == 403, denied.text + finally: + for team in teams: + write_rows('DELETE FROM "LiteLLM_VerificationToken" WHERE team_id = %s', (team,)) + write_rows('UPDATE "LiteLLM_TeamTable" SET organization_id = NULL WHERE team_id = %s', (team,)) + for org in orgs: + write_rows('DELETE FROM "LiteLLM_OrganizationMembership" WHERE organization_id = %s', (org,)) + write_rows('DELETE FROM "LiteLLM_OrganizationTable" WHERE organization_id = %s', (org,)) diff --git a/tests/unit/proxy/management_helpers/test_team_list_sorting.py b/tests/unit/proxy/management_helpers/test_team_list_sorting.py new file mode 100644 index 00000000000..bc4c77cc6a4 --- /dev/null +++ b/tests/unit/proxy/management_helpers/test_team_list_sorting.py @@ -0,0 +1,85 @@ +import json +import re +import sqlite3 +from typing import Final +from unittest.mock import AsyncMock, MagicMock + +import pytest +from pydantic import JsonValue + +from litellm.proxy.management_helpers.team_list_sorting import UnsupportedTeamFilter, _where_sql, sorted_team_ids + + +@pytest.mark.parametrize( + ("where", "expected"), + ( + ({}, ("a", "b", "c")), + ({"OR": []}, ()), + ({"AND": []}, ("a", "b", "c")), + ({"team_id": "a"}, ("a",)), + ({"team_alias": {"contains": "ALPHA", "mode": "insensitive"}}, ("a", "c")), + ({"team_id": {"startsWith": "b"}}, ("b",)), + ({"organization_id": {"in": ["org-1"]}}, ("a", "b")), + ( + { + "OR": [{"team_id": "b"}, {"team_alias": {"contains": "alpha", "mode": "insensitive"}}], + "AND": [{"OR": [{"organization_id": "org-1"}, {"team_id": "c"}]}], + "team_id": {"in": ["a", "c"]}, + }, + ("a", "c"), + ), + ({"team_id": "a' OR 1=1 --"}, ()), + ), +) +def test_team_filters_preserve_search_and_authorization_intersection( + where: dict[str, JsonValue], expected: tuple[str, ...] +) -> None: + result: Final = _where_sql(where) + assert not isinstance(result, UnsupportedTeamFilter) + predicate, params = result + sqlite_predicate: Final = re.sub( + r"= ANY\((\$\d+)::text\[\]\)", r"IN (SELECT value FROM json_each(\1))", predicate + ).replace(" ILIKE ", " LIKE ") + with sqlite3.connect(":memory:") as db: + db.execute("CREATE TABLE teams (team_id TEXT, team_alias TEXT, organization_id TEXT)") + db.executemany( + "INSERT INTO teams VALUES (?, ?, ?)", + (("a", "Alpha", "org-1"), ("b", "Beta", "org-1"), ("c", "Alpha", "org-2")), + ) + rows: Final = db.execute( + f"SELECT t.team_id FROM teams t WHERE {sqlite_predicate} ORDER BY t.team_id", + { + str(index): json.dumps(value) if isinstance(value, list) else value + for index, value in enumerate(params, 1) + }, + ).fetchall() + assert tuple(row[0] for row in rows) == expected + + +@pytest.mark.parametrize("where", ({"unknown": "x"}, {"team_id": {"not": "x"}}, {"AND": [False]})) +@pytest.mark.asyncio +async def test_unsupported_filter_fails_without_reading_unscoped_teams(where: dict[str, JsonValue]) -> None: + prisma: Final = MagicMock() + prisma.db.query_raw = AsyncMock() + result: Final = await sorted_team_ids(prisma, where=where, sort_by="keys_count", descending=False, skip=0, take=25) + assert isinstance(result, UnsupportedTeamFilter) + prisma.db.query_raw.assert_not_awaited() + + +@pytest.mark.parametrize("sort_by", ("keys_count", "members", "models", "organization_alias")) +@pytest.mark.parametrize("descending", (False, True)) +@pytest.mark.asyncio +async def test_sorted_ids_bind_search_and_pagination_without_interpolating_input( + sort_by: str, descending: bool +) -> None: + prisma: Final = MagicMock() + prisma.db.query_raw = AsyncMock(return_value=[]) + search: Final = "x' OR TRUE --" + result: Final = await sorted_team_ids( + prisma, where={"team_id": search}, sort_by=sort_by, descending=descending, skip=25, take=25 + ) + assert result == () + sql, *bindings = prisma.db.query_raw.await_args.args + assert search not in sql + assert bindings == [search, 25, 25] + assert ("DESC" if descending else "ASC") + " NULLS LAST, t.team_id ASC" in sql diff --git a/ui/litellm-dashboard/src/components/TeamsPage/TeamsTable.test.tsx b/ui/litellm-dashboard/src/components/TeamsPage/TeamsTable.integration.test.tsx similarity index 92% rename from ui/litellm-dashboard/src/components/TeamsPage/TeamsTable.test.tsx rename to ui/litellm-dashboard/src/components/TeamsPage/TeamsTable.integration.test.tsx index ad89b62d611..def6750071f 100644 --- a/ui/litellm-dashboard/src/components/TeamsPage/TeamsTable.test.tsx +++ b/ui/litellm-dashboard/src/components/TeamsPage/TeamsTable.integration.test.tsx @@ -163,16 +163,26 @@ describe("sort contract – only backend-sortable columns are sortable", () => { it("right-aligns Spend / Budget but not Team", () => { renderTable(); - expect(screen.getByRole("columnheader", { name: "Spend / Budget" })).toHaveClass("text-right"); + expect(screen.getByRole("columnheader", { name: /Spend\s*\/\s*Budget/ })).toHaveClass("text-right"); expect(screen.getByRole("columnheader", { name: "Team" })).not.toHaveClass("text-right"); }); - it("does not make Spend / Budget sortable (the backend rejects sort_by=spend)", () => { + it.each([ + ["Spend or Budget", "Spend", "spend"], + ["Spend or Budget", "Budget", "max_budget"], + ["Members or Models or Keys", "Members", "members"], + ["Members or Models or Keys", "Models", "models"], + ["Members or Models or Keys", "Keys", "keys_count"], + ])("sorts %s by %s through the API", async (menu, label, field) => { + const user = userEvent.setup(); renderTable(); - expect(screen.queryByText("Spend / Budget").closest("button")).toBeNull(); - // Team and Created are the only sortable headers. - expect(screen.getByText("Team").closest("button")).not.toBeNull(); - expect(screen.getByText("Created").closest("button")).not.toBeNull(); + await user.click(screen.getByRole("button", { name: `Sort options for ${menu}` })); + await user.click(await screen.findByRole("menuitem", { name: `${label} descending` })); + await waitFor(() => expect(lastOptions()).toMatchObject({ sortBy: field, sortOrder: "desc" })); + expect(mockUseTeamsTable.mock.calls.at(-1)?.[0]).toBe(1); + await user.click(screen.getByRole("button", { name: `Sort options for ${menu}` })); + await user.click(await screen.findByRole("menuitem", { name: `${label} ascending` })); + await waitFor(() => expect(lastOptions()).toMatchObject({ sortBy: field, sortOrder: "asc" })); }); }); diff --git a/ui/litellm-dashboard/src/components/TeamsPage/teamTableColumns.tsx b/ui/litellm-dashboard/src/components/TeamsPage/teamTableColumns.tsx index 05bfccf6417..e56e0d16522 100644 --- a/ui/litellm-dashboard/src/components/TeamsPage/teamTableColumns.tsx +++ b/ui/litellm-dashboard/src/components/TeamsPage/teamTableColumns.tsx @@ -3,7 +3,7 @@ import { ColumnDef } from "@tanstack/react-table"; import { Copy, KeyRound, Layers, MoreHorizontal, Pencil, Trash2, Users } from "lucide-react"; -import { DataTableSortHeader } from "@/components/shared/DataTable"; +import { DataTableMultiSortHeader, DataTableSortHeader } from "@/components/shared/DataTable"; import { DateCell, IdentityCell, SpendBudgetCell } from "@/components/shared/table_cells"; import { buttonVariants } from "@/components/ui/button"; import { @@ -174,9 +174,10 @@ export const getTeamTableColumns = ({ id: "organization_alias", accessorKey: "organization_id", meta: { title: "Organization" }, - header: "Organization", + header: ({ column }) => , + sortDescFirst: false, size: 160, - enableSorting: false, + enableSorting: true, cell: (info) => { const orgId = info.getValue() as string | null; if (!orgId) return —; @@ -202,7 +203,16 @@ export const getTeamTableColumns = ({ ), }, - header: "Resources", + header: ({ table }) => ( + + ), size: 210, enableSorting: false, cell: ({ row }) => , @@ -211,7 +221,15 @@ export const getTeamTableColumns = ({ id: "spend", accessorKey: "spend", meta: { title: "Spend / Budget", skeleton: "meter", numeric: true }, - header: "Spend / Budget", + header: ({ table }) => ( + + ), size: 200, enableSorting: false, cell: ({ row }) => ( @@ -234,24 +252,36 @@ export const getTeamTableColumns = ({ }, { id: "members", + accessorFn: teamMemberCount, meta: { title: "Members", numeric: true }, - header: "Members", + header: ({ column }) => , + sortDescFirst: true, size: 110, - enableSorting: false, + enableSorting: true, cell: ({ row }) => {teamMemberCount(row.original)}, }, { id: "models", + accessorFn: teamModelCount, meta: { title: "Models", numeric: true }, - header: "Models", + header: ({ column }) => , + sortDescFirst: true, size: 100, - enableSorting: false, + enableSorting: true, cell: ({ row }) => {teamModelCount(row.original)}, }, { id: "rate_limits", meta: { title: "Rate Limits", skeleton: "twoLine" }, - header: "Rate Limits", + header: ({ table }) => ( + + ), size: 140, enableSorting: false, cell: ({ row }) => ( @@ -265,9 +295,10 @@ export const getTeamTableColumns = ({ id: "updated_at", accessorKey: "updated_at", meta: { title: "Updated" }, - header: "Updated", + header: ({ column }) => , + sortDescFirst: false, size: 130, - enableSorting: false, + enableSorting: true, cell: (info) => , }, { From bffc166cd66b1661d0fc25e8f9783c370f1dda21 Mon Sep 17 00:00:00 2001 From: moe-berri Date: Fri, 23 Oct 2026 16:21:53 -0700 Subject: [PATCH 4/7] feat(ui): sort budgets by reset duration (#48107) * feat(ui): sort budgets by reset duration * fix(ui): order unlimited budget limits as infinity * fix(ui): sort localized budget reset durations --- .../management_v1/budgets.py | 46 +++++++-- .../management/test_budget_list_sorting.py | 94 +++++++++++++++++++ .../management_v1/test_budgets.py | 20 +--- .../budgets/_components/BudgetTable.test.tsx | 9 +- .../_components/BudgetTableColumns.tsx | 3 +- ui/litellm-dashboard/src/lib/http/schema.d.ts | 5 +- 6 files changed, 149 insertions(+), 28 deletions(-) create mode 100644 tests/integration/management/test_budget_list_sorting.py diff --git a/litellm/proxy/management_endpoints/management_v1/budgets.py b/litellm/proxy/management_endpoints/management_v1/budgets.py index 1970a0797ef..19a5408496a 100644 --- a/litellm/proxy/management_endpoints/management_v1/budgets.py +++ b/litellm/proxy/management_endpoints/management_v1/budgets.py @@ -3,6 +3,7 @@ from collections.abc import Mapping, Sequence from dataclasses import dataclass from datetime import datetime +from itertools import chain from types import MappingProxyType from typing import Annotated, Final @@ -76,6 +77,32 @@ _ROW_COUNTS: Final = TypeAdapter(tuple[_RowCount, ...]) SELECTED_COLUMNS: Final = ", ".join(f'"{name}"' for name in BudgetListItem.model_fields) +# Unicode 17.0 Nd zeroes, each followed by its other nine digits in UnicodeData.txt. +_DECIMAL_ZEROES: Final = "0٠۰߀०০੦૦୦௦౦೦൦෦๐໐༠၀႐០᠐᥆᧐᪀᪐᭐᮰᱀᱐꘠꣐꤀꧐꧰꩐꯰０𐒠𐴰𐵀𑁦𑃰𑄶𑇐𑋰𑑐𑓐𑙐𑛀𑛐𑛚𑜰𑣠𑥐𑯰𑱐𑵐𑶠\U00011de0𑽐𖄰𖩠𖫀𖭐𖵰𜳰𝟎𝟘𝟢𝟬𝟶𞅀𞋰𞓰𞗱𞥐🯰" +_DECIMAL_DIGITS: Final = "".join( + chain.from_iterable(tuple(chr(ord(zero) + digit) for digit in range(10)) for zero in _DECIMAL_ZEROES) +) +_ASCII_DIGITS: Final = "0123456789" * len(_DECIMAL_ZEROES) + + +_BUDGET_DURATION_SECONDS: Final = f"""( + SELECT EXTRACT(EPOCH FROM ( + (CURRENT_TIMESTAMP AT TIME ZONE 'UTC') + + (parts[1] || ' ' || CASE parts[2] + WHEN 's' THEN 'seconds' WHEN 'm' THEN 'minutes' WHEN 'h' THEN 'hours' + WHEN 'd' THEN 'days' WHEN 'w' THEN 'weeks' WHEN 'mo' THEN 'months' END)::interval + - (CURRENT_TIMESTAMP AT TIME ZONE 'UTC') + )) FROM regexp_match( + CASE lower(trim(budget_duration)) + WHEN 'hourly' THEN '1h' WHEN 'daily' THEN '24h' + WHEN 'weekly' THEN '7d' WHEN 'monthly' THEN '30d' + ELSE CASE WHEN octet_length(budget_duration) = length(budget_duration) THEN budget_duration + ELSE translate(budget_duration, '{_DECIMAL_DIGITS}', '{_ASCII_DIGITS}') END END, + '^([0-9]+)(mo|[smhdw])' + ) AS duration(parts) +)""" + + @dataclass(frozen=True, slots=True) class PrismaBudgetListExecutor: """The database half of the budgets list. Every caller-supplied value is bound to a @@ -91,10 +118,16 @@ class PrismaBudgetListExecutor: async def find_many(self, plan: QueryPlan) -> Sequence[BudgetListItem]: clauses, params = where_sql(plan.where) + ordering: Final = ", ".join( + order_by_sql((key,)).replace("NULLS LAST", "NULLS FIRST") + if key.descending and key.field in ("max_budget", "tpm_limit", "rpm_limit", "tpd_limit") + else order_by_sql((key,)) + for key in plan.order + ).replace('"budget_duration"', _BUDGET_DURATION_SECONDS) sql: Final = ( f"SELECT {SELECTED_COLUMNS} FROM {BUDGET_TABLE}" + (f" WHERE {clauses}" if clauses else "") - + f" ORDER BY {order_by_sql(plan.order)}" + + f" ORDER BY {ordering}" + f" LIMIT ${len(params) + 1} OFFSET ${len(params) + 2}" ) return _BUDGET_ROWS.validate_python(await self.prisma_client.db.query_raw(sql, *params, plan.take, plan.skip)) @@ -111,8 +144,6 @@ def _scope(caller: UserAPIKeyAuth) -> Scope: return ScopeDenied(reason=f"Only proxy admins can list budgets, your role={caller.user_role}") -# budget_duration is deliberately absent from `sortable`: the column holds strings -# like "7d" and "30d", so a lexicographic ORDER BY puts "30d" ahead of "7d". BUDGET_FILTERS: Final[Mapping[str, FilterSpec]] = MappingProxyType( { "budget_duration": FilterSpec(type=str, ops=frozenset(("in", "is_null"))), @@ -123,7 +154,9 @@ BUDGET_FILTERS: Final[Mapping[str, FilterSpec]] = MappingProxyType( BUDGETS_LIST_SPEC: Final[ListSpec[BudgetListItem, BudgetListItem]] = ListSpec( resource="budgets", - sortable=frozenset(("budget_id", "max_budget", "tpm_limit", "rpm_limit", "tpd_limit", "created_at")), + sortable=frozenset( + ("budget_id", "max_budget", "tpm_limit", "rpm_limit", "tpd_limit", "created_at", "budget_duration") + ), searchable=frozenset(("budget_id",)), filters=BUDGET_FILTERS, default_sort=(SortKey(field="created_at", descending=True),), @@ -154,8 +187,9 @@ async def list_budgets( way to page, sort or filter it. `sort` takes a comma-separated list of `budget_id`, `max_budget`, `tpm_limit`, - `rpm_limit`, `tpd_limit` or `created_at`, each optionally prefixed with `-` for descending, - and defaults to `-created_at`. `budget_id` is appended to every sort as the + `rpm_limit`, `tpd_limit`, `created_at` or `budget_duration`, each optionally prefixed with `-` for descending, + and defaults to `-created_at`. Reset periods sort by elapsed time, with calendar months + measured from the current UTC date and unset periods last. `budget_id` is appended to every sort as the tiebreaker. `q` is a case-insensitive substring match on `budget_id`. `page_size` defaults to 50 and is capped at 100. Filters are `filter[budget_duration][in|is_null]`, `filter[max_budget][gte|lte|is_null]` diff --git a/tests/integration/management/test_budget_list_sorting.py b/tests/integration/management/test_budget_list_sorting.py new file mode 100644 index 00000000000..aae0bd6e274 --- /dev/null +++ b/tests/integration/management/test_budget_list_sorting.py @@ -0,0 +1,94 @@ +from datetime import datetime, timezone +from itertools import chain +from typing import Final +from uuid import uuid4 + +import pytest + +from litellm.litellm_core_utils.duration_parser import duration_in_seconds_from +from tests.integration._support.client import Gateway, Scenario + + +@pytest.mark.parametrize("descending", (False, True)) +def test_budget_reset_sort_orders_durations_before_pagination( + descending: bool, + gateway: Gateway, + scenario: Scenario, +) -> None: + prefix: Final = f"sort-{uuid4().hex}" + durations: Final = ("30d", "7d", "1h", "60m", "1mo", "weekly", "monthly", "٥d", "۵d", "५d", "５d", "𝟝d", "5d", None) + ids: Final = tuple( + scenario.budget(budget_id=f"{prefix}-{index}", budget_duration=duration, max_budget=10) + for index, duration in enumerate(durations) + ) + now: Final = datetime.now(timezone.utc) + expected: Final = tuple( + budget_id + for _, budget_id in sorted( + zip(durations, ids, strict=True), + key=lambda item: ( + item[0] is None, + duration_in_seconds_from(now, item[0]) * (-1 if descending else 1) if item[0] else 0, + item[1], + ), + ) + ) + pages: Final = tuple( + gateway.request( + "GET", + "/management/v1/budgets", + params={ + "q": prefix, + "sort": "-budget_duration" if descending else "budget_duration", + "page_size": "3", + "page": str(page), + }, + ) + for page in range(1, -(-len(ids) // 3) + 1) + ) + assert all(page.status_code == 200 for page in pages), [page.text for page in pages] + assert tuple(row["budget_id"] for row in chain.from_iterable(page.json()["data"] for page in pages)) == expected + filtered: Final = gateway.request( + "GET", + "/management/v1/budgets", + params={ + "q": prefix, + "sort": "budget_duration", + "filter[budget_duration][in]": "7d,30d", + }, + ) + assert filtered.status_code == 200, filtered.text + assert [row["budget_duration"] for row in filtered.json()["data"]] == ["7d", "30d"] + + +@pytest.mark.parametrize("field", ("max_budget", "tpm_limit", "rpm_limit", "tpd_limit")) +@pytest.mark.parametrize("descending", (False, True)) +def test_budget_limits_sort_unlimited_as_infinity( + field: str, descending: bool, gateway: Gateway, scenario: Scenario +) -> None: + prefix: Final = f"sort-unlimited-{uuid4().hex}" + values: Final = (20, None, 0, 3) + ids: Final = tuple(scenario.budget(budget_id=f"{prefix}-{i}", **{field: value}) for i, value in enumerate(values)) + expected: Final = tuple( + identity + for value, identity in sorted( + zip(values, ids, strict=True), + key=lambda item: float("inf") if item[0] is None else item[0], + reverse=descending, + ) + ) + pages: Final = tuple( + gateway.request( + "GET", + "/management/v1/budgets", + params={ + "q": prefix, + "sort": ("-" if descending else "") + field, + "page_size": "2", + "page": str(page), + }, + ) + for page in (1, 2) + ) + assert all(page.status_code == 200 for page in pages), [page.text for page in pages] + assert tuple(row["budget_id"] for row in chain.from_iterable(page.json()["data"] for page in pages)) == expected diff --git a/tests/unit/proxy/management_endpoints/management_v1/test_budgets.py b/tests/unit/proxy/management_endpoints/management_v1/test_budgets.py index 2b438a9d370..4b42e6404a3 100644 --- a/tests/unit/proxy/management_endpoints/management_v1/test_budgets.py +++ b/tests/unit/proxy/management_endpoints/management_v1/test_budgets.py @@ -52,7 +52,7 @@ app.include_router(router) client = TestClient(app) BUDGETS_PATH = f"{MANAGEMENT_V1_PREFIX}/budgets" -SORTABLE = ["budget_id", "created_at", "max_budget", "rpm_limit", "tpd_limit", "tpm_limit"] +SORTABLE = ["budget_duration", "budget_id", "created_at", "max_budget", "rpm_limit", "tpd_limit", "tpm_limit"] def _row(budget_id: str, **overrides: Any) -> dict[str, Any]: @@ -172,22 +172,12 @@ def test_appends_the_tiebreaker_to_an_explicit_sort(query_raw, as_proxy_admin): _get("sort=-max_budget") - assert 'ORDER BY "max_budget" DESC NULLS LAST, "budget_id" ASC NULLS LAST' in _select_call(query_raw)[0] + assert 'ORDER BY "max_budget" DESC NULLS FIRST, "budget_id" ASC NULLS LAST' in _select_call(query_raw)[0] -def test_refuses_to_sort_on_budget_duration(query_raw, as_proxy_admin): - """The column holds "7d"/"30d", so a lexicographic ORDER BY would put "30d" - before "7d" and silently mis-order the page.""" - _serve(query_raw, []) - +def test_accepts_duration_sorting(query_raw, as_proxy_admin): response = _get("sort=budget_duration") - - assert response.status_code == 400 - assert response.headers["content-type"].startswith("application/problem+json") - body = response.json() - assert "budget_duration" in body["detail"] - assert body["allowed"] == SORTABLE - query_raw.assert_not_called() + assert response.status_code == 200 def test_the_advertised_sort_fields_are_the_ones_that_work(query_raw, as_proxy_admin): @@ -346,7 +336,7 @@ def test_q_does_not_search_any_other_column(query_raw, as_proxy_admin): def test_is_null_selects_the_unlimited_budgets(query_raw, as_proxy_admin): - """"Unlimited" is max_budget IS NULL; `max_budget = 0` would be a hard zero cap.""" + """ "Unlimited" is max_budget IS NULL; `max_budget = 0` would be a hard zero cap.""" _serve(query_raw, [_row("b-unlimited", max_budget=None)]) body = _get("filter[max_budget][is_null]=true").json() diff --git a/ui/litellm-dashboard/src/app/(dashboard)/budgets/_components/BudgetTable.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/budgets/_components/BudgetTable.test.tsx index ae898645de4..b4fb7586e6d 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/budgets/_components/BudgetTable.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/budgets/_components/BudgetTable.test.tsx @@ -134,13 +134,16 @@ describe("BudgetTable", () => { } }); - it("should not make the reset column sortable", async () => { + it("should request duration sorting from the server", async () => { const user = userEvent.setup(); - renderWithProviders(); + const onSortingChange = vi.fn(); + renderWithProviders(); await showColumn(user, "budget_duration"); const headers = screen.getAllByRole("columnheader").map((header) => header.textContent); expect(headers).toContain("Reset"); - expect(screen.queryByTestId("sort-header-budget_duration")).not.toBeInTheDocument(); + await user.click(screen.getByTestId("sort-header-budget_duration")); + expect(onSortingChange).toHaveBeenCalled(); + expect(onSortingChange.mock.calls[0][0]([])).toEqual([{ id: "budget_duration", desc: false }]); }); it("should ask the list for a new sort when a sortable header is clicked", async () => { diff --git a/ui/litellm-dashboard/src/app/(dashboard)/budgets/_components/BudgetTableColumns.tsx b/ui/litellm-dashboard/src/app/(dashboard)/budgets/_components/BudgetTableColumns.tsx index fb894322208..3ec291c1158 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/budgets/_components/BudgetTableColumns.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/budgets/_components/BudgetTableColumns.tsx @@ -139,8 +139,7 @@ export const getBudgetTableColumns = ({ accessorKey: "budget_duration", filterFn: serverFilter, meta: { title: "Reset" }, - // "7d"/"30d" sort lexicographically, not chronologically, so the route does not offer it. - enableSorting: false, + sortDescFirst: false, header: ({ column }) => , size: 110, cell: ({ row }) => , diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index a71933626fc..1d36f5fd9a2 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -9280,8 +9280,9 @@ export interface paths { * way to page, sort or filter it. * * `sort` takes a comma-separated list of `budget_id`, `max_budget`, `tpm_limit`, - * `rpm_limit`, `tpd_limit` or `created_at`, each optionally prefixed with `-` for descending, - * and defaults to `-created_at`. `budget_id` is appended to every sort as the + * `rpm_limit`, `tpd_limit`, `created_at` or `budget_duration`, each optionally prefixed with `-` for descending, + * and defaults to `-created_at`. Reset periods sort by elapsed time, with calendar months + * measured from the current UTC date and unset periods last. `budget_id` is appended to every sort as the * tiebreaker. `q` is a case-insensitive substring match on `budget_id`. * `page_size` defaults to 50 and is capped at 100. Filters are * `filter[budget_duration][in|is_null]`, `filter[max_budget][gte|lte|is_null]` From acfa7a29addc4cc44f7f26d9e6b4309872737aac Mon Sep 17 00:00:00 2001 From: Igor Tavares Date: Fri, 23 Oct 2026 20:27:33 -0300 Subject: [PATCH 5/7] fix(proxy): record spend_logs_metadata on /interactions (#47848) The router plants its deployment fields in litellm_metadata for Interactions calls, while the proxy put the request's metadata (spend_logs_metadata, tags, key fields) under metadata. The spend log resolves litellm_metadata and only carries user_api_key* fields over, so spend_logs_metadata never reached LiteLLM_SpendLogs. Route /interactions like /responses so both land in litellm_metadata. Fixes #47841 Co-authored-by: Igor Tavares <2163472+igorborgest@users.noreply.github.com> Co-authored-by: Claude Opus 5.5 --- litellm/proxy/litellm_pre_call_utils.py | 1 + .../unit/proxy/test_litellm_pre_call_utils.py | 47 +++++++++++++++++++ 2 files changed, 48 insertions(+) diff --git a/litellm/proxy/litellm_pre_call_utils.py b/litellm/proxy/litellm_pre_call_utils.py index 8a45f658f36..da849ee4fad 100644 --- a/litellm/proxy/litellm_pre_call_utils.py +++ b/litellm/proxy/litellm_pre_call_utils.py @@ -334,6 +334,7 @@ LITELLM_METADATA_ROUTES: Final = ( "/v1/messages", "responses", "files", + "interactions", ) LITELLM_TRACE_CONTROL_METADATA_FIELDS: Final = frozenset( diff --git a/tests/unit/proxy/test_litellm_pre_call_utils.py b/tests/unit/proxy/test_litellm_pre_call_utils.py index d9def23bde8..04f6381f32f 100644 --- a/tests/unit/proxy/test_litellm_pre_call_utils.py +++ b/tests/unit/proxy/test_litellm_pre_call_utils.py @@ -106,6 +106,13 @@ class TestGetMetadataVariableName: request = self._make_request("/v1/files") assert get_metadata_variable_name(request) == "litellm_metadata" + @pytest.mark.parametrize( + "path", ["/v1beta/interactions", "/interactions", "/v1beta/interactions/v1_abc123", "/interactions/v1_abc123/cancel"] + ) + def test_returns_litellm_metadata_for_interactions_routes(self, path): + request = self._make_request(path) + assert get_metadata_variable_name(request) == "litellm_metadata" + def test_returns_metadata_for_chat_completions(self): request = self._make_request("/chat/completions") assert get_metadata_variable_name(request) == "metadata" @@ -708,6 +715,46 @@ async def test_add_litellm_data_to_request_parses_json_object_string_litellm_met assert updated["litellm_metadata"]["cost_centre"] == "research" +@pytest.mark.asyncio +async def test_interactions_spend_logs_metadata_survives_the_router_metadata_bucket(): + """GH#47841: on /v1beta/interactions the router plants its deployment fields in + `litellm_metadata`. The proxy must put the request's metadata there too: otherwise the + spend log reads only the router's bucket and `spend_logs_metadata` from the + `x-litellm-spend-logs-metadata` header never reaches LiteLLM_SpendLogs.""" + request_mock = MagicMock(spec=Request) + request_mock.url = MagicMock() + request_mock.url.__str__.return_value = "http://localhost/v1beta/interactions" + request_mock.url.path = "/v1beta/interactions" + request_mock.method = "POST" + request_mock.query_params = {} + request_mock.headers = { + "Content-Type": "application/json", + "x-litellm-spend-logs-metadata": '{"cost_centre": "research"}', + } + request_mock.client = MagicMock() + request_mock.client.host = "127.0.0.1" + request_mock.state.parent_otel_span = None + + data = await add_litellm_data_to_request( + data={"model": "gemini-3.1-flash-lite", "input": "Reply OK.", "store": True}, + request=request_mock, + user_api_key_dict=UserAPIKeyAuth(api_key="hashed-key"), + proxy_config=MagicMock(), + general_settings={}, + version="test-version", + ) + assert data["litellm_metadata"]["spend_logs_metadata"] == {"cost_centre": "research"} + + # The router then adds its deployment fields to the same bucket, and the spend log + # resolves that bucket. + data["litellm_metadata"].update({"model_group": "gemini-3.1-flash-lite", "deployment": "gemini/gemini-3.1-flash-lite"}) + resolved = get_litellm_metadata_from_kwargs( + {"litellm_params": {"metadata": data.get("metadata") or {}, "litellm_metadata": data["litellm_metadata"]}} + ) + assert resolved["spend_logs_metadata"] == {"cost_centre": "research"} + assert resolved["model_group"] == "gemini-3.1-flash-lite" + + @pytest.mark.asyncio async def test_add_litellm_data_to_request_proxy_server_request_body_is_post_strip(): """Regression: proxy_server_request['body'] used to be snapshotted before From 23ccff69a105bf70499c6003393e3390d6985aeb Mon Sep 17 00:00:00 2001 From: shrey-berri Date: Fri, 23 Oct 2026 16:34:55 -0700 Subject: [PATCH 6/7] fix(azure_ai): forward the thinking-binding-controls beta header to Foundry (#48033) --- .../src/providers/anthropic/beta.rs | 6 +- litellm/anthropic_beta_headers_config.json | 2 +- litellm/proxy/litellm_pre_call_utils.py | 5 +- .../unit/proxy/test_litellm_pre_call_utils.py | 61 +++++++++++++++++++ .../test_anthropic_beta_headers_filtering.py | 4 +- 5 files changed, 72 insertions(+), 6 deletions(-) diff --git a/litellm-rust/crates/llms-types/src/providers/anthropic/beta.rs b/litellm-rust/crates/llms-types/src/providers/anthropic/beta.rs index b1509313b88..8429c98a88a 100644 --- a/litellm-rust/crates/llms-types/src/providers/anthropic/beta.rs +++ b/litellm-rust/crates/llms-types/src/providers/anthropic/beta.rs @@ -235,7 +235,10 @@ impl AnthropicBeta { BetaProvider::Databricks, ], Self::Compact20260112 => &[BetaProvider::AzureAi, BetaProvider::BedrockConverse], - Self::ComputerUse20250124 | Self::ComputerUse20251124 | Self::Context1m20250807 => &[], + Self::ComputerUse20250124 + | Self::ComputerUse20251124 + | Self::Context1m20250807 + | Self::ThinkingBindingControls20260801 => &[], Self::DangerousToolUse20260903 => &[BetaProvider::BedrockConverse], Self::Effort20251124 => &[BetaProvider::VertexAi], Self::FastMode20260201 | Self::Oauth20250420 => &[ @@ -272,7 +275,6 @@ impl AnthropicBeta { BetaProvider::Databricks, ], Self::StructuredOutputs20251113 => &[BetaProvider::Bedrock, BetaProvider::VertexAi], - Self::ThinkingBindingControls20260801 => &[BetaProvider::AzureAi], Self::ThinkingDisplayUpdates20260818 => &[BetaProvider::Databricks], Self::ToolSearchTool20251019 => &[ BetaProvider::Anthropic, diff --git a/litellm/anthropic_beta_headers_config.json b/litellm/anthropic_beta_headers_config.json index 236b18fe428..d68de6c9d4d 100644 --- a/litellm/anthropic_beta_headers_config.json +++ b/litellm/anthropic_beta_headers_config.json @@ -67,7 +67,7 @@ "structured-outputs-2025-11-13": "structured-outputs-2025-11-13", "text_editor_20241022": null, "text_editor_20250124": null, - "thinking-binding-controls-2026-08-01": null, + "thinking-binding-controls-2026-08-01": "thinking-binding-controls-2026-08-01", "thinking-display-updates-2026-08-18": "thinking-display-updates-2026-08-18", "token-efficient-tools-2025-02-19": null, "web-fetch-2025-09-10": "web-fetch-2025-09-10", diff --git a/litellm/proxy/litellm_pre_call_utils.py b/litellm/proxy/litellm_pre_call_utils.py index da849ee4fad..8dc7501bc3b 100644 --- a/litellm/proxy/litellm_pre_call_utils.py +++ b/litellm/proxy/litellm_pre_call_utils.py @@ -4278,6 +4278,9 @@ _ANTHROPIC_API_HEADER_PROVIDERS: Final = ",".join( LlmProviders.VERTEX_AI.value, ) ) +_NATIVE_MESSAGES_ANTHROPIC_API_HEADER_PROVIDERS: Final = ",".join( + (LlmProviders.AZURE_AI.value, LlmProviders.DATABRICKS.value) +) _ANTHROPIC_OAUTH_CREDENTIAL_PROVIDERS: Final = ",".join(sorted(ANTHROPIC_OAUTH_FORWARD_PROVIDERS)) @@ -4306,7 +4309,7 @@ def add_provider_specific_headers_to_request( custom_llm_provider=_ANTHROPIC_API_HEADER_PROVIDERS, extra_headers=anthropic_api_headers ), ProviderSpecificHeader( - custom_llm_provider=LlmProviders.DATABRICKS.value, + custom_llm_provider=_NATIVE_MESSAGES_ANTHROPIC_API_HEADER_PROVIDERS, extra_headers=anthropic_api_headers, native_anthropic_messages_only=True, ), diff --git a/tests/unit/proxy/test_litellm_pre_call_utils.py b/tests/unit/proxy/test_litellm_pre_call_utils.py index 04f6381f32f..64ff8392c02 100644 --- a/tests/unit/proxy/test_litellm_pre_call_utils.py +++ b/tests/unit/proxy/test_litellm_pre_call_utils.py @@ -8111,6 +8111,67 @@ def test_client_anthropic_api_headers_stay_off_openai_compatible_providers(): assert forwarded == {} +def test_client_anthropic_api_headers_stay_off_azure_ai_outside_native_messages(): + forwarded = _headers_forwarded_to({"anthropic-beta": "thinking-binding-controls-2026-08-01"}, "azure_ai") + + assert forwarded == {} + + +@pytest.mark.asyncio +@pytest.mark.usefixtures("local_beta_headers_config") +async def test_client_thinking_binding_beta_reaches_azure_ai_foundry_claude(): + import httpx + + from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler + + data: Final[dict] = {"model": "claude-opus-4-8"} # mutable-ok: the function under test writes into it + add_provider_specific_headers_to_request( + data=data, + headers={ + "content-type": "application/json", + "anthropic-beta": "thinking-binding-controls-2026-08-01", + "anthropic-version": "2026-08-01", + }, + ) + received: Final[list[httpx.Request]] = [] # mutable-ok: the fake Foundry endpoint records what the proxy sent + + def foundry_answers(request: httpx.Request) -> httpx.Response: + received.append(request) + return httpx.Response( + status_code=200, + json={ + "id": "msg_01", + "type": "message", + "role": "assistant", + "model": "claude-opus-4-8", + "content": [{"type": "text", "text": "hello"}], + "stop_reason": "end_turn", + "stop_sequence": None, + "usage": {"input_tokens": 12, "output_tokens": 6}, + }, + request=request, + ) + + upstream: Final = AsyncHTTPHandler(transport=httpx.MockTransport(foundry_answers)) + try: + await litellm.anthropic.messages.acreate( + max_tokens=16, + messages=[{"role": "user", "content": "hi"}], + thinking={"type": "adaptive", "block_binding": {"prefix_mismatch_behavior": "drop_block"}}, + model="azure_ai/claude-opus-4-8", + api_key="foundry-key", + api_base="https://my-resource.services.ai.azure.com/anthropic", + client=upstream, + provider_specific_header=data["provider_specific_header"], + ) + finally: + await upstream.close() + + assert received[0].headers["anthropic-beta"] == "thinking-binding-controls-2026-08-01" + assert received[0].headers["anthropic-version"] == "2026-08-01" + assert json.loads(received[0].content)["thinking"]["block_binding"] == {"prefix_mismatch_behavior": "drop_block"} + + TYPESAFE_DANGER_ZONE_HEADERS = { "TypeSafe-Danger-Zone-Images": "true", "typesafe-danger-zone-long-context": "true", diff --git a/tests/unit/test_anthropic_beta_headers_filtering.py b/tests/unit/test_anthropic_beta_headers_filtering.py index 2a85122d91c..4d5ba1c5897 100644 --- a/tests/unit/test_anthropic_beta_headers_filtering.py +++ b/tests/unit/test_anthropic_beta_headers_filtering.py @@ -429,13 +429,13 @@ class TestAnthropicBetaHeadersFiltering: assert filtered == ["fine-grained-tool-streaming-2025-05-14"] @pytest.mark.parametrize( - "provider", ["anthropic", "bedrock", "bedrock_converse", "vertex_ai", "databricks"] + "provider", ["anthropic", "azure_ai", "bedrock", "bedrock_converse", "vertex_ai", "databricks"] ) def test_thinking_binding_controls_forwarded(self, provider): """`thinking.block_binding` (preserved thinking, Claude Fable 5.1) is only accepted alongside thinking-binding-controls-2026-08-01. The body field is forwarded untouched, so stripping the header (previously unknown, hence - dropped) makes Bedrock and Vertex reject the request with + dropped) makes Bedrock, Vertex and Azure AI Foundry reject the request with "thinking.adaptive.block_binding: Extra inputs are not permitted".""" filtered = filter_and_transform_beta_headers( beta_headers=["thinking-binding-controls-2026-08-01"], From ca64a3bef2ee0c6430f338c6967288d998ed7f7e Mon Sep 17 00:00:00 2001 From: Ishaan Date: Fri, 23 Oct 2026 16:39:23 -0700 Subject: [PATCH 7/7] feat(chat): add Moyai UI, model hub, images, and usage metrics (#40557) (#48108) * feat(chat): port Moyai UI with Model Hub and images * feat(chat): add Models sidebar with embedded Model Hub * feat(chat): polish Model Hub with provider cards and logos * refactor(chat): use a provider dropdown in Model Hub * feat(chat): keep workspace switching visible with a Chat identity * Add minimal opt-in Chat usage metrics and prioritize Chat in workspace selector * fix(chat): enable usage metrics by default only when chat is used * fix(chat): count requesting users with frontend-only metrics * feat(chat): send default-on request metrics to litellm-tel --- tests/unit/proxy/telemetry/test_endpoints.py | 5 +- tests/unit/proxy/telemetry/test_settings.py | 1 - .../src/app/chat/layout.integration.test.tsx | 142 ++++ .../src/app/chat/layout.test.tsx | 79 +- ui/litellm-dashboard/src/app/chat/layout.tsx | 16 +- .../app/chat/models/page.integration.test.tsx | 172 +++++ .../src/app/chat/models/page.tsx | 37 + .../app/chat/page.images.integration.test.tsx | 355 +++++++++ .../src/app/chat/page.integration.test.tsx | 4 +- ui/litellm-dashboard/src/app/chat/page.tsx | 719 +++++++++-------- .../chat/page.url-state.integration.test.tsx | 30 +- ui/litellm-dashboard/src/app/globals.css | 375 +++++++++ .../ModelHubProviders.integration.test.tsx | 85 ++ .../components/AIHub/ModelHubProviders.tsx | 103 +++ .../components/AIHub/ModelHubTable.test.tsx | 101 +++ .../src/components/AIHub/ModelHubTable.tsx | 167 ++-- .../AIHub/ModelHubTableColumns.test.tsx | 21 +- .../components/AIHub/ModelHubTableColumns.tsx | 354 +++++---- .../DashboardHeader.integration.test.tsx | 4 +- .../src/components/DashboardHeader.test.tsx | 8 +- .../Navbar/ViewSwitcher.integration.test.tsx | 179 +++++ .../components/Navbar/ViewSwitcher.test.tsx | 234 ------ .../src/components/Navbar/ViewSwitcher.tsx | 142 ++-- .../components/PublicModelHubTableColumns.tsx | 83 +- .../src/components/chat/ChatMessages.test.tsx | 104 +++ .../src/components/chat/ChatMessages.tsx | 104 ++- .../src/components/chat/ChatModelPicker.tsx | 122 +++ .../chat/ChatShell.integration.test.tsx | 173 +++++ .../chat/ChatShell.serverRootPath.test.ts | 6 +- .../src/components/chat/ChatShell.test.tsx | 91 --- .../src/components/chat/ChatShell.tsx | 316 +++++--- .../ConversationList.integration.test.tsx | 89 +++ .../src/components/chat/ConversationList.tsx | 119 ++- .../src/components/chat/chatImageStorage.ts | 125 +++ .../src/components/chat/design.md | 63 +- .../src/components/chat/imageHelpers.test.tsx | 113 +++ .../src/components/chat/imageHelpers.ts | 83 ++ .../components/chat/moyai/FONT-LICENSE.txt | 93 +++ .../src/components/chat/moyai/LICENSE.txt | 21 + .../chat/moyai/MoyaiChatComposer.test.tsx | 116 +++ .../chat/moyai/MoyaiChatComposer.tsx | 81 ++ .../chat/moyai/MoyaiChatHome.test.tsx | 86 +++ .../components/chat/moyai/MoyaiChatHome.tsx | 108 +++ .../chat/moyai/geist-mono-variable.woff2 | Bin 0 -> 23108 bytes .../chat/moyai/geist-variable.woff2 | Bin 0 -> 29288 bytes .../src/components/chat/moyai/moyai-chat.css | 729 ++++++++++++++++++ .../components/chat/runChatRequest.test.ts | 132 ++++ .../src/components/chat/runChatRequest.ts | 86 +++ .../src/components/chat/types.ts | 27 +- .../components/chat/useChatHistory.test.ts | 218 +++++- .../src/components/chat/useChatHistory.ts | 86 ++- .../src/components/model_filters.tsx | 135 +++- .../src/components/public_model_hub.test.tsx | 85 +- .../src/components/public_model_hub.tsx | 354 +++++---- .../shared/DataTable/DataTablePagination.tsx | 7 +- .../src/components/shared/MultiSelect.tsx | 4 +- .../src/contexts/PluginModeContext.test.tsx | 43 +- .../src/contexts/PluginModeContext.tsx | 16 +- .../src/lib/http/chatTelemetry.test.ts | 43 ++ .../src/lib/http/chatTelemetry.ts | 16 + .../src/lib/telemetry/chat.test.ts | 117 +++ .../src/lib/telemetry/chat.ts | 20 + 62 files changed, 6165 insertions(+), 1382 deletions(-) create mode 100644 ui/litellm-dashboard/src/app/chat/layout.integration.test.tsx create mode 100644 ui/litellm-dashboard/src/app/chat/models/page.integration.test.tsx create mode 100644 ui/litellm-dashboard/src/app/chat/models/page.tsx create mode 100644 ui/litellm-dashboard/src/app/chat/page.images.integration.test.tsx create mode 100644 ui/litellm-dashboard/src/components/AIHub/ModelHubProviders.integration.test.tsx create mode 100644 ui/litellm-dashboard/src/components/AIHub/ModelHubProviders.tsx create mode 100644 ui/litellm-dashboard/src/components/Navbar/ViewSwitcher.integration.test.tsx delete mode 100644 ui/litellm-dashboard/src/components/Navbar/ViewSwitcher.test.tsx create mode 100644 ui/litellm-dashboard/src/components/chat/ChatMessages.test.tsx create mode 100644 ui/litellm-dashboard/src/components/chat/ChatModelPicker.tsx create mode 100644 ui/litellm-dashboard/src/components/chat/ChatShell.integration.test.tsx delete mode 100644 ui/litellm-dashboard/src/components/chat/ChatShell.test.tsx create mode 100644 ui/litellm-dashboard/src/components/chat/ConversationList.integration.test.tsx create mode 100644 ui/litellm-dashboard/src/components/chat/chatImageStorage.ts create mode 100644 ui/litellm-dashboard/src/components/chat/imageHelpers.test.tsx create mode 100644 ui/litellm-dashboard/src/components/chat/imageHelpers.ts create mode 100644 ui/litellm-dashboard/src/components/chat/moyai/FONT-LICENSE.txt create mode 100644 ui/litellm-dashboard/src/components/chat/moyai/LICENSE.txt create mode 100644 ui/litellm-dashboard/src/components/chat/moyai/MoyaiChatComposer.test.tsx create mode 100644 ui/litellm-dashboard/src/components/chat/moyai/MoyaiChatComposer.tsx create mode 100644 ui/litellm-dashboard/src/components/chat/moyai/MoyaiChatHome.test.tsx create mode 100644 ui/litellm-dashboard/src/components/chat/moyai/MoyaiChatHome.tsx create mode 100644 ui/litellm-dashboard/src/components/chat/moyai/geist-mono-variable.woff2 create mode 100644 ui/litellm-dashboard/src/components/chat/moyai/geist-variable.woff2 create mode 100644 ui/litellm-dashboard/src/components/chat/moyai/moyai-chat.css create mode 100644 ui/litellm-dashboard/src/components/chat/runChatRequest.test.ts create mode 100644 ui/litellm-dashboard/src/components/chat/runChatRequest.ts create mode 100644 ui/litellm-dashboard/src/lib/http/chatTelemetry.test.ts create mode 100644 ui/litellm-dashboard/src/lib/http/chatTelemetry.ts create mode 100644 ui/litellm-dashboard/src/lib/telemetry/chat.test.ts create mode 100644 ui/litellm-dashboard/src/lib/telemetry/chat.ts diff --git a/tests/unit/proxy/telemetry/test_endpoints.py b/tests/unit/proxy/telemetry/test_endpoints.py index 3b5f66da0ed..c962c8df6a5 100644 --- a/tests/unit/proxy/telemetry/test_endpoints.py +++ b/tests/unit/proxy/telemetry/test_endpoints.py @@ -1,13 +1,12 @@ -from dataclasses import dataclass, field from collections.abc import Mapping +from dataclasses import dataclass, field from typing import Final -from typing_extensions import LiteralString - import pytest from fastapi import FastAPI from fastapi.testclient import TestClient from httpx import ASGITransport, AsyncClient, Response +from typing_extensions import LiteralString from litellm.proxy._types import CommonProxyErrors, LitellmUserRoles, UserAPIKeyAuth from litellm.proxy.auth.user_api_key_auth import user_api_key_auth diff --git a/tests/unit/proxy/telemetry/test_settings.py b/tests/unit/proxy/telemetry/test_settings.py index 0e9c180da4f..b88e6811113 100644 --- a/tests/unit/proxy/telemetry/test_settings.py +++ b/tests/unit/proxy/telemetry/test_settings.py @@ -1,7 +1,6 @@ from typing import Final import pytest - from pydantic import ValidationError from litellm.proxy.telemetry.settings import EnvPolicy, TelemetrySettings, describe_errors, env_policy, load_settings diff --git a/ui/litellm-dashboard/src/app/chat/layout.integration.test.tsx b/ui/litellm-dashboard/src/app/chat/layout.integration.test.tsx new file mode 100644 index 00000000000..12acf9192aa --- /dev/null +++ b/ui/litellm-dashboard/src/app/chat/layout.integration.test.tsx @@ -0,0 +1,142 @@ +import userEvent from "@testing-library/user-event"; +import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"; +import { renderWithProviders, screen, testQueryClient } from "@/../tests/test-utils"; +import { setServerRootPath } from "@/lib/serverRootPath"; +import { recordChatQuery } from "@/lib/telemetry/chat"; +import ChatLayout from "./layout"; + +vi.mock("@/lib/telemetry/chat", () => ({ recordChatQuery: vi.fn(async () => {}) })); + +const { router, route } = vi.hoisted(() => ({ + router: { push: vi.fn(), replace: vi.fn() }, + route: { pathname: "/ui/chat" }, +})); + +vi.mock("next/navigation", () => ({ + useRouter: () => router, + usePathname: () => route.pathname, +})); + +vi.mock("@/app/(dashboard)/hooks/useAuthorized", () => ({ + default: () => ({ + accessToken: "chat-layout-test-token", + userId: "chat-layout-user", + userEmail: "developer@example.com", + userRole: "internal_user", + premiumUser: false, + isLoading: false, + isAuthorized: true, + }), +})); + +vi.mock("@/app/(dashboard)/hooks/uiSettings/useUISettings", () => ({ + useUISettings: () => ({ data: { values: { enable_chat_ui: true } }, isLoading: false }), +})); + +const plugin = { name: "observability", display_name: "Observability", url: "https://plugin.example.com" }; +const originalLocation = Object.getOwnPropertyDescriptor(window, "location")!; +const assign = vi.fn(); + +describe("Chat layout product navigation", () => { + beforeEach(() => { + route.pathname = "/ui/chat"; + localStorage.clear(); + localStorage.setItem("litellm_plugin_mode", plugin.name); + testQueryClient.clear(); + assign.mockReset(); + Object.defineProperty(window, "location", { + configurable: true, + value: { ...window.location, assign }, + }); + vi.stubGlobal( + "fetch", + vi.fn(async (input: RequestInfo | URL, init?: RequestInit) => { + const url = input instanceof Request ? input.url : String(input); + const path = new URL(url, "http://localhost").pathname; + switch (path) { + case "/api/plugins": + expect(new Headers(init?.headers).get("Authorization")).toBe("Bearer chat-layout-test-token"); + return Response.json([plugin]); + case "/get/ui_theme_settings": + return Response.json({ values: {} }); + case "/sso/get/ui_settings": + return Response.json({}); + default: + throw new Error(`Unexpected request: ${path}`); + } + }), + ); + }); + + afterEach(() => { + vi.restoreAllMocks(); + testQueryClient.clear(); + localStorage.clear(); + setServerRootPath("/"); + Object.defineProperty(window, "location", originalLocation); + vi.unstubAllGlobals(); + vi.clearAllMocks(); + }); + + it.each([ + { pathname: "/ui/chat", root: "/", gateway: "/ui/" }, + { pathname: "/gateway/ui/chat/models/", root: "/gateway", gateway: "/gateway/ui/" }, + ])( + "should render Chat at $pathname and clear a saved plugin when returning to AI Gateway", + async ({ pathname, root, gateway }) => { + route.pathname = pathname; + setServerRootPath(root); + const user = userEvent.setup(); + renderWithProviders( + +

Chat content

+ , + ); + + expect(screen.getByRole("heading", { name: "Chat content" })).toBeVisible(); + expect(recordChatQuery).not.toHaveBeenCalled(); + await user.click(screen.getByRole("button", { name: "Switch workspace, current: Chat" })); + expect(await screen.findByRole("menuitem", { name: "Observability" })).toBeInTheDocument(); + expect(localStorage.getItem("litellm_plugin_mode")).toBe(plugin.name); + + await user.click(screen.getByRole("menuitem", { name: "AI Gateway", exact: true })); + + expect(localStorage.getItem("litellm_plugin_mode")).toBe("ai-gateway"); + expect(assign).toHaveBeenCalledWith(gateway); + expect(recordChatQuery).not.toHaveBeenCalled(); + }, + ); + + it.each(["SecurityError", "QuotaExceededError"])( + "should still navigate to AI Gateway when browser storage throws %s", + async (errorName) => { + if (errorName === "SecurityError") { + vi.spyOn(Storage.prototype, "getItem").mockImplementation(() => { + throw new DOMException("Storage is unavailable", errorName); + }); + vi.spyOn(Storage.prototype, "removeItem").mockImplementation(() => { + throw new DOMException("Storage is unavailable", errorName); + }); + } + vi.spyOn(Storage.prototype, "setItem").mockImplementation(() => { + throw new DOMException("Storage is unavailable", errorName); + }); + const user = userEvent.setup(); + renderWithProviders( + +

Chat content

+ , + ); + + expect(screen.getByRole("heading", { name: "Chat content" })).toBeVisible(); + await user.click(screen.getByRole("button", { name: "Switch workspace, current: Chat" })); + expect(await screen.findByRole("menuitem", { name: "Observability" })).toBeInTheDocument(); + await user.click(screen.getByRole("menuitem", { name: "AI Gateway", exact: true })); + + expect(assign).toHaveBeenCalledWith("/ui/"); + if (errorName === "QuotaExceededError") { + expect(localStorage.getItem("litellm_plugin_mode")).toBeNull(); + } + }, + ); +}); diff --git a/ui/litellm-dashboard/src/app/chat/layout.test.tsx b/ui/litellm-dashboard/src/app/chat/layout.test.tsx index d455b84afa3..0b0cb96ddc9 100644 --- a/ui/litellm-dashboard/src/app/chat/layout.test.tsx +++ b/ui/litellm-dashboard/src/app/chat/layout.test.tsx @@ -1,23 +1,25 @@ -import { afterEach, describe, expect, it, vi } from "vitest"; -import { screen } from "@testing-library/react"; +import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"; +import { act, fireEvent, screen } from "@testing-library/react"; import { renderWithProviders as render } from "@/../tests/test-utils"; import ChatLayout from "./layout"; -const { mockUseAuthorized, mockUseUISettings, mockReplace, mockUiHref, state } = vi.hoisted(() => { +const { mockUseAuthorized, mockUseUISettings, mockReplace, mockUiHref, recordChatQuery, state } = vi.hoisted(() => { const state = { enableChatUI: false, isUISettingsLoading: false, isIdentityLoading: false, isAuthorized: true, + userId: "user-1" as string | null, }; return { state, mockReplace: vi.fn(), + recordChatQuery: vi.fn(async () => {}), mockUiHref: vi.fn((segment: string) => `/mocked-ui/${segment}`), mockUseAuthorized: vi.fn(() => ({ accessToken: "token-123", userRole: "Internal User", - userId: "user-1", + userId: state.userId, userEmail: "user@example.com", premiumUser: false, isLoading: state.isIdentityLoading, @@ -36,6 +38,7 @@ vi.mock("next/navigation", () => ({ vi.mock("@/app/(dashboard)/hooks/useAuthorized", () => ({ default: mockUseAuthorized })); vi.mock("@/app/(dashboard)/hooks/uiSettings/useUISettings", () => ({ useUISettings: mockUseUISettings })); vi.mock("@/utils/uiHref", () => ({ uiHref: mockUiHref })); +vi.mock("@/lib/telemetry/chat", () => ({ recordChatQuery })); vi.mock("@/components/navbar", () => ({ default: () =>
})); vi.mock("@/contexts/ThemeContext", () => ({ ThemeProvider: ({ children }: { children: React.ReactNode }) => <>{children}, @@ -43,21 +46,32 @@ vi.mock("@/contexts/ThemeContext", () => ({ vi.mock("@/contexts/ChatShellContext", () => ({ ChatShellProvider: ({ children }: { children: React.ReactNode }) => <>{children}, })); +vi.mock("@/contexts/PluginModeContext", () => ({ + PluginModeProvider: ({ children }: { children: React.ReactNode }) => <>{children}, +})); vi.mock("@/components/chat/ChatShell", () => ({ default: ({ children }: { children: React.ReactNode }) =>
{children}
, })); describe("ChatLayout", () => { + beforeEach(() => { + vi.stubGlobal("fetch", vi.fn()); + }); + afterEach(() => { state.enableChatUI = false; state.isUISettingsLoading = false; state.isIdentityLoading = false; state.isAuthorized = true; + state.userId = "user-1"; mockReplace.mockClear(); mockUiHref.mockClear(); + recordChatQuery.mockClear(); + vi.restoreAllMocks(); + vi.unstubAllGlobals(); }); - it("renders the chat shell when enable_chat_ui is on", () => { + it("should render the chat shell without recording a query when enable_chat_ui is on", () => { state.enableChatUI = true; render( @@ -67,9 +81,10 @@ describe("ChatLayout", () => { expect(screen.getByTestId("chat-shell")).toBeInTheDocument(); expect(screen.getByTestId("page-content")).toBeInTheDocument(); expect(mockReplace).not.toHaveBeenCalled(); + expect(recordChatQuery).not.toHaveBeenCalled(); }); - it("redirects to the dashboard when enable_chat_ui is off", () => { + it("should redirect without recording a query when enable_chat_ui is off", () => { state.enableChatUI = false; render( @@ -78,9 +93,10 @@ describe("ChatLayout", () => { ); expect(screen.queryByTestId("chat-shell")).not.toBeInTheDocument(); expect(mockReplace).toHaveBeenCalledWith("/mocked-ui/"); + expect(recordChatQuery).not.toHaveBeenCalled(); }); - it("renders nothing while UI settings are still loading", () => { + it("should wait for UI settings without recording a query", () => { state.isUISettingsLoading = true; render( @@ -89,9 +105,10 @@ describe("ChatLayout", () => { ); expect(screen.queryByTestId("chat-shell")).not.toBeInTheDocument(); expect(mockReplace).not.toHaveBeenCalled(); + expect(recordChatQuery).not.toHaveBeenCalled(); }); - it("renders nothing until the signed-in user's identity has loaded", () => { + it("should wait for the signed-in user's identity without recording a query", () => { state.enableChatUI = true; state.isIdentityLoading = true; render( @@ -101,9 +118,10 @@ describe("ChatLayout", () => { ); expect(screen.queryByTestId("chat-shell")).not.toBeInTheDocument(); expect(screen.queryByTestId("page-content")).not.toBeInTheDocument(); + expect(recordChatQuery).not.toHaveBeenCalled(); }); - it("renders nothing when the signed-in user's identity could not be confirmed", () => { + it("should hide Chat without recording a query when the signed-in user's identity cannot be confirmed", () => { state.enableChatUI = true; state.isAuthorized = false; render( @@ -113,5 +131,48 @@ describe("ChatLayout", () => { ); expect(screen.queryByTestId("chat-shell")).not.toBeInTheDocument(); expect(screen.queryByTestId("page-content")).not.toBeInTheDocument(); + expect(recordChatQuery).not.toHaveBeenCalled(); + }); + + it("should not record a query or fetch telemetry when opening, focusing, or revisiting Chat", () => { + state.enableChatUI = true; + const visibility = vi.spyOn(document, "visibilityState", "get").mockReturnValue("hidden"); + const view = render(Chat content); + + act(() => fireEvent.focus(window)); + expect(recordChatQuery).not.toHaveBeenCalled(); + + visibility.mockReturnValue("visible"); + act(() => fireEvent(document, new Event("visibilitychange"))); + expect(recordChatQuery).not.toHaveBeenCalled(); + + view.rerender(Another Chat page); + act(() => fireEvent.focus(window)); + + view.unmount(); + render(Chat revisited); + act(() => { + fireEvent.focus(window); + fireEvent(document, new Event("visibilitychange")); + }); + expect(recordChatQuery).not.toHaveBeenCalled(); + expect(fetch).not.toHaveBeenCalled(); + }); + + it("should not record a query or fetch telemetry when the signed-in account changes", () => { + state.enableChatUI = true; + const view = render(Chat content); + expect(recordChatQuery).not.toHaveBeenCalled(); + + state.userId = "user-2"; + view.rerender(Chat content); + expect(screen.getByTestId("chat-shell")).toBeInTheDocument(); + + state.userId = null; + state.isAuthorized = false; + view.rerender(Chat content); + act(() => fireEvent.focus(window)); + expect(recordChatQuery).not.toHaveBeenCalled(); + expect(fetch).not.toHaveBeenCalled(); }); }); diff --git a/ui/litellm-dashboard/src/app/chat/layout.tsx b/ui/litellm-dashboard/src/app/chat/layout.tsx index bfe494d2234..1da33d80ec5 100644 --- a/ui/litellm-dashboard/src/app/chat/layout.tsx +++ b/ui/litellm-dashboard/src/app/chat/layout.tsx @@ -4,11 +4,12 @@ import { Suspense, useEffect } from "react"; import { useRouter } from "next/navigation"; import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized"; import { useUISettings } from "@/app/(dashboard)/hooks/uiSettings/useUISettings"; -import Navbar from "@/components/navbar"; import { ThemeProvider } from "@/contexts/ThemeContext"; import { ChatShellProvider } from "@/contexts/ChatShellContext"; +import { PluginModeProvider } from "@/contexts/PluginModeContext"; import ChatShell from "@/components/chat/ChatShell"; import { uiHref } from "@/utils/uiHref"; +import "@/components/chat/moyai/moyai-chat.css"; // The nuqs Next adapter uses useSearchParams, so keep the chat tree behind Suspense. function ChatLayoutContent({ children }: { children: React.ReactNode }) { @@ -32,13 +33,16 @@ function ChatLayoutContent({ children }: { children: React.ReactNode }) { }, [blocked, router]); const identityReady = !isIdentityLoading && isAuthorized; + const chatReady = identityReady && !isUISettingsLoading && chatEnabled; - if (!identityReady || isUISettingsLoading || blocked) return null; + if (!chatReady) return null; return ( -
- + + + +
- {children} + + {children} +
diff --git a/ui/litellm-dashboard/src/app/chat/models/page.integration.test.tsx b/ui/litellm-dashboard/src/app/chat/models/page.integration.test.tsx new file mode 100644 index 00000000000..d1868558626 --- /dev/null +++ b/ui/litellm-dashboard/src/app/chat/models/page.integration.test.tsx @@ -0,0 +1,172 @@ +import userEvent from "@testing-library/user-event"; +import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"; + +import { renderWithProviders, screen, testQueryClient, waitFor, within } from "@/../tests/test-utils"; +import type { ModelHubData } from "@/components/AIHub/ModelHubTableColumns"; + +import ChatModelsPage from "./page"; + +const { apiGetMock, modelHubCallMock, shellState } = vi.hoisted(() => ({ + apiGetMock: vi.fn(), + modelHubCallMock: vi.fn(), + shellState: { accessToken: "chat-session-token", userRole: "proxy_admin", premiumUser: false }, +})); + +vi.mock("@/contexts/ChatShellContext", () => ({ + useChatShell: () => shellState, +})); + +vi.mock("next/navigation", () => ({ + useRouter: () => ({ push: vi.fn(), replace: vi.fn(), refresh: vi.fn() }), +})); + +vi.mock("@/app/(dashboard)/hooks/uiSettings/useUISettings", () => ({ + useUISettings: () => ({ data: { values: {} }, isLoading: false }), +})); + +vi.mock("@/components/networking", async (importOriginal) => { + const actual = await importOriginal(); + const hubInfo = { + docs_title: "Gateway models", + custom_docs_description: null, + litellm_version: "1.0.0", + useful_links: {}, + }; + return { + ...actual, + apiClient: { ...actual.apiClient, get: apiGetMock }, + modelHubCall: modelHubCallMock, + getConfigFieldSetting: vi.fn().mockResolvedValue({ field_value: true }), + getUiConfig: vi.fn().mockResolvedValue({}), + getProxyBaseUrl: () => "https://gateway.example.com", + getGatewayBaseUrl: () => "https://gateway.example.com", + getPublicModelHubInfo: vi.fn().mockResolvedValue(hubInfo), + getAgentsList: vi.fn().mockResolvedValue({ agents: [] }), + fetchMCPServers: vi.fn().mockResolvedValue([]), + getClaudeCodePluginsList: vi.fn().mockResolvedValue({ plugins: [] }), + agentHubPublicModelsCall: vi.fn().mockResolvedValue([]), + mcpHubPublicServersCall: vi.fn().mockResolvedValue([]), + skillHubPublicCall: vi.fn().mockResolvedValue({ plugins: [] }), + }; +}); + +const PRIVATE_MODEL = { + model_group: "research-private", + description: "An internal model for reviewing research reports.", + providers: ["openai"], + mode: "chat", + max_input_tokens: 128_000, + max_output_tokens: 4_096, + input_cost_per_token: 0.000002, + output_cost_per_token: 0.000008, + supports_function_calling: true, + supports_parallel_function_calling: false, + supports_vision: true, + is_public_model_group: false, +} satisfies ModelHubData; + +const IMAGE_MODEL: ModelHubData = { + ...PRIVATE_MODEL, + model_group: "image-studio", + description: "Create illustrations for presentations.", + mode: "image_generation", +}; + +const PUBLISHED_MODEL = { + ...PRIVATE_MODEL, + model_group: "published-chat", + description: "The shared model published for gateway users.", + is_public_model_group: true, +} satisfies ModelHubData; + +const PUBLIC_MODELS_PATH = "/public/v1/model_hub"; +const FACETS: Record = { + [`${PUBLIC_MODELS_PATH}/providers`]: ["openai"], + [`${PUBLIC_MODELS_PATH}/modes`]: ["chat"], + [`${PUBLIC_MODELS_PATH}/features`]: ["function_calling", "vision"], +}; + +describe("Chat Models page", () => { + beforeEach(() => { + shellState.userRole = "proxy_admin"; + modelHubCallMock.mockResolvedValue({ data: [PRIVATE_MODEL, IMAGE_MODEL] }); + apiGetMock.mockImplementation(async (path: string) => { + if (path in FACETS) { + return { + data: FACETS[path], + meta: { page: 1, page_size: 100, has_more: false }, + links: { self: path, prev: null, next: null }, + }; + } + if (path === PUBLIC_MODELS_PATH) { + return { + data: [PUBLISHED_MODEL], + meta: { total_count: 1, page: 1, page_size: 50, total_pages: 1 }, + links: { + self: path, + first: path, + prev: null, + next: null, + last: path, + }, + }; + } + throw new Error(`Unexpected catalog request: ${path}`); + }); + vi.stubGlobal( + "fetch", + vi.fn(async (input: RequestInfo | URL) => { + if (String(input) !== "https://gateway.example.com/get/ui_theme_settings") { + throw new Error(`Unexpected fetch: ${String(input)}`); + } + return new Response(JSON.stringify({ values: {} }), { + status: 200, + headers: { "Content-Type": "application/json" }, + }); + }), + ); + }); + + afterEach(() => { + testQueryClient.clear(); + vi.clearAllMocks(); + vi.unstubAllGlobals(); + }); + + it("should render gateway models for admins and let them search and inspect model details", async () => { + const user = userEvent.setup(); + renderWithProviders(); + + expect(screen.getByRole("heading", { name: "Models", level: 1 })).toBeInTheDocument(); + expect(await screen.findByRole("button", { name: /^research-private/ })).toBeInTheDocument(); + expect(screen.getByRole("button", { name: /^image-studio/ })).toBeInTheDocument(); + + await user.type(screen.getByPlaceholderText("Search model names..."), "research"); + await waitFor(() => expect(screen.queryByRole("button", { name: /^image-studio/ })).not.toBeInTheDocument()); + await user.click(screen.getByRole("button", { name: /^research-private/ })); + + const details = await screen.findByRole("dialog", { name: PRIVATE_MODEL.model_group }); + expect(within(details).getByText(PRIVATE_MODEL.description)).toBeInTheDocument(); + expect(within(details).getByText("128,000")).toBeInTheDocument(); + expect(within(details).getByText("$2.00")).toBeInTheDocument(); + expect(within(details).getByText("Vision")).toBeInTheDocument(); + expect(within(details).getByText("Function Calling")).toBeInTheDocument(); + }); + + it("should show regular users the published catalog and its model capabilities", async () => { + shellState.userRole = "internal_user"; + const user = userEvent.setup(); + renderWithProviders(); + + const publishedModel = await screen.findByRole("button", { name: /^published-chat/ }); + expect(screen.queryByRole("button", { name: /^research-private/ })).not.toBeInTheDocument(); + expect(screen.queryByRole("button", { name: /^image-studio/ })).not.toBeInTheDocument(); + await user.click(publishedModel); + + const details = await screen.findByRole("dialog", { name: PUBLISHED_MODEL.model_group }); + expect(within(details).getByText(PUBLISHED_MODEL.description)).toBeInTheDocument(); + expect(within(details).getByText("128,000")).toBeInTheDocument(); + expect(within(details).getByText("$2.0000")).toBeInTheDocument(); + expect(within(details).getByText("Vision")).toBeInTheDocument(); + }); +}); diff --git a/ui/litellm-dashboard/src/app/chat/models/page.tsx b/ui/litellm-dashboard/src/app/chat/models/page.tsx new file mode 100644 index 00000000000..8011235821e --- /dev/null +++ b/ui/litellm-dashboard/src/app/chat/models/page.tsx @@ -0,0 +1,37 @@ +"use client"; + +import ModelHubTable from "@/components/AIHub/ModelHubTable"; +import PublicModelHub from "@/components/public_model_hub"; +import { useChatShell } from "@/contexts/ChatShellContext"; +import { isAdminRole } from "@/utils/roles"; + +export default function ChatModelsPage() { + const { accessToken, userRole, premiumUser } = useChatShell(); + + return ( +
+
+
+

Model Hub

+

Models

+

+ Find the right model. Explore providers, compare pricing, and see what each model can do. +

+
+ {isAdminRole(userRole) ? ( + + ) : ( + + )} +
+
+ ); +} diff --git a/ui/litellm-dashboard/src/app/chat/page.images.integration.test.tsx b/ui/litellm-dashboard/src/app/chat/page.images.integration.test.tsx new file mode 100644 index 00000000000..4f34fe22e58 --- /dev/null +++ b/ui/litellm-dashboard/src/app/chat/page.images.integration.test.tsx @@ -0,0 +1,355 @@ +import { act, fireEvent, screen, waitFor } from "@testing-library/react"; +import userEvent from "@testing-library/user-event"; +import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"; +import { renderWithProviders, testQueryClient } from "@/../tests/test-utils"; +import { ChatShellProvider } from "@/contexts/ChatShellContext"; +import { recordChatQuery } from "@/lib/telemetry/chat"; +import ChatConversationPage from "./page"; + +vi.mock("@/lib/telemetry/chat", () => ({ recordChatQuery: vi.fn(async () => {}) })); + +vi.mock("next/navigation", () => ({ + useRouter: () => ({ push: vi.fn(), replace: vi.fn() }), +})); + +type RecordedRequest = { pathname: string; body: Record }; +const requests: RecordedRequest[] = []; +const pixel = "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAQAAAC1HAwCAAAAC0lEQVR42mP8/x8AAwMCAO+aRZQAAAAASUVORK5CYII="; +let models: { model_group: string; mode: string; providers: string[] }[]; +let imageResponse: () => Response; +let pendingGeneration: boolean; +let generationSignal: AbortSignal | null; +let includeMCPEvents: boolean; + +function responseStream() { + return new Response( + [ + { type: "response.output_text.delta", delta: "I can see " }, + ...(includeMCPEvents + ? ["search_notes", "read_note"].map((name) => ({ + type: "response.output_item.done", + item: { id: name, type: "mcp_call", name, server_label: "research-tools", output: "Found a note" }, + })) + : []), + { type: "response.output_text.delta", delta: "the uploaded image." }, + { type: "response.completed", response: { id: `response-${requests.length}`, output: [] } }, + ] + .map((event) => `event: ${event.type}\ndata: ${JSON.stringify(event)}\n\n`) + .join(""), + { headers: { "Content-Type": "text/event-stream" } }, + ); +} + +function renderChat() { + return renderWithProviders( + + + , + ); +} + +async function uploadImage(user: ReturnType, name = "diagram.png") { + const bytes = Uint8Array.from(atob(pixel), (character) => character.charCodeAt(0)); + await user.upload(screen.getByLabelText("Upload images"), new File([bytes], name, { type: "image/png" })); + await screen.findByRole("img", { name }); +} + +async function chooseModel(user: ReturnType, current: string, next: string) { + await user.click(screen.getByRole("button", { name: current, exact: true })); + await user.click(await screen.findByRole("button", { name: new RegExp(`^${next}`) })); + await waitFor(() => expect(screen.queryByRole("button", { name: "All", exact: true })).not.toBeInTheDocument()); +} + +beforeEach(() => { + localStorage.clear(); + testQueryClient.clear(); + requests.length = 0; + vi.mocked(recordChatQuery).mockClear(); + models = [ + { model_group: "research-chat", mode: "chat", providers: ["openai"] }, + { model_group: "studio-image", mode: "image_generation", providers: ["openai"] }, + ]; + imageResponse = () => Response.json({ created: 1, data: [{ b64_json: pixel }] }); + pendingGeneration = false; + generationSignal = null; + includeMCPEvents = false; + vi.stubGlobal( + "fetch", + vi.fn(async (input: RequestInfo | URL, init?: RequestInit) => { + const request = + input instanceof Request ? input : new Request(new URL(String(input), window.location.origin), init); + const pathname = new URL(request.url).pathname; + if (pathname.endsWith("/model_group/info")) return Response.json({ data: models }); + if (pathname.endsWith("/v1/mcp/server")) { + return Response.json([{ server_id: "research-tools", server_name: "research-tools" }]); + } + if (pathname.endsWith("/mcp-rest/tools/list")) return Response.json({ tools: [] }); + if (pathname.endsWith("/responses") || pathname.endsWith("/images/generations")) { + requests.push({ pathname, body: JSON.parse(await request.text()) }); + if (pathname.endsWith("/responses")) return responseStream(); + if (pendingGeneration) { + generationSignal = request.signal; + return new Promise((_resolve, reject) => { + request.signal.addEventListener("abort", () => reject(new DOMException("Stopped", "AbortError")), { + once: true, + }); + }); + } + return imageResponse(); + } + return Response.json({}, { status: 404 }); + }), + ); +}); + +afterEach(() => { + vi.unstubAllGlobals(); + localStorage.clear(); +}); + +describe("chat image workflows", () => { + it("should render the permitted models and image controls for a research team member", async () => { + renderChat(); + expect(await screen.findByRole("button", { name: "research-chat" })).toBeInTheDocument(); + expect(screen.getByRole("button", { name: "Attach images" })).toBeInTheDocument(); + await userEvent.setup().click(screen.getByRole("button", { name: "research-chat", exact: true })); + expect(await screen.findByRole("button", { name: /^studio-image/ })).toBeInTheDocument(); + expect(screen.getByRole("button", { name: "Send" })).toBeDisabled(); + expect(recordChatQuery).not.toHaveBeenCalled(); + }); + + it("should restore the selected chat model after switching modes and reopening the page", async () => { + const user = userEvent.setup(); + const view = renderChat(); + await screen.findByRole("button", { name: "research-chat" }); + await chooseModel(user, "research-chat", "studio-image"); + expect(screen.getByRole("button", { name: "studio-image" })).toBeInTheDocument(); + await chooseModel(user, "studio-image", "research-chat"); + view.unmount(); + renderChat(); + expect(await screen.findByRole("button", { name: "research-chat" })).toBeInTheDocument(); + expect(screen.getByRole("textbox", { name: "Message", exact: true })).toBeInTheDocument(); + }); + + it("should preview and remove an uploaded image before sending", async () => { + const user = userEvent.setup(); + renderChat(); + await screen.findByRole("button", { name: "research-chat" }); + await uploadImage(user); + expect(screen.getByRole("button", { name: "Send" })).toBeEnabled(); + await user.click(screen.getByRole("button", { name: "Remove diagram.png" })); + expect(screen.queryByRole("img", { name: "diagram.png" })).not.toBeInTheDocument(); + expect(screen.getByRole("button", { name: "Send" })).toBeDisabled(); + expect(requests).toHaveLength(0); + expect(recordChatQuery).not.toHaveBeenCalled(); + }); + + it("should send the image and preserve it when editing and resending the message", async () => { + const user = userEvent.setup(); + renderChat(); + await screen.findByRole("button", { name: "research-chat" }); + await uploadImage(user); + await user.type(screen.getByRole("textbox", { name: "Message", exact: true }), "Explain this diagram"); + await user.click(screen.getByRole("button", { name: "Send" })); + await screen.findByText("I can see the uploaded image."); + expect(recordChatQuery).toHaveBeenCalledExactlyOnceWith("media-test-user"); + expect(requests[0].body).toMatchObject({ + model: "research-chat", + input: [ + { + role: "user", + content: [ + { type: "input_text", text: "Explain this diagram" }, + { type: "input_image", image_url: `data:image/png;base64,${pixel}` }, + ], + }, + ], + }); + expect(screen.getByRole("img", { name: "diagram.png" })).toBeInTheDocument(); + await user.hover(screen.getAllByText("Explain this diagram").at(-1)!); + const editButton = await screen.findByRole("button", { name: "Edit message" }); + act(() => fireEvent.click(editButton)); + await user.clear(screen.getByRole("textbox", { name: "Edit message" })); + await user.type(screen.getByRole("textbox", { name: "Edit message" }), "Explain the arrows"); + await user.click(screen.getByRole("button", { name: "Save & Send" })); + await waitFor(() => expect(requests).toHaveLength(2)); + expect(requests[1].body).toMatchObject({ + input: [ + { + role: "user", + content: [ + { type: "input_text", text: "Explain the arrows" }, + { type: "input_image", image_url: `data:image/png;base64,${pixel}` }, + ], + }, + ], + }); + expect(requests[1].body).not.toHaveProperty("previous_response_id"); + expect(screen.queryByText("Explain this diagram", { selector: "p" })).not.toBeInTheDocument(); + expect(recordChatQuery).toHaveBeenCalledTimes(2); + }); + + it("should preserve a chat turn when editing it while an image model is selected", async () => { + const user = userEvent.setup(); + renderChat(); + await screen.findByRole("button", { name: "research-chat" }); + await user.type(screen.getByRole("textbox", { name: "Message", exact: true }), "Explain the experiment"); + await user.click(screen.getByRole("button", { name: "Send" })); + await screen.findByText("I can see the uploaded image."); + await chooseModel(user, "research-chat", "studio-image"); + await user.hover(screen.getAllByText("Explain the experiment").at(-1)!); + const editButton = await screen.findByRole("button", { name: "Edit message" }); + act(() => fireEvent.click(editButton)); + await user.clear(screen.getByRole("textbox", { name: "Edit message" })); + await user.type(screen.getByRole("textbox", { name: "Edit message" }), "Explain the findings"); + await user.click(screen.getByRole("button", { name: "Save & Send" })); + expect(await screen.findByRole("alert")).toHaveTextContent("Choose a chat model before resending this message."); + expect(screen.getByRole("heading", { name: "Explain the experiment" })).toBeInTheDocument(); + expect(screen.getByText("I can see the uploaded image.")).toBeInTheDocument(); + expect(requests).toHaveLength(1); + expect(recordChatQuery).toHaveBeenCalledExactlyOnceWith("media-test-user"); + }); + + it("should continue an image conversation using the completed response context", async () => { + const user = userEvent.setup(); + renderChat(); + await screen.findByRole("button", { name: "research-chat" }); + await uploadImage(user); + await user.click(screen.getByRole("button", { name: "Send" })); + await screen.findByText("I can see the uploaded image."); + await user.type(screen.getByRole("textbox", { name: "Message", exact: true }), "What is the main color?"); + await user.click(screen.getByRole("button", { name: "Send" })); + await waitFor(() => expect(requests).toHaveLength(2)); + expect(requests[1].body).toMatchObject({ + previous_response_id: "response-1", + input: [{ role: "user", content: "What is the main color?" }], + }); + expect(screen.getByRole("img", { name: "diagram.png" })).toBeInTheDocument(); + }); + + it("should generate and render an image through the selected Model Hub model", async () => { + vi.mocked(recordChatQuery).mockReturnValueOnce(new Promise(() => {})); + const user = userEvent.setup(); + renderChat(); + await screen.findByRole("button", { name: "research-chat" }); + await chooseModel(user, "research-chat", "studio-image"); + expect(screen.getByRole("button", { name: "studio-image" })).toBeInTheDocument(); + await user.type(screen.getByRole("textbox", { name: "Message", exact: true }), "A blue research lab"); + await user.click(screen.getByRole("button", { name: "Generate image" })); + const result = await screen.findByRole("img", { name: "Generated image 1" }); + expect(result).toHaveAttribute("src", `data:image/png;base64,${pixel}`); + expect(requests).toEqual([ + { pathname: "/images/generations", body: { model: "studio-image", prompt: "A blue research lab" } }, + ]); + expect(recordChatQuery).toHaveBeenCalledExactlyOnceWith("media-test-user"); + }); + + it("should list no image models when the account has none permitted", async () => { + models = models.filter((model) => model.mode === "chat"); + const user = userEvent.setup(); + renderChat(); + await user.click(await screen.findByRole("button", { name: "research-chat" })); + await user.click(screen.getByRole("button", { name: "Images", exact: true })); + expect(screen.getByText("No models found.")).toBeInTheDocument(); + expect(screen.queryByRole("button", { name: /^studio-image/ })).not.toBeInTheDocument(); + expect(screen.queryByRole("button", { name: "Generate image" })).not.toBeInTheDocument(); + expect(requests).toHaveLength(0); + }); + + it("should show the provider error when image generation fails", async () => { + imageResponse = () => + Response.json({ error: { message: "Image access is disabled for this key" } }, { status: 403 }); + const user = userEvent.setup(); + renderChat(); + await screen.findByRole("button", { name: "research-chat" }); + await chooseModel(user, "research-chat", "studio-image"); + await user.type(screen.getByRole("textbox", { name: "Message", exact: true }), "A blue lab"); + await user.click(screen.getByRole("button", { name: "Generate image" })); + expect(await screen.findByText(/Image access is disabled for this key/)).toBeInTheDocument(); + expect(screen.queryByRole("img", { name: "Generated image 1" })).not.toBeInTheDocument(); + expect(recordChatQuery).toHaveBeenCalledExactlyOnceWith("media-test-user"); + }); + + it("should abort image generation when the user stops the request", async () => { + pendingGeneration = true; + const user = userEvent.setup(); + renderChat(); + await screen.findByRole("button", { name: "research-chat" }); + await chooseModel(user, "research-chat", "studio-image"); + await user.type(screen.getByRole("textbox", { name: "Message", exact: true }), "A blue lab"); + await user.click(screen.getByRole("button", { name: "Generate image" })); + await waitFor(() => expect(generationSignal).not.toBeNull()); + act(() => { + fireEvent.change(screen.getByRole("textbox", { name: "Message", exact: true }), { + target: { value: "Another image" }, + }); + fireEvent.keyDown(screen.getByRole("textbox", { name: "Message", exact: true }), { key: "Enter" }); + }); + expect(requests).toHaveLength(1); + expect(recordChatQuery).toHaveBeenCalledExactlyOnceWith("media-test-user"); + await user.click(screen.getByRole("button", { name: "Stop request" })); + expect(generationSignal?.aborted).toBe(true); + expect(await screen.findByText("[stopped]")).toBeInTheDocument(); + expect(screen.queryByRole("img", { name: "Generated image 1" })).not.toBeInTheDocument(); + expect(recordChatQuery).toHaveBeenCalledTimes(1); + }); + + it("should reject unsupported pasted files without issuing a request", async () => { + renderChat(); + await screen.findByRole("button", { name: "research-chat" }); + act(() => { + fireEvent.paste(screen.getByRole("textbox", { name: "Message", exact: true }), { + clipboardData: { files: [new File(["svg"], "diagram.svg", { type: "image/svg+xml" })] }, + }); + }); + expect(await screen.findByRole("alert")).toHaveTextContent("Choose a PNG, JPEG, WebP, or GIF image."); + expect(requests).toHaveLength(0); + expect(recordChatQuery).not.toHaveBeenCalled(); + }); + + it("should not count a blank submission as a query", async () => { + renderChat(); + await screen.findByRole("button", { name: "research-chat" }); + act(() => { + fireEvent.change(screen.getByRole("textbox", { name: "Message", exact: true }), { + target: { value: " " }, + }); + fireEvent.keyDown(screen.getByRole("textbox", { name: "Message", exact: true }), { key: "Enter" }); + }); + expect(requests).toHaveLength(0); + expect(recordChatQuery).not.toHaveBeenCalled(); + }); + + it("should count one query for a message that streams text and executes multiple MCP tools", async () => { + includeMCPEvents = true; + const user = userEvent.setup(); + renderChat(); + await screen.findByRole("button", { name: "research-chat" }); + await user.click(screen.getByRole("button", { name: "Connect tools" })); + const serverSwitch = await screen.findByRole("switch"); + await user.click(serverSwitch); + await waitFor(() => expect(screen.getByRole("switch")).toBeChecked()); + await user.click(screen.getByRole("button", { name: "Connect tools" })); + expect(recordChatQuery).not.toHaveBeenCalled(); + + act(() => { + fireEvent.change(screen.getByRole("textbox", { name: "Message", exact: true }), { + target: { value: "Search the research notes" }, + }); + }); + await user.click(screen.getByRole("button", { name: "Send" })); + await screen.findByText("I can see the uploaded image."); + + expect(requests).toHaveLength(1); + expect(requests[0].body).toMatchObject({ + tools: [{ type: "mcp", server_label: "research-tools", require_approval: "never" }], + }); + expect(recordChatQuery).toHaveBeenCalledExactlyOnceWith("media-test-user"); + }); +}); diff --git a/ui/litellm-dashboard/src/app/chat/page.integration.test.tsx b/ui/litellm-dashboard/src/app/chat/page.integration.test.tsx index 2020a66bd1c..e49f3a17094 100644 --- a/ui/litellm-dashboard/src/app/chat/page.integration.test.tsx +++ b/ui/litellm-dashboard/src/app/chat/page.integration.test.tsx @@ -72,7 +72,7 @@ const ON_TOTAL_LATENCY_INDEX = 24; async function sendOneMessage(onUrlUpdate?: OnUrlUpdateFunction): Promise { renderWithProviders(, { onUrlUpdate }); expect(await screen.findByRole("button", { name: /gpt-5\.4-mini/ })).toBeInTheDocument(); - fireEvent.change(screen.getByPlaceholderText("How can I help you today?"), { + fireEvent.change(screen.getByPlaceholderText("Describe a change, paste an error, or ask a question…"), { target: { value: "How much did this cost?" }, }); fireEvent.click(screen.getByRole("button", { name: "Send" })); @@ -157,7 +157,7 @@ describe("/ui/chat storage banner", () => { it("keeps the dismiss control amber on hover instead of the ghost variant's foreground", async () => { renderWithProviders(); - const banner = await screen.findByText("Chat history won't be saved in this browser session"); + const banner = await screen.findByText("Some chat history or images could not be saved or restored."); const dismiss = within(banner.parentElement!).getByRole("button"); expect(dismiss).toHaveClass("hover:text-warning/80"); expect(dismiss).not.toHaveClass("hover:text-foreground"); diff --git a/ui/litellm-dashboard/src/app/chat/page.tsx b/ui/litellm-dashboard/src/app/chat/page.tsx index fdf6a1e02a3..dbb578a0e3d 100644 --- a/ui/litellm-dashboard/src/app/chat/page.tsx +++ b/ui/litellm-dashboard/src/app/chat/page.tsx @@ -1,11 +1,8 @@ "use client"; import React, { useCallback, useEffect, useRef, useState, useLayoutEffect } from "react"; -import { Plus, ChevronDown, Check, X } from "lucide-react"; +import { Plus, ChevronDown, Paperclip, ImageIcon, X } from "lucide-react"; import { Popover, PopoverContent, PopoverTrigger } from "@/components/ui/popover"; -import { Skeleton } from "@/components/ui/skeleton"; -import { ScrollArea } from "@/components/ui/scroll-area"; -import { Input } from "@/components/ui/input"; import { Button } from "@/components/ui/button"; import { toast } from "@/lib/toast"; import { useRouter } from "next/navigation"; @@ -13,41 +10,22 @@ import { useChatShell } from "@/contexts/ChatShellContext"; import { getChatRoutes } from "@/components/chat/ChatShell"; import ChatMessages from "@/components/chat/ChatMessages"; import MCPConnectPicker from "@/components/chat/MCPConnectPicker"; -import { fetchAvailableModels } from "@/components/llm_calls/fetch_models"; -import { makeOpenAIResponsesRequest } from "@/components/llm_calls/responses_api"; -import type { TokenUsage } from "@/components/chat_ui/ResponseMetrics"; -import type { MCPEvent } from "@/components/chat/types"; -import { getProviderLogoAndName } from "@/components/provider_info_helpers"; +import { fetchAvailableModels, type ModelGroup } from "@/components/llm_calls/fetch_models"; +import type { AssistantMessageUpdate, ChatImageAttachment, ChatMessage } from "@/components/chat/types"; +import ChatModelPicker, { isChatModel, modelChatMode, type ChatMode } from "@/components/chat/ChatModelPicker"; +import MoyaiChatComposer from "@/components/chat/moyai/MoyaiChatComposer"; +import MoyaiChatHome from "@/components/chat/moyai/MoyaiChatHome"; +import { runChatRequest } from "@/components/chat/runChatRequest"; +import { readChatImage, MAX_CHAT_IMAGES } from "@/components/chat/imageHelpers"; import { useQueryState } from "nuqs"; +import { recordChatQuery } from "@/lib/telemetry/chat"; -const SUGGESTIONS = ["Write", "Learn", "Code", "Brainstorm"]; const LOCALSTORAGE_MODEL_KEY = "litellm_chat_selected_model"; -function getGreeting(): string { - const h = new Date().getHours(); - if (h >= 5 && h < 12) return "Good morning"; - if (h >= 12 && h < 17) return "Good afternoon"; - return "Good evening"; -} +type ReplayMessage = Pick; -// Extract provider from model name for logo lookup. -// Handles prefixed models ("groq/llama-3"), and detects well-known providers by keyword. -function getProviderFromModelName(modelName: string): string { - if (!modelName) return ""; - const lower = modelName.toLowerCase(); - const slash = lower.indexOf("/"); - if (slash > 0) return lower.slice(0, slash); - // Keyword matching — order matters (more specific first) - if (lower.includes("claude")) return "anthropic"; - if (lower.includes("gemini")) return "gemini"; - if (lower.includes("gpt") || lower.includes("chatgpt") || /^o[0-9]/.test(lower)) return "openai"; - if (lower.includes("mistral") || lower.includes("codestral")) return "mistral"; - if (lower.includes("llama")) return "meta_llama"; - if (lower.includes("deepseek")) return "deepseek"; - if (lower.includes("grok")) return "xai"; - if (lower.includes("command")) return "cohere"; - if (lower.includes("nova") || lower.includes("titan")) return "bedrock"; - return ""; +function hasMessageContent(text: string, images: ChatImageAttachment[], mode: ChatMode): boolean { + return Boolean(text.trim()) || (mode === "chat" && images.length > 0); } export default function ChatConversationPage() { @@ -56,7 +34,7 @@ export default function ChatConversationPage() { const { accessToken, userId, - userEmail, + conversations, selectedMCPServers, setSelectedMCPServers, activeConversationId, @@ -70,10 +48,15 @@ export default function ChatConversationPage() { } = useChatShell(); const [selectedModel, setSelectedModel] = useState(null); - const [models, setModels] = useState([]); + const [models, setModels] = useState([]); + const [mode, setMode] = useState("chat"); + const [attachments, setAttachments] = useState([]); + const [attachmentError, setAttachmentError] = useState(null); + const [isReadingImages, setIsReadingImages] = useState(false); + const imageInputRef = useRef(null); + const readImagesRef = useRef(false); + const draftVersionRef = useRef(0); const [isLoadingModels, setIsLoadingModels] = useState(true); - const [modelSelectorOpen, setModelSelectorOpen] = useState(false); - const [modelSearchText, setModelSearchText] = useState(""); const [responsesSessionId, setResponsesSessionId] = useState(null); const [prevConversationIdForSessionReset, setPrevConversationIdForSessionReset] = useState(activeConversationId); @@ -83,7 +66,7 @@ export default function ChatConversationPage() { const [storageBannerDismissed, setStorageBannerDismissed] = useState(false); const abortControllerRef = useRef(null); - const textareaRef = useRef(null); + const requestConversationRef = useRef(null); const messagesScrollRef = useRef(null); const [showScrollButton, setShowScrollButton] = useState(false); const streamScrollLock = useRef(null); @@ -97,12 +80,14 @@ export default function ChatConversationPage() { if (!accessToken) return; fetchAvailableModels(accessToken) .then((data) => { - const names = (data || []).map((m: { model_group?: string }) => m.model_group ?? "").filter(Boolean); - setModels(names); + const available = (data || []).filter(isChatModel); + const names = available.map((model) => model.model_group); + setModels(available); try { const saved = localStorage.getItem(LOCALSTORAGE_MODEL_KEY); if (saved && names.includes(saved)) { setSelectedModel(saved); + setMode(modelChatMode(available.find((model) => model.model_group === saved)!)); return; } } catch { @@ -110,7 +95,7 @@ export default function ChatConversationPage() { } if (names.length > 0) { setSelectedModel(names[0]); - localStorage.setItem(LOCALSTORAGE_MODEL_KEY, names[0]); + setMode(modelChatMode(available[0])); } }) .catch(() => toast.error("Could not load models")) @@ -122,20 +107,87 @@ export default function ChatConversationPage() { if (activeConversationId !== prevConversationIdForSessionReset) { setPrevConversationIdForSessionReset(activeConversationId); setResponsesSessionId(null); + setAttachments([]); + setAttachmentError(null); + setInputText(""); } - const selectModel = useCallback((model: string) => { - setSelectedModel(model); - localStorage.setItem(LOCALSTORAGE_MODEL_KEY, model); - setModelSelectorOpen(false); - setModelSearchText(""); - }, []); + useEffect(() => { + draftVersionRef.current += 1; + if (abortControllerRef.current && requestConversationRef.current !== activeConversationId) { + abortControllerRef.current.abort(); + } + }, [activeConversationId]); + + const selectModel = useCallback( + (model: ModelGroup) => { + if (modelChatMode(model) === "image" && attachments.length > 0) { + setAttachmentError("Remove the attached images before switching to image generation."); + return; + } + setSelectedModel(model.model_group); + setMode(modelChatMode(model)); + setResponsesSessionId(null); + setAttachmentError(null); + try { + localStorage.setItem(LOCALSTORAGE_MODEL_KEY, model.model_group); + } catch { + return; + } + }, + [attachments.length], + ); + + const addImages = async (files: File[]) => { + if (mode !== "chat" || files.length === 0) return; + if (isStreaming || readImagesRef.current) return; + if (attachments.length + files.length > MAX_CHAT_IMAGES) { + setAttachmentError(`Attach up to ${MAX_CHAT_IMAGES} images per message.`); + return; + } + const draftVersion = draftVersionRef.current; + readImagesRef.current = true; + setIsReadingImages(true); + setAttachmentError(null); + try { + const images = await Promise.all(files.map(readChatImage)); + if (draftVersion === draftVersionRef.current) setAttachments((current) => [...current, ...images]); + } catch (error) { + if (draftVersion === draftVersionRef.current) + setAttachmentError(error instanceof Error ? error.message : "Could not read these images."); + } finally { + readImagesRef.current = false; + setIsReadingImages(false); + } + }; + + const handleImageInput = (event: React.ChangeEvent) => { + const files = Array.from(event.target.files ?? []); + event.target.value = ""; + void addImages(files); + }; + + useEffect( + () => () => { + abortControllerRef.current?.abort(); + draftVersionRef.current += 1; + }, + [], + ); const handleSend = useCallback( - async (text: string, historyOverride?: Array<{ role: "user" | "assistant"; content: string }>) => { + async ( + text: string, + historyOverride?: ReplayMessage[], + attachmentOverride?: ChatImageAttachment[], + generationModel?: string, + ) => { const trimmed = text.trim(); - if (!trimmed || !selectedModel || isStreaming) return; - const model = selectedModel; + const images = attachmentOverride ?? attachments; + const requestMode = generationModel ? "image" : mode; + const model = generationModel ?? selectedModel; + if (!hasMessageContent(trimmed, images, requestMode) || !model) return; + if (isStreaming || abortControllerRef.current || readImagesRef.current) return; setInputText(""); let convId = activeConversationId; @@ -145,11 +197,21 @@ export default function ChatConversationPage() { setConversationIdInUrl(convId); } - appendMessage(convId, { role: "user", content: trimmed }); - appendMessage(convId, { role: "assistant", content: "" }); + const imageGeneration = requestMode === "image" ? { model, prompt: trimmed } : undefined; + const userMessage = { + role: "user" as const, + content: trimmed, + attachments: images.length > 0 ? images : undefined, + imageGeneration, + }; + appendMessage(convId, userMessage); + appendMessage(convId, { role: "assistant", content: "", imageGeneration }); + setAttachments([]); + setAttachmentError(null); setIsStreaming(true); abortControllerRef.current = new AbortController(); + requestConversationRef.current = convId; // When historyOverride is set (edit / retry), the existing server-side // session chain covers messages that were just truncated and is no longer @@ -167,86 +229,33 @@ export default function ChatConversationPage() { // On the very first turn (no session yet), we send the full history. const previousResponseId = historyOverride ? null : responsesSessionId; - const history: Array<{ role: "user" | "assistant"; content: string }> = historyOverride - ? [...historyOverride, { role: "user" as const, content: trimmed }] - : previousResponseId - ? [{ role: "user" as const, content: trimmed }] - : [ - // Explicitly filter to only user/assistant roles — tool messages - // lack a required tool_call_id and would cause API errors. - ...(activeConversation?.messages ?? []) - .filter( - (m): m is typeof m & { role: "user" | "assistant" } => m.role === "user" || m.role === "assistant", - ) - .map((m) => ({ role: m.role, content: m.content })), - { role: "user" as const, content: trimmed }, - ]; - - let accumulatedContent = ""; - let accumulatedReasoning = ""; - // MCP events accumulated locally so we can persist them to the message - // without relying on component state (which would cause stale closures). - const accumulatedMCPEvents: MCPEvent[] = []; - // Track clean completion so partial events are not shown on error/abort. - let streamCompletedCleanly = false; + const history: ReplayMessage[] = [ + ...(historyOverride ?? (previousResponseId ? [] : activeConversation?.messages ?? [])), + userMessage, + ]; + const requestController = abortControllerRef.current; + const request = { + model, + mode: requestMode, + prompt: trimmed, + accessToken, + signal: requestController.signal, + messages: history, + previousResponseId, + mcpServers: selectedMCPServers, + onResponseId: setResponsesSessionId, + onUpdate: (updates: AssistantMessageUpdate) => updateLastAssistantMessage(convId, updates), + }; try { - await makeOpenAIResponsesRequest( - history, - (_role: string, chunk: string) => { - accumulatedContent += chunk; - updateLastAssistantMessage(convId!, { content: accumulatedContent }); - }, - model, - accessToken, - undefined, // tags - abortControllerRef.current.signal, - (rc: string) => { - accumulatedReasoning += rc; - updateLastAssistantMessage(convId!, { reasoningContent: accumulatedReasoning }); - }, - (timeToFirstToken: number) => updateLastAssistantMessage(convId!, { timeToFirstToken }), - (usage: TokenUsage) => updateLastAssistantMessage(convId!, { usage }), - undefined, - undefined, - undefined, - undefined, - selectedMCPServers.length > 0 ? selectedMCPServers : undefined, - previousResponseId, - (id: string) => setResponsesSessionId(id), - (event: MCPEvent) => { - // Accumulate locally only — persisted once in finally to avoid - // one full localStorage write per MCP event during streaming. - accumulatedMCPEvents.push(event); - }, - undefined, - undefined, - undefined, - undefined, - undefined, - undefined, - true, - (totalLatency: number) => updateLastAssistantMessage(convId!, { totalLatency }), - ); - streamCompletedCleanly = true; - } catch (err: unknown) { - if (err instanceof Error && err.name === "AbortError") { - updateLastAssistantMessage(convId!, { - content: accumulatedContent + " [stopped]", - }); - } else { - updateLastAssistantMessage(convId!, { - content: "[Something went wrong. The partial response has been saved.]", - }); - } + void recordChatQuery(userId); + await runChatRequest(request); } finally { - // Only persist MCP events on clean completion — partial events from an - // aborted or errored turn would show incomplete tool calls to the user. - if (accumulatedMCPEvents.length > 0 && streamCompletedCleanly) { - updateLastAssistantMessage(convId!, { mcpEvents: accumulatedMCPEvents }); + if (abortControllerRef.current === requestController) { + setIsStreaming(false); + abortControllerRef.current = null; + requestConversationRef.current = null; } - setIsStreaming(false); - abortControllerRef.current = null; } }, [ @@ -254,7 +263,10 @@ export default function ChatConversationPage() { activeConversation, selectedModel, selectedMCPServers, + mode, + attachments, accessToken, + userId, createConversation, appendMessage, updateLastAssistantMessage, @@ -271,34 +283,25 @@ export default function ChatConversationPage() { const handleEditAndResend = useCallback( (messageId: string, newContent: string) => { if (!activeConversationId || isStreaming) return; + if (abortControllerRef.current || readImagesRef.current) return; // Compute the truncated history synchronously before the async state update lands, // so handleSend receives the correct pre-edit context rather than the stale closure value. const msgs = activeConversation?.messages ?? []; const idx = msgs.findIndex((m) => m.id === messageId); - const priorMessages = (idx === -1 ? msgs : msgs.slice(0, idx)) - .filter((m) => m.role === "user" || m.role === "assistant") - .map((m) => ({ role: m.role as "user" | "assistant", content: m.content })); + if (idx < 0) return; + const original = msgs[idx]; + if (!original.imageGeneration && (mode !== "chat" || !selectedModel)) { + setAttachmentError("Choose a chat model before resending this message."); + return; + } + if (!newContent.trim() && !original.attachments?.length) return; + const priorMessages = msgs.slice(0, idx); truncateFromMessage(activeConversationId, messageId); - handleSend(newContent, priorMessages); + handleSend(newContent, priorMessages, original.attachments ?? [], original.imageGeneration?.model); }, - [activeConversationId, isStreaming, activeConversation, truncateFromMessage, handleSend], + [activeConversationId, isStreaming, activeConversation, truncateFromMessage, handleSend, mode, selectedModel], ); - const handleKeyDown = (e: React.KeyboardEvent) => { - if (e.key === "Enter" && !e.shiftKey) { - e.preventDefault(); - handleSend(inputText); - } - }; - - // Auto-resize textarea - useEffect(() => { - const ta = textareaRef.current; - if (!ta) return; - ta.style.height = "auto"; - ta.style.height = `${Math.min(ta.scrollHeight, 180)}px`; - }, [inputText]); - // Track scroll position to show/hide scroll-to-bottom button useEffect(() => { const el = messagesScrollRef.current; @@ -345,178 +348,165 @@ export default function ChatConversationPage() { }, [activeConversation?.messages]); const showBlankState = !activeConversation || activeConversation.messages.length === 0; - const displayName = userEmail?.split("@")[0] ?? userId ?? ""; - const greeting = displayName ? `${getGreeting()}, ${displayName}` : getGreeting(); - // Filtered models: selected one floats to the top, then alphabetical - const filteredModels = ( - modelSearchText ? models.filter((m) => m.toLowerCase().includes(modelSearchText.toLowerCase())) : models - ).sort((a, b) => { - if (a === selectedModel) return -1; - if (b === selectedModel) return 1; - return 0; - }); - - const modelSelectorContent = ( -
-
- setModelSearchText(e.target.value)} - placeholder="Search models..." - className="h-8 text-[13px]" - /> -
- - {filteredModels.map((m) => { - const checked = m === selectedModel; - const provider = getProviderFromModelName(m); - const { logo } = provider ? getProviderLogoAndName(provider) : { logo: "" }; - return ( - - ); - })} - -
- ); - - const modelSelectorTrigger = isLoadingModels ? ( - - ) : ( - { - setModelSelectorOpen(open); - if (!open) setModelSearchText(""); - }} - > - - {selectedModel ? ( + const inputBar = (inConversation: boolean) => { + const preparing = isLoadingModels || isReadingImages; + const submitDisabled = !hasMessageContent(inputText, attachments, mode) || !selectedModel || preparing; + return ( +
{ + if (mode === "chat" && event.dataTransfer.types.includes("Files")) event.preventDefault(); + }} + onDrop={(event) => { + if (mode === "chat" && event.dataTransfer.files.length) { + event.preventDefault(); + void addImages(Array.from(event.dataTransfer.files)); + } + }} + > + void handleSend(inputText)} + onCancel={handleStop} + placeholder={ + mode === "image" + ? "Describe the image you want to create…" + : inConversation + ? "Send a message..." + : "Describe a change, paste an error, or ask a question…" + } + isLoading={isStreaming} + submitLabel={mode === "image" ? "Generate image" : "Send"} + submitDisabled={submitDisabled} + onPaste={(event) => { + const images = Array.from(event.clipboardData.files).filter((file) => file.type.startsWith("image/")); + if (mode === "chat" && images.length > 0) { + event.preventDefault(); + void addImages(images); + } + }} + attachmentPreview={ + attachments.length > 0 ? ( +
+ {attachments.map((image) => ( +
+ {image.name} + +

{image.name}

+
+ ))} +
+ ) : undefined + } + modelPicker={ + + } + tools={ + mode === "chat" ? ( <> - {(() => { - const provider = getProviderFromModelName(selectedModel); - const { logo } = provider ? getProviderLogoAndName(provider) : { logo: "" }; - return logo ? ( - { - (e.currentTarget as HTMLImageElement).style.display = "none"; - }} + + + + + + Tools + {selectedMCPServers.length > 0 && ( + {selectedMCPServers.length} + )} + + } + /> + + - ) : null; - })()} - {selectedModel} + + ) : ( - Select model - )} - - - } - /> - - {modelSelectorContent} - - - ); - - const inputBar = (inConversation: boolean) => ( -
-