Merge branch 'main' into litellm_oss_staging_03_19_2026

litellm/proxy/_types.py

@@ -1672,6 +1672,7 @@ class NewTeamRequest(TeamBase):
         int
     ] = None  # allow user to set TPM limit for all team members
     team_member_key_duration: Optional[str] = None  # e.g. "1d", "1w", "1m"
+    team_member_budget_duration: Optional[str] = None  # e.g. "30d", "1mo"
     allowed_vector_store_indexes: Optional[List[AllowedVectorStoreIndexItem]] = None
     enforced_batch_output_expires_after: Optional[dict] = None
     enforced_file_expires_after: Optional[dict] = None

litellm/proxy/discovery_endpoints/ui_discovery_endpoints.py

@@ -27,10 +27,17 @@ async def get_ui_config():
     admin_ui_disabled = os.getenv("DISABLE_ADMIN_UI", "false").lower() == "true"
 
     sso_configured = _has_user_setup_sso()
+
+    from litellm.proxy.proxy_server import proxy_config
+
+    is_control_plane = len(proxy_config.worker_registry) > 0
+
     return UiDiscoveryEndpoints(
         server_root_path=get_server_root_path(),
         proxy_base_url=get_proxy_base_url(),
         auto_redirect_to_sso=sso_configured and auto_redirect_ui_login_to_sso,
         admin_ui_disabled=admin_ui_disabled,
         sso_configured=sso_configured,
+        is_control_plane=is_control_plane,
+        workers=proxy_config.worker_registry if is_control_plane else [],
     )

litellm/proxy/management_endpoints/mcp_management_endpoints.py

@@ -1214,17 +1214,9 @@ if MCP_AVAILABLE:
                     "error": "User does not have permission to create mcp servers. You can only create mcp servers if you are a PROXY_ADMIN."
                 },
             )
-        elif payload.server_id is not None:
-            # fail if the mcp server with id already exists
-            mcp_server = await get_mcp_server(prisma_client, payload.server_id)
-            if mcp_server is not None:
-                raise HTTPException(
-                    status_code=status.HTTP_400_BAD_REQUEST,
-                    detail={
-                        "error": f"MCP Server with id {payload.server_id} already exists. Cannot create another."
-                    },
-                )
-        elif (
+
+        # Block reserved special server IDs
+        if (
             SpecialMCPServerName.all_team_servers == payload.server_id
             or SpecialMCPServerName.all_proxy_servers == payload.server_id
         ):
@@ -1235,6 +1227,17 @@ if MCP_AVAILABLE:
                 },
             )
 
+        if payload.server_id is not None:
+            # fail if the mcp server with id already exists
+            mcp_server = await get_mcp_server(prisma_client, payload.server_id)
+            if mcp_server is not None:
+                raise HTTPException(
+                    status_code=status.HTTP_400_BAD_REQUEST,
+                    detail={
+                        "error": f"MCP Server with id {payload.server_id} already exists. Cannot create another."
+                    },
+                )
+
         # TODO: audit log for create
 
         # Admin-created servers are always active — clear any submission lifecycle

litellm/proxy/management_endpoints/team_endpoints.py

@@ -724,6 +724,7 @@ async def new_team(  # noqa: PLR0915
     - disable_global_guardrails: Optional[bool] - Whether to disable global guardrails for the key.
     - object_permission: Optional[LiteLLM_ObjectPermissionBase] - team-specific object permission. Example - {"vector_stores": ["vector_store_1", "vector_store_2"], "agents": ["agent_1", "agent_2"], "agent_access_groups": ["dev_group"]}. IF null or {} then no object permission.
     - team_member_budget: Optional[float] - The maximum budget allocated to an individual team member.
+    - team_member_budget_duration: Optional[str] - The duration of the budget for the team member. Doc [here](https://docs.litellm.ai/docs/proxy/team_budgets)
     - team_member_rpm_limit: Optional[int] - The RPM (Requests Per Minute) limit for individual team members.
     - team_member_tpm_limit: Optional[int] - The TPM (Tokens Per Minute) limit for individual team members.
     - team_member_key_duration: Optional[str] - The duration for a team member's key. e.g. "1d", "1w", "1mo"
@@ -934,6 +935,7 @@ async def new_team(  # noqa: PLR0915
             team_member_budget=data.team_member_budget,
             team_member_rpm_limit=data.team_member_rpm_limit,
             team_member_tpm_limit=data.team_member_tpm_limit,
+            team_member_budget_duration=data.team_member_budget_duration,
         ):
             data_json = await TeamMemberBudgetHandler.create_team_member_budget_table(
                 data=data,
@@ -942,6 +944,7 @@ async def new_team(  # noqa: PLR0915
                 team_member_budget=data.team_member_budget,
                 team_member_rpm_limit=data.team_member_rpm_limit,
                 team_member_tpm_limit=data.team_member_tpm_limit,
+                team_member_budget_duration=data.team_member_budget_duration,
             )
 
         ## ADD TO TEAM TABLE

litellm/proxy/management_endpoints/ui_sso.py

@@ -16,6 +16,7 @@ import os
 import secrets
 from copy import deepcopy
 from typing import TYPE_CHECKING, Any, Dict, List, Literal, Optional, Tuple, Union, cast
+from urllib.parse import urlencode, urlparse
 
 if TYPE_CHECKING:
     import httpx
@@ -301,6 +302,7 @@ async def google_login(
     source: Optional[str] = None,
     key: Optional[str] = None,
     existing_key: Optional[str] = None,
+    return_to: Optional[str] = None,
 ):  # noqa: PLR0915
     """
     Create Proxy API Keys using Google Workspace SSO. Requires setting PROXY_BASE_URL in .env
@@ -394,13 +396,23 @@ async def google_login(
         is True
     ):
         verbose_proxy_logger.info(f"Redirecting to SSO login for {redirect_url}")
-        return await SSOAuthenticationHandler.get_sso_login_redirect(
+        sso_redirect = await SSOAuthenticationHandler.get_sso_login_redirect(
             redirect_url=redirect_url,
             microsoft_client_id=microsoft_client_id,
             google_client_id=google_client_id,
             generic_client_id=generic_client_id,
             state=cli_state,
         )
+        if return_to is not None and sso_redirect is not None:
+            SSOAuthenticationHandler._validate_return_to(return_to)
+            sso_redirect.set_cookie(
+                key="litellm_cp_return_to",
+                value=return_to,
+                max_age=600,
+                httponly=True,
+                samesite="lax",
+            )
+        return sso_redirect
     elif ui_username is not None:
         # No Google, Microsoft SSO
         # Use UI Credentials set in .env
@@ -1312,12 +1324,17 @@ async def auth_callback(request: Request, state: Optional[str] = None):  # noqa:
             request=request, key=key_id, existing_key=existing_key, result=result
         )
 
+    # Control-plane cross-origin: read return_to from cookie.
+    # Starlette's cookie_parser already handles RFC 2109 unquoting.
+    cp_return_to: Optional[str] = request.cookies.get("litellm_cp_return_to")
+
     return await SSOAuthenticationHandler.get_redirect_response_from_openid(
         result=result,
         request=request,
         received_response=received_response,
         generic_client_id=generic_client_id,
         ui_access_mode=ui_access_mode,
+        return_to=cp_return_to,
     )
 
 
@@ -1760,6 +1777,38 @@ class SSOAuthenticationHandler:
     Handler for SSO Authentication across all SSO providers
     """
 
+    @staticmethod
+    def _validate_return_to(return_to: str) -> None:
+        """
+        Validate that return_to matches the configured control_plane_url origin.
+
+        Raises HTTPException(400) if:
+        - control_plane_url is not configured in general_settings
+        - return_to origin does not match control_plane_url origin
+        """
+        from litellm.proxy.proxy_server import general_settings
+
+        control_plane_url = general_settings.get("control_plane_url")
+        if control_plane_url is None:
+            raise HTTPException(
+                status_code=400,
+                detail="return_to is not allowed: control_plane_url is not configured",
+            )
+
+        def _origin(url: str) -> tuple:
+            parsed = urlparse(url)
+            scheme = (parsed.scheme or "").lower()
+            hostname = (parsed.hostname or "").lower()
+            default_port = 443 if scheme == "https" else 80
+            port = parsed.port if parsed.port is not None else default_port
+            return (scheme, hostname, port)
+
+        if _origin(return_to) != _origin(control_plane_url):
+            raise HTTPException(
+                status_code=400,
+                detail="return_to does not match the configured control_plane_url",
+            )
+
     @staticmethod
     async def get_sso_login_redirect(
         redirect_url: str,
@@ -2358,6 +2407,7 @@ class SSOAuthenticationHandler:
         received_response: Optional[dict] = None,
         generic_client_id: Optional[str] = None,
         ui_access_mode: Optional[Dict] = None,
+        return_to: Optional[str] = None,
     ) -> RedirectResponse:
         import jwt
 
@@ -2367,6 +2417,7 @@ class SSOAuthenticationHandler:
             master_key,
             premium_user,
             proxy_logging_obj,
+            redis_usage_cache,
             user_api_key_cache,
             user_custom_sso,
         )
@@ -2534,6 +2585,36 @@ class SSOAuthenticationHandler:
             master_key or "",
             algorithm="HS256",
         )
+
+        # Control-plane cross-origin: store JWT behind a single-use opaque
+        # code (60s TTL) so the token never appears in browser history / logs.
+        # The control plane redeems it via POST /v3/login/exchange.
+        if return_to is not None:
+            SSOAuthenticationHandler._validate_return_to(return_to)
+
+            code = secrets.token_urlsafe(32)
+            cache_key = f"login_code:{code}"
+            cache_value = {"token": jwt_token, "redirect_url": return_to}
+            if redis_usage_cache is not None:
+                await redis_usage_cache.async_set_cache(
+                    key=cache_key, value=cache_value, ttl=60
+                )
+            else:
+                await user_api_key_cache.async_set_cache(
+                    key=cache_key, value=cache_value, ttl=60
+                )
+
+            separator = "&" if "?" in return_to else "?"
+            redirect_url = (
+                return_to + separator + urlencode({"login": "success", "code": code})
+            )
+            verbose_proxy_logger.info(
+                "Cross-origin SSO: redirecting to control plane with login code"
+            )
+            redirect_response = RedirectResponse(url=redirect_url, status_code=303)
+            redirect_response.delete_cookie("litellm_cp_return_to")
+            return redirect_response
+
         if user_id is not None and isinstance(user_id, str):
             litellm_dashboard_ui += "?login=success"
         verbose_proxy_logger.info(f"Redirecting to {litellm_dashboard_ui}")

litellm/proxy/proxy_server.py

@@ -541,6 +541,7 @@ from litellm.types.llms.anthropic import (
     AnthropicResponseUsageBlock,
 )
 from litellm.types.llms.openai import HttpxBinaryResponseContent
+from litellm.types.proxy.control_plane_endpoints import WorkerRegistryEntry
 from litellm.types.proxy.management_endpoints.model_management_endpoints import (
     ModelGroupInfoProxy,
 )
@@ -1546,6 +1547,7 @@ user_custom_key_generate = None
 # Sentinel: prevents PKCE-no-Redis advisory from re-logging on config hot-reload.
 # Tests that need to reset it can patch 'litellm.proxy.proxy_server._pkce_no_redis_warning_emitted'.
 _pkce_no_redis_warning_emitted: bool = False
+_cp_no_redis_warning_emitted: bool = False
 user_custom_sso = None
 user_custom_ui_sso_sign_in_handler = None
 use_background_health_checks = None
@@ -2295,6 +2297,7 @@ class ProxyConfig:
         self.config: Dict[str, Any] = {}
         self._last_semantic_filter_config: Optional[Dict[str, Any]] = None
         self._last_hashicorp_vault_config: Optional[Dict[str, Any]] = None
+        self.worker_registry: List["WorkerRegistryEntry"] = []
 
     def is_yaml(self, config_file_path: str) -> bool:
         if not os.path.isfile(config_file_path):
@@ -3095,6 +3098,21 @@ class ProxyConfig:
                         "Set PKCE_STRICT_CACHE_MISS=true to fail fast with a 401 on cache misses "
                         "instead of continuing without a code_verifier."
                     )
+
+            ### CONTROL PLANE CODE-EXCHANGE PREREQUISITE CHECK ###
+            cp_url = general_settings.get("control_plane_url")
+            if cp_url and redis_usage_cache is None:
+                global _cp_no_redis_warning_emitted
+                if not _cp_no_redis_warning_emitted:
+                    _cp_no_redis_warning_emitted = True
+                    verbose_proxy_logger.warning(
+                        "control_plane_url is configured but Redis is not configured for LiteLLM caching. "
+                        "Login codes (SSO and /v3/login) will not be shared across instances — "
+                        "the /v3/login/exchange call may land on a different pod and fail with 401. "
+                        "Configure Redis via the 'cache' section in your proxy config, "
+                        "or ensure sticky sessions for single-instance deployments."
+                    )
+
             ### STORE MODEL IN DB ### feature flag for `/model/new`
             store_model_in_db = general_settings.get("store_model_in_db", False)
             if store_model_in_db is None:
@@ -3385,7 +3403,15 @@ class ProxyConfig:
             litellm.vector_store_registry.load_vector_stores_from_config(
                 vector_store_registry_config
             )
-        pass
+
+        ## WORKER REGISTRY (Control Plane)
+        worker_registry_config = config.get("worker_registry", None)
+        if worker_registry_config:
+            self.worker_registry = [
+                WorkerRegistryEntry(**e) for e in worker_registry_config
+            ]
+        else:
+            self.worker_registry = []
 
     async def _init_policy_engine(
         self,
@@ -11095,6 +11121,165 @@ async def login_v2(request: Request):  # noqa: PLR0915
             )
 
 
+@router.post(
+    "/v3/login", include_in_schema=False
+)  # control-plane login — always returns token in body for cross-origin use
+async def login_v3(request: Request):  # noqa: PLR0915
+    global premium_user, general_settings, master_key
+    from litellm.proxy.auth.login_utils import authenticate_user, create_ui_token_object
+    from litellm.proxy.utils import get_custom_url
+
+    try:
+        if not general_settings.get("control_plane_url"):
+            raise ProxyException(
+                message="/v3/login is only available on workers with control_plane_url configured",
+                type=ProxyErrorTypes.not_found_error,
+                param="control_plane_url",
+                code=status.HTTP_404_NOT_FOUND,
+            )
+
+        body = await request.json()
+        username = str(body.get("username"))
+        password = str(body.get("password"))
+
+        login_result = await authenticate_user(
+            username=username,
+            password=password,
+            master_key=master_key,
+            prisma_client=prisma_client,
+        )
+
+        returned_ui_token_object = create_ui_token_object(
+            login_result=login_result,
+            general_settings=general_settings,
+            premium_user=premium_user,
+        )
+
+        import jwt
+
+        jwt_token = jwt.encode(
+            cast(dict, returned_ui_token_object),
+            cast(str, master_key),
+            algorithm="HS256",
+        )
+
+        litellm_dashboard_ui = get_custom_url(str(request.base_url))
+        if litellm_dashboard_ui.endswith("/"):
+            litellm_dashboard_ui += "ui/"
+        else:
+            litellm_dashboard_ui += "/ui/"
+        litellm_dashboard_ui += "?login=success"
+
+        # Store JWT behind a single-use opaque code (60s TTL)
+        code = secrets.token_urlsafe(32)
+        cache_key = f"login_code:{code}"
+        cache_value = {"token": jwt_token, "redirect_url": litellm_dashboard_ui}
+        if redis_usage_cache is not None:
+            await redis_usage_cache.async_set_cache(
+                key=cache_key, value=cache_value, ttl=60
+            )
+        else:
+            await user_api_key_cache.async_set_cache(
+                key=cache_key, value=cache_value, ttl=60
+            )
+
+        return JSONResponse(
+            content={"code": code, "expires_in": 60},
+            status_code=status.HTTP_200_OK,
+        )
+    except Exception as e:
+        verbose_proxy_logger.exception(
+            "litellm.proxy.proxy_server.login_v3(): Exception occurred - {}".format(
+                str(e)
+            )
+        )
+        if isinstance(e, ProxyException):
+            raise e
+        elif isinstance(e, HTTPException):
+            raise ProxyException(
+                message=getattr(e, "detail", str(e)),
+                type=ProxyErrorTypes.auth_error,
+                param=getattr(e, "param", "None"),
+                code=getattr(e, "status_code", status.HTTP_500_INTERNAL_SERVER_ERROR),
+            )
+        else:
+            error_msg = f"{str(e)}"
+            raise ProxyException(
+                message=error_msg,
+                type=ProxyErrorTypes.auth_error,
+                param="None",
+                code=status.HTTP_500_INTERNAL_SERVER_ERROR,
+            )
+
+
+@router.post(
+    "/v3/login/exchange", include_in_schema=False
+)  # exchange single-use opaque code for JWT
+async def login_v3_exchange(request: Request):
+    try:
+        if not general_settings.get("control_plane_url"):
+            raise ProxyException(
+                message="/v3/login/exchange is only available on workers with control_plane_url configured",
+                type=ProxyErrorTypes.not_found_error,
+                param="control_plane_url",
+                code=status.HTTP_404_NOT_FOUND,
+            )
+
+        body = await request.json()
+        code = body.get("code")
+        if not code:
+            raise ProxyException(
+                message="Missing 'code' parameter",
+                type=ProxyErrorTypes.auth_error,
+                param="code",
+                code=status.HTTP_400_BAD_REQUEST,
+            )
+
+        cache_key = f"login_code:{code}"
+        if redis_usage_cache is not None:
+            cached_data = await redis_usage_cache.async_get_cache(key=cache_key)
+        else:
+            cached_data = await user_api_key_cache.async_get_cache(key=cache_key)
+
+        if not cached_data or not isinstance(cached_data, dict):
+            raise ProxyException(
+                message="Invalid or expired login code",
+                type=ProxyErrorTypes.auth_error,
+                param="code",
+                code=status.HTTP_401_UNAUTHORIZED,
+            )
+
+        # Single-use: delete immediately
+        if redis_usage_cache is not None:
+            await redis_usage_cache.async_delete_cache(key=cache_key)
+        else:
+            await user_api_key_cache.async_delete_cache(key=cache_key)
+
+        json_response = JSONResponse(
+            content={
+                "token": cached_data["token"],
+                "redirect_url": cached_data["redirect_url"],
+            },
+            status_code=status.HTTP_200_OK,
+        )
+        json_response.set_cookie(key="token", value=cached_data["token"])
+        return json_response
+    except ProxyException:
+        raise
+    except Exception as e:
+        verbose_proxy_logger.exception(
+            "litellm.proxy.proxy_server.login_v3_exchange(): Exception occurred - {}".format(
+                str(e)
+            )
+        )
+        raise ProxyException(
+            message=str(e),
+            type=ProxyErrorTypes.auth_error,
+            param="None",
+            code=status.HTTP_500_INTERNAL_SERVER_ERROR,
+        )
+
+
 @app.get("/onboarding/get_token", include_in_schema=False)
 async def onboarding(invite_link: str, request: Request):
     """

litellm/types/proxy/control_plane_endpoints.py

@@ -0,0 +1,14 @@
+from pydantic import BaseModel, field_validator
+
+
+class WorkerRegistryEntry(BaseModel):
+    worker_id: str
+    name: str
+    url: str
+
+    @field_validator("url")
+    @classmethod
+    def url_must_be_http(cls, v: str) -> str:
+        if not v.startswith(("http://", "https://")):
+            raise ValueError("Worker URL must start with http:// or https://")
+        return v

litellm/types/proxy/discovery_endpoints/ui_discovery_endpoints.py

@@ -1,7 +1,9 @@
-from typing import Optional
+from typing import List, Optional
 
 from pydantic import BaseModel
 
+from litellm.types.proxy.control_plane_endpoints import WorkerRegistryEntry
+
 
 class UiDiscoveryEndpoints(BaseModel):
     server_root_path: str
@@ -9,3 +11,5 @@ class UiDiscoveryEndpoints(BaseModel):
     auto_redirect_to_sso: bool
     admin_ui_disabled: bool
     sso_configured: bool
+    is_control_plane: bool = False
+    workers: List[WorkerRegistryEntry] = []

tests/test_litellm/llms/anthropic/test_anthropic_files_and_batches.py

@@ -336,7 +336,7 @@ class TestAnthropicFilesHandler:
             "extra_body": None
         }
 
-        with patch.object(handler.anthropic_model_info, "get_api_key", return_value=None):
+        with patch.object(handler.anthropic_model_info, "get_auth_header", return_value=None):
             with pytest.raises(ValueError, match="Missing Anthropic API Key"):
                 await handler.afile_content(
                     file_content_request=file_content_request,

tests/test_litellm/proxy/discovery_endpoints/test_ui_discovery_endpoints.py

@@ -1,6 +1,6 @@
 import os
 import sys
-from unittest.mock import patch
+from unittest.mock import MagicMock, patch
 
 import pytest
 from fastapi import FastAPI
@@ -11,6 +11,7 @@ sys.path.insert(
 )
 
 from litellm.proxy.discovery_endpoints.ui_discovery_endpoints import router
+from litellm.types.proxy.control_plane_endpoints import WorkerRegistryEntry
 
 
 def test_ui_discovery_endpoints_with_defaults():
@@ -245,9 +246,9 @@ def test_ui_discovery_endpoints_with_admin_ui_enabled():
          patch("litellm.proxy.utils.get_proxy_base_url", return_value=None), \
          patch("litellm.proxy.auth.auth_utils._has_user_setup_sso", return_value=False), \
          patch.dict(os.environ, {"DISABLE_ADMIN_UI": "false"}, clear=False):
-        
+
         response = client.get("/.well-known/litellm-ui-config")
-        
+
         assert response.status_code == 200
         data = response.json()
         assert data["server_root_path"] == "/"
@@ -256,3 +257,53 @@ def test_ui_discovery_endpoints_with_admin_ui_enabled():
         assert data["admin_ui_disabled"] is False
         assert data["sso_configured"] is False
 
+
+def test_ui_discovery_endpoints_is_control_plane_true_when_workers_configured():
+    app = FastAPI()
+    app.include_router(router)
+    client = TestClient(app)
+
+    mock_config = MagicMock()
+    mock_config.worker_registry = [
+        WorkerRegistryEntry(
+            worker_id="team-a", name="Team A", url="https://worker-1:4001"
+        ),
+    ]
+
+    with patch("litellm.proxy.utils.get_server_root_path", return_value="/"), \
+         patch("litellm.proxy.utils.get_proxy_base_url", return_value=None), \
+         patch("litellm.proxy.auth.auth_utils._has_user_setup_sso", return_value=False), \
+         patch("litellm.proxy.proxy_server.proxy_config", mock_config), \
+         patch.dict(os.environ, {"DISABLE_ADMIN_UI": "false"}, clear=False):
+
+        response = client.get("/.well-known/litellm-ui-config")
+
+        assert response.status_code == 200
+        data = response.json()
+        assert data["is_control_plane"] is True
+        assert len(data["workers"]) == 1
+        assert data["workers"][0]["worker_id"] == "team-a"
+        assert data["workers"][0]["name"] == "Team A"
+        assert data["workers"][0]["url"] == "https://worker-1:4001"
+
+
+def test_ui_discovery_endpoints_is_control_plane_false_when_no_workers():
+    app = FastAPI()
+    app.include_router(router)
+    client = TestClient(app)
+
+    mock_config = MagicMock()
+    mock_config.worker_registry = []
+
+    with patch("litellm.proxy.utils.get_server_root_path", return_value="/"), \
+         patch("litellm.proxy.utils.get_proxy_base_url", return_value=None), \
+         patch("litellm.proxy.auth.auth_utils._has_user_setup_sso", return_value=False), \
+         patch("litellm.proxy.proxy_server.proxy_config", mock_config), \
+         patch.dict(os.environ, {"DISABLE_ADMIN_UI": "false"}, clear=False):
+
+        response = client.get("/.well-known/litellm-ui-config")
+
+        assert response.status_code == 200
+        data = response.json()
+        assert data["is_control_plane"] is False
+        assert data["workers"] == []

tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py

@@ -6441,3 +6441,53 @@ async def test_list_team_v1_batches_key_queries():
         assert result[0].keys == [key1, key2]
         assert result[1].team_id == "team-2"
         assert result[1].keys == [key3]
+
+
+def test_new_team_request_accepts_team_member_budget_duration():
+    """Test that NewTeamRequest does not silently drop team_member_budget_duration."""
+    from litellm.proxy._types import NewTeamRequest
+
+    request = NewTeamRequest(
+        team_member_budget=20.0,
+        team_member_budget_duration="30d",
+    )
+    assert request.team_member_budget == 20.0
+    assert request.team_member_budget_duration == "30d"
+
+
+@pytest.mark.asyncio
+async def test_create_team_member_budget_table_with_duration():
+    """Verify that create_team_member_budget_table passes budget_duration
+    through to the new_budget call when team_member_budget_duration is provided."""
+    from litellm.proxy._types import NewTeamRequest, UserAPIKeyAuth, LitellmUserRoles
+    from litellm.proxy.management_endpoints.team_endpoints import TeamMemberBudgetHandler
+
+    mock_budget_response = MagicMock(budget_id="budget-abc")
+    mock_admin = UserAPIKeyAuth(
+        user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN
+    )
+
+    data = NewTeamRequest(
+        team_alias="test-team",
+        team_member_budget=20.0,
+        team_member_budget_duration="30d",
+    )
+
+    with patch(
+        "litellm.proxy.management_endpoints.budget_management_endpoints.new_budget",
+        new_callable=AsyncMock,
+        return_value=mock_budget_response,
+    ) as mock_new_budget:
+        result = await TeamMemberBudgetHandler.create_team_member_budget_table(
+            data=data,
+            new_team_data_json={"metadata": None},
+            user_api_key_dict=mock_admin,
+            team_member_budget=20.0,
+            team_member_budget_duration="30d",
+        )
+
+        mock_new_budget.assert_awaited_once()
+        budget_request = mock_new_budget.call_args.kwargs["budget_obj"]
+        assert budget_request.budget_duration == "30d"
+        assert budget_request.max_budget == 20.0
+        assert result["metadata"]["team_member_budget_id"] == "budget-abc"

tests/test_litellm/proxy/management_endpoints/test_ui_sso.py

@@ -6,7 +6,7 @@ from unittest.mock import AsyncMock, MagicMock, patch
 
 import httpx
 import pytest
-from fastapi import Request
+from fastapi import HTTPException, Request
 
 from litellm._uuid import uuid
 
@@ -5160,3 +5160,99 @@ def test_generic_response_convertor_extra_attributes_missing_field(monkeypatch):
     assert result.extra_fields["missing_field"] is None
     assert result.extra_fields["another_missing"] is None
 
+
+class TestValidateReturnTo:
+    """Tests for SSOAuthenticationHandler._validate_return_to"""
+
+    def test_rejects_when_no_control_plane_url_configured(self, monkeypatch):
+        """return_to should be rejected if control_plane_url is not in general_settings."""
+        monkeypatch.setattr(
+            "litellm.proxy.proxy_server.general_settings", {}
+        )
+        with pytest.raises(HTTPException) as exc_info:
+            SSOAuthenticationHandler._validate_return_to("https://cp.example.com/ui")
+        assert exc_info.value.status_code == 400
+        assert "not configured" in exc_info.value.detail
+
+    def test_allows_matching_origin(self, monkeypatch):
+        """return_to matching the configured control_plane_url origin should pass."""
+        monkeypatch.setattr(
+            "litellm.proxy.proxy_server.general_settings",
+            {"control_plane_url": "https://cp.example.com"},
+        )
+        # Should not raise
+        SSOAuthenticationHandler._validate_return_to("https://cp.example.com/ui?page=models")
+
+    def test_allows_matching_origin_with_trailing_slash(self, monkeypatch):
+        """Trailing slash on control_plane_url should not affect origin comparison."""
+        monkeypatch.setattr(
+            "litellm.proxy.proxy_server.general_settings",
+            {"control_plane_url": "https://cp.example.com/"},
+        )
+        SSOAuthenticationHandler._validate_return_to("https://cp.example.com/ui")
+
+    def test_rejects_prefix_attack(self, monkeypatch):
+        """return_to like cp.example.com.evil.com must be rejected (not just prefix match)."""
+        monkeypatch.setattr(
+            "litellm.proxy.proxy_server.general_settings",
+            {"control_plane_url": "https://cp.example.com"},
+        )
+        with pytest.raises(HTTPException) as exc_info:
+            SSOAuthenticationHandler._validate_return_to("https://cp.example.com.evil.com/steal")
+        assert exc_info.value.status_code == 400
+
+    def test_rejects_different_origin(self, monkeypatch):
+        """return_to pointing to a completely different domain should be rejected."""
+        monkeypatch.setattr(
+            "litellm.proxy.proxy_server.general_settings",
+            {"control_plane_url": "https://cp.example.com"},
+        )
+        with pytest.raises(HTTPException) as exc_info:
+            SSOAuthenticationHandler._validate_return_to("https://evil.com/phish")
+        assert exc_info.value.status_code == 400
+
+    def test_case_insensitive_hostname(self, monkeypatch):
+        """Hostname comparison should be case-insensitive per RFC 3986."""
+        monkeypatch.setattr(
+            "litellm.proxy.proxy_server.general_settings",
+            {"control_plane_url": "https://CP.Example.COM"},
+        )
+        # Should not raise
+        SSOAuthenticationHandler._validate_return_to("https://cp.example.com/ui")
+
+    def test_rejects_scheme_mismatch(self, monkeypatch):
+        """http:// must be rejected when control_plane_url uses https://."""
+        monkeypatch.setattr(
+            "litellm.proxy.proxy_server.general_settings",
+            {"control_plane_url": "https://cp.example.com"},
+        )
+        with pytest.raises(HTTPException) as exc_info:
+            SSOAuthenticationHandler._validate_return_to("http://cp.example.com/ui")
+        assert exc_info.value.status_code == 400
+
+    def test_rejects_port_mismatch(self, monkeypatch):
+        """Non-default port must be rejected."""
+        monkeypatch.setattr(
+            "litellm.proxy.proxy_server.general_settings",
+            {"control_plane_url": "https://cp.example.com"},
+        )
+        with pytest.raises(HTTPException) as exc_info:
+            SSOAuthenticationHandler._validate_return_to("https://cp.example.com:8443/ui")
+        assert exc_info.value.status_code == 400
+
+    def test_allows_explicit_default_port(self, monkeypatch):
+        """https://host:443 should match https://host (default port normalisation)."""
+        monkeypatch.setattr(
+            "litellm.proxy.proxy_server.general_settings",
+            {"control_plane_url": "https://cp.example.com"},
+        )
+        SSOAuthenticationHandler._validate_return_to("https://cp.example.com:443/ui")
+
+    def test_allows_matching_custom_port(self, monkeypatch):
+        """Both sides on the same custom port should match."""
+        monkeypatch.setattr(
+            "litellm.proxy.proxy_server.general_settings",
+            {"control_plane_url": "https://cp.example.com:3000"},
+        )
+        SSOAuthenticationHandler._validate_return_to("https://cp.example.com:3000/ui")
+

tests/test_litellm/proxy/test_proxy_server.py

@@ -236,6 +236,217 @@ def test_login_v2_returns_json_on_invalid_json_body(monkeypatch):
     assert isinstance(data["error"], dict)
 
 
+def test_login_v3_rejected_without_control_plane_url(monkeypatch):
+    """v3/login returns 404 when control_plane_url is not configured."""
+    mock_prisma_client = MagicMock()
+    monkeypatch.setattr("litellm.proxy.proxy_server.master_key", "test-master-key")
+    monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", {})
+    monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client)
+
+    client = TestClient(app)
+    response = client.post(
+        "/v3/login",
+        json={"username": "alice", "password": "secret"},
+    )
+
+    assert response.status_code == 404
+    assert "control_plane_url" in response.json()["error"]["message"]
+
+
+def test_login_v3_returns_code(monkeypatch):
+    """v3/login returns an opaque code, not the JWT directly."""
+    mock_prisma_client = MagicMock()
+    monkeypatch.setattr(
+        "litellm.proxy.auth.login_utils.authenticate_user",
+        AsyncMock(return_value={"user_id": "test-user"}),
+    )
+    monkeypatch.setattr(
+        "litellm.proxy.auth.login_utils.create_ui_token_object",
+        MagicMock(return_value={"user_id": "test-user"}),
+    )
+    monkeypatch.setattr("jwt.encode", MagicMock(return_value="signed-token"))
+    monkeypatch.setattr("litellm.proxy.proxy_server.master_key", "test-master-key")
+    monkeypatch.setattr(
+        "litellm.proxy.proxy_server.general_settings",
+        {"control_plane_url": "https://cp.example.com"},
+    )
+    monkeypatch.setattr("litellm.proxy.proxy_server.premium_user", False)
+    monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client)
+    mock_config = MagicMock()
+    mock_config.worker_registry = []
+    monkeypatch.setattr("litellm.proxy.proxy_server.proxy_config", mock_config)
+    monkeypatch.setattr("litellm.proxy.utils.get_server_root_path", lambda: "")
+    monkeypatch.setattr("litellm.proxy.utils.get_proxy_base_url", lambda: None)
+
+    client = TestClient(app)
+    response = client.post(
+        "/v3/login",
+        json={"username": "alice", "password": "secret"},
+    )
+
+    assert response.status_code == 200
+    data = response.json()
+    assert "code" in data
+    assert data["expires_in"] == 60
+    assert "token" not in data
+
+
+def test_login_v3_exchange_happy_path(monkeypatch):
+    """Full flow: v3/login returns code, v3/login/exchange redeems it for JWT."""
+    mock_prisma_client = MagicMock()
+    monkeypatch.setattr(
+        "litellm.proxy.auth.login_utils.authenticate_user",
+        AsyncMock(return_value={"user_id": "test-user"}),
+    )
+    monkeypatch.setattr(
+        "litellm.proxy.auth.login_utils.create_ui_token_object",
+        MagicMock(return_value={"user_id": "test-user"}),
+    )
+    monkeypatch.setattr("jwt.encode", MagicMock(return_value="signed-token"))
+    monkeypatch.setattr("litellm.proxy.proxy_server.master_key", "test-master-key")
+    monkeypatch.setattr(
+        "litellm.proxy.proxy_server.general_settings",
+        {"control_plane_url": "https://cp.example.com"},
+    )
+    monkeypatch.setattr("litellm.proxy.proxy_server.premium_user", False)
+    monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client)
+    mock_config = MagicMock()
+    mock_config.worker_registry = []
+    monkeypatch.setattr("litellm.proxy.proxy_server.proxy_config", mock_config)
+    monkeypatch.setattr("litellm.proxy.utils.get_server_root_path", lambda: "")
+    monkeypatch.setattr("litellm.proxy.utils.get_proxy_base_url", lambda: None)
+
+    client = TestClient(app)
+
+    # Step 1: login — get code
+    login_response = client.post(
+        "/v3/login",
+        json={"username": "alice", "password": "secret"},
+    )
+    assert login_response.status_code == 200
+    code = login_response.json()["code"]
+
+    # Step 2: exchange — get JWT
+    exchange_response = client.post(
+        "/v3/login/exchange",
+        json={"code": code},
+    )
+    assert exchange_response.status_code == 200
+    exchange_data = exchange_response.json()
+    assert exchange_data["token"] == "signed-token"
+    assert "redirect_url" in exchange_data
+    assert exchange_response.cookies.get("token") == "signed-token"
+
+
+def test_login_v3_exchange_single_use(monkeypatch):
+    """Code can only be redeemed once."""
+    mock_prisma_client = MagicMock()
+    monkeypatch.setattr(
+        "litellm.proxy.auth.login_utils.authenticate_user",
+        AsyncMock(return_value={"user_id": "test-user"}),
+    )
+    monkeypatch.setattr(
+        "litellm.proxy.auth.login_utils.create_ui_token_object",
+        MagicMock(return_value={"user_id": "test-user"}),
+    )
+    monkeypatch.setattr("jwt.encode", MagicMock(return_value="signed-token"))
+    monkeypatch.setattr("litellm.proxy.proxy_server.master_key", "test-master-key")
+    monkeypatch.setattr(
+        "litellm.proxy.proxy_server.general_settings",
+        {"control_plane_url": "https://cp.example.com"},
+    )
+    monkeypatch.setattr("litellm.proxy.proxy_server.premium_user", False)
+    monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client)
+    mock_config = MagicMock()
+    mock_config.worker_registry = []
+    monkeypatch.setattr("litellm.proxy.proxy_server.proxy_config", mock_config)
+    monkeypatch.setattr("litellm.proxy.utils.get_server_root_path", lambda: "")
+    monkeypatch.setattr("litellm.proxy.utils.get_proxy_base_url", lambda: None)
+
+    client = TestClient(app)
+
+    login_response = client.post(
+        "/v3/login",
+        json={"username": "alice", "password": "secret"},
+    )
+    code = login_response.json()["code"]
+
+    # First exchange succeeds
+    first = client.post("/v3/login/exchange", json={"code": code})
+    assert first.status_code == 200
+
+    # Second exchange fails
+    second = client.post("/v3/login/exchange", json={"code": code})
+    assert second.status_code == 401
+
+
+def test_login_v3_exchange_invalid_code(monkeypatch):
+    """Random code returns 401."""
+    monkeypatch.setattr(
+        "litellm.proxy.proxy_server.general_settings",
+        {"control_plane_url": "https://cp.example.com"},
+    )
+    client = TestClient(app)
+    response = client.post(
+        "/v3/login/exchange",
+        json={"code": "nonexistent-code"},
+    )
+    assert response.status_code == 401
+
+
+def test_login_v3_exchange_rejected_without_control_plane_url(monkeypatch):
+    """v3/login/exchange returns 404 when control_plane_url is not configured."""
+    monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", {})
+
+    client = TestClient(app)
+    response = client.post(
+        "/v3/login/exchange",
+        json={"code": "some-code"},
+    )
+
+    assert response.status_code == 404
+    assert "control_plane_url" in response.json()["error"]["message"]
+
+
+def test_login_v3_returns_json_on_proxy_exception(monkeypatch):
+    """Test that /v3/login returns JSON error when ProxyException is raised"""
+    from litellm.proxy._types import ProxyErrorTypes, ProxyException
+
+    mock_prisma_client = MagicMock()
+    mock_authenticate_user = AsyncMock(
+        side_effect=ProxyException(
+            message="Invalid credentials",
+            type=ProxyErrorTypes.auth_error,
+            param="password",
+            code=401,
+        )
+    )
+
+    monkeypatch.setattr(
+        "litellm.proxy.auth.login_utils.authenticate_user",
+        mock_authenticate_user,
+    )
+    monkeypatch.setattr("litellm.proxy.proxy_server.master_key", "test-master-key")
+    monkeypatch.setattr(
+        "litellm.proxy.proxy_server.general_settings",
+        {"control_plane_url": "https://cp.example.com"},
+    )
+    monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client)
+
+    client = TestClient(app)
+    response = client.post(
+        "/v3/login",
+        json={"username": "alice", "password": "wrong"},
+    )
+
+    assert response.status_code == 401
+    assert response.headers["content-type"] == "application/json"
+    data = response.json()
+    assert "error" in data
+    assert data["error"]["message"] == "Invalid credentials"
+    assert data["error"]["type"] == "auth_error"
+
+
 def test_fallback_login_has_no_deprecation_banner(client_no_auth):
     response = client_no_auth.get("/fallback/login")
 

ui/litellm-dashboard/src/app/(dashboard)/api-reference/page.tsx

@@ -1,16 +1,10 @@
 "use client";
 
 import APIReferenceView from "@/app/(dashboard)/api-reference/APIReferenceView";
-import { useState } from "react";
-
-interface ProxySettings {
-  PROXY_BASE_URL: string;
-  PROXY_LOGOUT_URL: string;
-  LITELLM_UI_API_DOC_BASE_URL?: string | null;
-}
+import useProxySettings from "@/app/(dashboard)/hooks/proxySettings/useProxySettings";
 
 const APIReferencePage = () => {
-  const [proxySettings, setProxySettings] = useState<ProxySettings>({ PROXY_BASE_URL: "", PROXY_LOGOUT_URL: "" });
+  const proxySettings = useProxySettings();
 
   return <APIReferenceView proxySettings={proxySettings} />;
 };

ui/litellm-dashboard/src/app/(dashboard)/components/Sidebar2.tsx

@@ -195,7 +195,7 @@ const menuItems: MenuItemCfg[] = [
     icon: <UserOutlined style={{ fontSize: 18 }} />,
     roles: all_admin_roles,
   },
-  { key: "14", page: "api_ref", label: "API Reference", icon: <ApiOutlined style={{ fontSize: 18 }} /> },
+  { key: "14", page: "api-reference", label: "API Reference", icon: <ApiOutlined style={{ fontSize: 18 }} /> },
   {
     key: "16",
     page: "model-hub-table",

ui/litellm-dashboard/src/app/(dashboard)/hooks/common/queryKeysFactory.test.ts

@@ -0,0 +1,34 @@
+import { describe, it, expect } from "vitest";
+import { createQueryKeys } from "./queryKeysFactory";
+
+describe("createQueryKeys", () => {
+  const keys = createQueryKeys("books");
+
+  it("should return the resource name as the base key", () => {
+    expect(keys.all).toEqual(["books"]);
+  });
+
+  it("should generate a lists key", () => {
+    expect(keys.lists()).toEqual(["books", "list"]);
+  });
+
+  it("should generate a list key with params", () => {
+    expect(keys.list({ page: 1, limit: 10 })).toEqual([
+      "books",
+      "list",
+      { params: { page: 1, limit: 10 } },
+    ]);
+  });
+
+  it("should generate a list key with undefined params when none provided", () => {
+    expect(keys.list()).toEqual(["books", "list", { params: undefined }]);
+  });
+
+  it("should generate a details key", () => {
+    expect(keys.details()).toEqual(["books", "detail"]);
+  });
+
+  it("should generate a detail key for a specific ID", () => {
+    expect(keys.detail("123")).toEqual(["books", "detail", "123"]);
+  });
+});

ui/litellm-dashboard/src/app/(dashboard)/hooks/login/useLogin.ts

@@ -3,8 +3,8 @@ import { loginCall, LoginRequest } from "@/components/networking";
 
 export const useLogin = () => {
   return useMutation({
-    mutationFn: async ({ username, password }: LoginRequest) => {
-      const result = await loginCall(username, password);
+    mutationFn: async ({ username, password, useV3 }: LoginRequest) => {
+      const result = await loginCall(username, password, useV3);
       return result;
     },
   });

ui/litellm-dashboard/src/app/(dashboard)/hooks/proxySettings/useProxySettings.ts

@@ -0,0 +1,21 @@
+import { useState, useEffect } from "react";
+import { fetchProxySettings } from "@/utils/proxyUtils";
+import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized";
+
+export default function useProxySettings() {
+  const { accessToken } = useAuthorized();
+  const [proxySettings, setProxySettings] = useState({
+    PROXY_BASE_URL: "",
+    PROXY_LOGOUT_URL: "",
+    LITELLM_UI_API_DOC_BASE_URL: null as string | null,
+  });
+
+  useEffect(() => {
+    if (!accessToken) return;
+    fetchProxySettings(accessToken).then((settings) => {
+      if (settings) setProxySettings(settings);
+    });
+  }, [accessToken]);
+
+  return proxySettings;
+}

ui/litellm-dashboard/src/app/(dashboard)/hooks/uiConfig/useUIConfig.test.ts

@@ -28,6 +28,8 @@ const mockUIConfig: LiteLLMWellKnownUiConfig = {
   proxy_base_url: "https://proxy.example.com",
   auto_redirect_to_sso: true,
   admin_ui_disabled: false,
+  is_control_plane: false,
+  workers: [],
 };
 
 describe("useUIConfig", () => {
@@ -102,6 +104,8 @@ describe("useUIConfig", () => {
       auto_redirect_to_sso: false,
       sso_configured: false,
       admin_ui_disabled: true,
+      is_control_plane: false,
+      workers: [],
     };
 
     // Mock successful API call with different data

ui/litellm-dashboard/src/app/(dashboard)/teams/components/TeamsHeaderTabs.test.tsx

@@ -0,0 +1,54 @@
+import { render, screen } from "@testing-library/react";
+import React from "react";
+import { describe, expect, it, vi } from "vitest";
+import TeamsHeaderTabs from "./TeamsHeaderTabs";
+
+vi.mock("@tremor/react", () => ({
+  TabGroup: ({ children, ...props }: any) => <div data-testid="tab-group" {...props}>{children}</div>,
+  TabList: ({ children, ...props }: any) => <div data-testid="tab-list" {...props}>{children}</div>,
+  Tab: ({ children, ...props }: any) => <button {...props}>{children}</button>,
+  TabPanels: ({ children, ...props }: any) => <div data-testid="tab-panels" {...props}>{children}</div>,
+  Text: ({ children, ...props }: any) => <span {...props}>{children}</span>,
+  Icon: ({ onClick, ...props }: any) => <button data-testid="refresh-icon" onClick={onClick} />,
+}));
+
+vi.mock("@heroicons/react/outline", () => ({
+  RefreshIcon: () => <svg data-testid="refresh-svg" />,
+}));
+
+const renderTabs = (props: Partial<Parameters<typeof TeamsHeaderTabs>[0]> = {}) => {
+  const defaults = {
+    lastRefreshed: "",
+    onRefresh: vi.fn(),
+    userRole: "Internal User",
+    children: <div data-testid="panel-content">Panel</div>,
+  };
+  return render(<TeamsHeaderTabs {...defaults} {...props} />);
+};
+
+describe("TeamsHeaderTabs", () => {
+  it("should render 'Your Teams' and 'Available Teams' tabs", () => {
+    renderTabs();
+
+    expect(screen.getByText("Your Teams")).toBeInTheDocument();
+    expect(screen.getByText("Available Teams")).toBeInTheDocument();
+  });
+
+  it("should render 'Default Team Settings' tab when user is Admin", () => {
+    renderTabs({ userRole: "Admin" });
+
+    expect(screen.getByText("Default Team Settings")).toBeInTheDocument();
+  });
+
+  it("should not render 'Default Team Settings' tab for non-admin users", () => {
+    renderTabs({ userRole: "Internal User" });
+
+    expect(screen.queryByText("Default Team Settings")).not.toBeInTheDocument();
+  });
+
+  it("should display last refreshed time when provided", () => {
+    renderTabs({ lastRefreshed: "2024-06-01 12:00:00" });
+
+    expect(screen.getByText("Last Refreshed: 2024-06-01 12:00:00")).toBeInTheDocument();
+  });
+});

ui/litellm-dashboard/src/app/(dashboard)/teams/components/TeamsTable/TeamsTable.test.tsx

@@ -0,0 +1,129 @@
+import { render, screen } from "@testing-library/react";
+import userEvent from "@testing-library/user-event";
+import React from "react";
+import { describe, expect, it, vi } from "vitest";
+import { Team } from "@/components/key_team_helpers/key_list";
+import TeamsTable from "./TeamsTable";
+
+vi.mock("@tremor/react", () => ({
+  Button: React.forwardRef<HTMLButtonElement, any>(({ children, ...props }, ref) =>
+    React.createElement("button", { ...props, ref }, children),
+  ),
+  Icon: ({ onClick, ...props }: any) => <button data-testid={props["data-testid"] || "icon-btn"} onClick={onClick} aria-label={props["aria-label"]} />,
+  Table: ({ children }: any) => <table>{children}</table>,
+  TableHead: ({ children }: any) => <thead>{children}</thead>,
+  TableBody: ({ children }: any) => <tbody>{children}</tbody>,
+  TableRow: ({ children }: any) => <tr>{children}</tr>,
+  TableHeaderCell: ({ children }: any) => <th>{children}</th>,
+  TableCell: ({ children, ...props }: any) => <td {...props}>{children}</td>,
+  Text: ({ children }: any) => <span>{children}</span>,
+}));
+
+vi.mock("antd", () => ({
+  Tooltip: ({ children }: any) => <>{children}</>,
+}));
+
+vi.mock("@heroicons/react/outline", () => ({
+  PencilAltIcon: () => <svg data-testid="pencil-icon" />,
+  TrashIcon: () => <svg data-testid="trash-icon" />,
+}));
+
+vi.mock("@/utils/dataUtils", () => ({
+  formatNumberWithCommas: (val: number, decimals: number) =>
+    val != null ? val.toFixed(decimals) : "N/A",
+}));
+
+vi.mock("@/app/(dashboard)/teams/components/TeamsTable/ModelsCell", () => ({
+  default: ({ team }: any) => <td data-testid="models-cell">{team.models.join(",")}</td>,
+}));
+
+vi.mock("@/app/(dashboard)/teams/components/TeamsTable/YourRoleCell/YourRoleCell", () => ({
+  default: ({ team }: any) => <td data-testid="role-cell">{team.team_id}</td>,
+}));
+
+const makeTeam = (overrides: Partial<Team> = {}): Team => ({
+  team_id: "team-abc1234",
+  team_alias: "Platform",
+  models: ["gpt-4"],
+  max_budget: 500,
+  budget_duration: null,
+  tpm_limit: null,
+  rpm_limit: null,
+  organization_id: "org-1",
+  created_at: "2024-06-01T00:00:00Z",
+  keys: [],
+  members_with_roles: [],
+  spend: 123.4567,
+  ...overrides,
+});
+
+const defaultPerTeamInfo = {
+  "team-abc1234": {
+    keys: [{ token: "tok-1" } as any, { token: "tok-2" } as any],
+    team_info: {
+      members_with_roles: [{ user_id: "u1", role: "admin" } as any],
+    },
+  },
+};
+
+const renderTable = (overrides: Partial<Parameters<typeof TeamsTable>[0]> = {}) => {
+  const defaults = {
+    teams: [makeTeam()],
+    currentOrg: null,
+    perTeamInfo: defaultPerTeamInfo,
+    userRole: "Admin",
+    userId: "user-1",
+    setSelectedTeamId: vi.fn(),
+    setEditTeam: vi.fn(),
+    onDeleteTeam: vi.fn(),
+  };
+  return render(<TeamsTable {...defaults} {...overrides} />);
+};
+
+describe("TeamsTable", () => {
+  it("should render table headers", () => {
+    renderTable();
+
+    expect(screen.getByText("Team Name")).toBeInTheDocument();
+    expect(screen.getByText("Team ID")).toBeInTheDocument();
+    expect(screen.getByText("Created")).toBeInTheDocument();
+    expect(screen.getByText("Spend (USD)")).toBeInTheDocument();
+    expect(screen.getByText("Budget (USD)")).toBeInTheDocument();
+    expect(screen.getByText("Models")).toBeInTheDocument();
+    expect(screen.getByText("Organization")).toBeInTheDocument();
+    expect(screen.getByText("Your Role")).toBeInTheDocument();
+    expect(screen.getByText("Info")).toBeInTheDocument();
+  });
+
+  it("should render team rows with team data", () => {
+    renderTable();
+
+    expect(screen.getByText("Platform")).toBeInTheDocument();
+    expect(screen.getByText("team-ab...")).toBeInTheDocument();
+    expect(screen.getByText("org-1")).toBeInTheDocument();
+  });
+
+  it("should show edit and delete icons for Admin users", () => {
+    renderTable({ userRole: "Admin" });
+
+    expect(screen.getAllByTestId("icon-btn").length).toBeGreaterThanOrEqual(2);
+  });
+
+  it("should not show edit and delete icons for non-Admin users", () => {
+    renderTable({ userRole: "Internal User" });
+
+    // Only the team ID button should be present, no icon-btn for edit/delete
+    const iconBtns = screen.queryAllByTestId("icon-btn");
+    expect(iconBtns).toHaveLength(0);
+  });
+
+  it("should call setSelectedTeamId when team ID button is clicked", async () => {
+    const user = userEvent.setup();
+    const setSelectedTeamId = vi.fn();
+    renderTable({ setSelectedTeamId });
+
+    await user.click(screen.getByText("team-ab..."));
+
+    expect(setSelectedTeamId).toHaveBeenCalledWith("team-abc1234");
+  });
+});

ui/litellm-dashboard/src/app/login/LoginPage.test.tsx

@@ -41,6 +41,17 @@ vi.mock("@/app/(dashboard)/hooks/login/useLogin", () => ({
   })),
 }));
 
+vi.mock("@/hooks/useWorker", () => ({
+  useWorker: vi.fn(() => ({
+    isControlPlane: false,
+    workers: [],
+    selectedWorkerId: null,
+    selectedWorker: null,
+    selectWorker: vi.fn(),
+    disconnectFromWorker: vi.fn(),
+  })),
+}));
+
 import { useUIConfig } from "@/app/(dashboard)/hooks/uiConfig/useUIConfig";
 import { getCookie } from "@/utils/cookieUtils";
 import { isJwtExpired } from "@/utils/jwtUtils";
@@ -108,7 +119,7 @@ describe("LoginPage", () => {
     );
 
     await waitFor(() => {
-      expect(mockReplace).toHaveBeenCalledWith("http://localhost:4000/ui");
+      expect(mockReplace).toHaveBeenCalledWith("/ui");
     });
   });
 
@@ -189,7 +200,7 @@ describe("LoginPage", () => {
     );
 
     await waitFor(() => {
-      expect(mockReplace).toHaveBeenCalledWith("http://localhost:4000/ui");
+      expect(mockReplace).toHaveBeenCalledWith("/ui");
     });
 
     expect(mockPush).not.toHaveBeenCalled();

ui/litellm-dashboard/src/app/login/LoginPage.tsx

@@ -3,14 +3,15 @@
 import { useLogin } from "@/app/(dashboard)/hooks/login/useLogin";
 import { useUIConfig } from "@/app/(dashboard)/hooks/uiConfig/useUIConfig";
 import LoadingScreen from "@/components/common_components/LoadingScreen";
-import { getProxyBaseUrl } from "@/components/networking";
-import { getCookie } from "@/utils/cookieUtils";
+import { exchangeLoginCode, getProxyBaseUrl, switchToWorkerUrl } from "@/components/networking";
+import { clearTokenCookies, getCookie } from "@/utils/cookieUtils";
 import { isJwtExpired } from "@/utils/jwtUtils";
 import { consumeReturnUrl, getReturnUrl, isValidReturnUrl } from "@/utils/returnUrlUtils";
-import { InfoCircleOutlined } from "@ant-design/icons";
-import { Alert, Button, Card, Form, Input, Popover, Space, Typography } from "antd";
+import { InfoCircleOutlined, CloudServerOutlined } from "@ant-design/icons";
+import { Alert, Button, Card, Form, Input, Popover, Select, Space, Typography } from "antd";
 import { useRouter } from "next/navigation";
 import { useEffect, useState } from "react";
+import { useWorker } from "@/hooks/useWorker";
 
 function LoginPageContent() {
   const [username, setUsername] = useState("");
@@ -19,6 +20,17 @@ function LoginPageContent() {
   const { data: uiConfig, isLoading: isConfigLoading } = useUIConfig();
   const loginMutation = useLogin();
   const router = useRouter();
+  const { workers, selectWorker } = useWorker();
+  const [selectedWorkerId, setSelectedWorkerId] = useState<string | null>(null);
+
+  // Pre-select worker from URL param (e.g. /ui/login?worker=team-b)
+  useEffect(() => {
+    const params = new URLSearchParams(window.location.search);
+    const workerParam = params.get("worker");
+    if (workerParam) {
+      setSelectedWorkerId(workerParam);
+    }
+  }, []);
 
   useEffect(() => {
     if (isConfigLoading) {
@@ -31,6 +43,44 @@ function LoginPageContent() {
       return;
     }
 
+    // Cross-origin SSO: worker redirected back with a single-use code.
+    // Exchange it for the JWT via the worker's /v3/login/exchange endpoint.
+    const params = new URLSearchParams(window.location.search);
+    const ssoCode = params.get("code");
+    if (ssoCode) {
+      const workerUrl = localStorage.getItem("litellm_worker_url");
+      exchangeLoginCode(ssoCode, workerUrl).then(() => {
+        params.delete("code");
+        const cleanSearch = params.toString();
+        window.history.replaceState(null, "", window.location.pathname + (cleanSearch ? `?${cleanSearch}` : ""));
+        router.replace("/ui/?login=success");
+      });
+      return;
+    }
+
+    // Backwards compat: handle direct token in URL (legacy flow)
+    const urlToken = params.get("token");
+    if (urlToken && !isJwtExpired(urlToken)) {
+      document.cookie = `token=${urlToken}; path=/; SameSite=Lax`;
+      params.delete("token");
+      const cleanSearch = params.toString();
+      window.history.replaceState(
+        null,
+        "",
+        window.location.pathname + (cleanSearch ? `?${cleanSearch}` : ""),
+      );
+      router.replace("/ui/?login=success");
+      return;
+    }
+
+    // If switching workers on a control plane, clear the old token and show login
+    const switchingWorker = params.has("worker");
+    if (switchingWorker && uiConfig?.is_control_plane) {
+      clearTokenCookies();
+      setIsLoading(false);
+      return;
+    }
+
     const rawToken = getCookie("token");
     if (rawToken && !isJwtExpired(rawToken)) {
       // User already logged in - redirect to return URL or default
@@ -38,7 +88,7 @@ function LoginPageContent() {
       if (returnUrl) {
         router.replace(returnUrl);
       } else {
-        router.replace(`${getProxyBaseUrl()}/ui`);
+        router.replace("/ui");
       }
       return;
     }
@@ -58,16 +108,35 @@ function LoginPageContent() {
   }, [isConfigLoading, router, uiConfig]);
 
   const handleSubmit = () => {
+    // If a worker is selected, point proxyBaseUrl at it before login
+    const selectedWorker = workers.find((w) => w.worker_id === selectedWorkerId);
+    if (selectedWorker) {
+      switchToWorkerUrl(selectedWorker.url);
+    }
+
     loginMutation.mutate(
-      { username, password },
+      { username, password, useV3: !!selectedWorker },
       {
         onSuccess: (data) => {
-          // Check if we have a return URL to use instead of the default redirect
-          const returnUrl = consumeReturnUrl();
-          if (returnUrl) {
-            router.push(returnUrl);
+          // Update the worker context with the selected worker
+          if (selectedWorker) {
+            selectWorker(selectedWorker.worker_id);
+            // Stay on the CP's UI — proxyBaseUrl already points at the worker
+            router.push("/ui/?login=success");
           } else {
-            router.push(data.redirect_url);
+            // Normal (non-control-plane) login — follow the server's redirect
+            const returnUrl = consumeReturnUrl();
+            if (returnUrl) {
+              router.push(returnUrl);
+            } else {
+              router.push(data.redirect_url);
+            }
+          }
+        },
+        onError: () => {
+          // Reset proxyBaseUrl on login failure
+          if (selectedWorker) {
+            switchToWorkerUrl(null);
           }
         },
       },
@@ -154,6 +223,22 @@ function LoginPageContent() {
           {error && <Alert message={error} type="error" showIcon />}
 
           <Form onFinish={handleSubmit} layout="vertical" requiredMark={true}>
+            {uiConfig?.is_control_plane && workers.length > 0 && (
+              <Form.Item label="Worker" style={{ marginBottom: 16 }}>
+                <Select
+                  value={selectedWorkerId || undefined}
+                  onChange={(value) => setSelectedWorkerId(value)}
+                  placeholder="Choose a worker to connect to"
+                  size="large"
+                  suffixIcon={<CloudServerOutlined />}
+                  options={workers.map((w) => ({
+                    label: w.name,
+                    value: w.worker_id,
+                  }))}
+                />
+              </Form.Item>
+            )}
+
             <Form.Item
               label="Username"
               name="username"
@@ -209,10 +294,20 @@ function LoginPageContent() {
                 </Popover>
               ) : (
                 <Button
-                  disabled={isLoginLoading}
-                  onClick={() =>
-                    router.push(`${getProxyBaseUrl()}/sso/key/generate`)
-                  }
+                  disabled={isLoginLoading || (!!selectedWorkerId && workers.length === 0)}
+                  onClick={() => {
+                    const selectedWorker = workers.find((w) => w.worker_id === selectedWorkerId);
+                    if (selectedWorker) {
+                      // Store worker selection so useWorker hook restores it after redirect
+                      localStorage.setItem("litellm_selected_worker_id", selectedWorkerId!);
+                      switchToWorkerUrl(selectedWorker.url);
+                    }
+                    // SSO on the worker (or this instance if no worker), always
+                    // include return_to so the callback redirects back here
+                    const ssoBase = selectedWorker?.url ?? getProxyBaseUrl();
+                    const returnTo = encodeURIComponent(window.location.origin + "/ui/login");
+                    router.push(`${ssoBase}/sso/key/generate?return_to=${returnTo}`);
+                  }}
                   block
                   size="large"
                 >

ui/litellm-dashboard/src/app/onboarding/OnboardingForm.test.tsx

@@ -0,0 +1,95 @@
+import { render, screen } from "@testing-library/react";
+import React from "react";
+import { describe, expect, it, vi } from "vitest";
+import { OnboardingForm } from "./OnboardingForm";
+
+const mockUseOnboardingCredentials = vi.fn();
+const mockClaimToken = vi.fn();
+
+vi.mock("next/navigation", () => ({
+  useSearchParams: () => new URLSearchParams("invitation_id=inv-123"),
+}));
+
+vi.mock("jwt-decode", () => ({
+  jwtDecode: vi.fn(() => ({
+    user_email: "alice@example.com",
+    user_id: "user-1",
+    key: "access-tok",
+  })),
+}));
+
+vi.mock("@/app/(dashboard)/hooks/onboarding/useOnboarding", () => ({
+  useOnboardingCredentials: (...args: unknown[]) => mockUseOnboardingCredentials(...args),
+  useClaimOnboardingToken: () => ({ mutate: mockClaimToken, isPending: false }),
+}));
+
+vi.mock("@/components/networking", () => ({
+  getProxyBaseUrl: vi.fn(() => ""),
+}));
+
+vi.mock("./OnboardingLoadingView", () => ({
+  OnboardingLoadingView: () => <div data-testid="loading-view">Loading</div>,
+}));
+
+vi.mock("./OnboardingErrorView", () => ({
+  OnboardingErrorView: () => <div data-testid="error-view">Error</div>,
+}));
+
+vi.mock("./OnboardingFormBody", () => ({
+  OnboardingFormBody: ({ variant, userEmail }: { variant: string; userEmail: string }) => (
+    <div data-testid="form-body" data-variant={variant} data-email={userEmail}>
+      Form Body
+    </div>
+  ),
+}));
+
+describe("OnboardingForm", () => {
+  it("should render loading view when credentials are loading", () => {
+    mockUseOnboardingCredentials.mockReturnValue({
+      data: undefined,
+      isLoading: true,
+      isError: false,
+    });
+
+    render(<OnboardingForm variant="signup" />);
+
+    expect(screen.getByTestId("loading-view")).toBeInTheDocument();
+  });
+
+  it("should render error view when credentials fail to load", () => {
+    mockUseOnboardingCredentials.mockReturnValue({
+      data: undefined,
+      isLoading: false,
+      isError: true,
+    });
+
+    render(<OnboardingForm variant="signup" />);
+
+    expect(screen.getByTestId("error-view")).toBeInTheDocument();
+  });
+
+  it("should render form body with decoded email when credentials are loaded", () => {
+    mockUseOnboardingCredentials.mockReturnValue({
+      data: { token: "fake-jwt-token" },
+      isLoading: false,
+      isError: false,
+    });
+
+    render(<OnboardingForm variant="signup" />);
+
+    expect(screen.getByTestId("form-body")).toBeInTheDocument();
+    expect(screen.getByTestId("form-body")).toHaveAttribute("data-email", "alice@example.com");
+  });
+
+  it("should pass variant prop to OnboardingFormBody", () => {
+    mockUseOnboardingCredentials.mockReturnValue({
+      data: { token: "fake-jwt-token" },
+      isLoading: false,
+      isError: false,
+    });
+
+    render(<OnboardingForm variant="reset_password" />);
+
+    expect(screen.getByTestId("form-body")).toHaveAttribute("data-variant", "reset_password");
+  });
+});

ui/litellm-dashboard/src/app/page.tsx

@@ -1,6 +1,5 @@
 "use client";
 
-import APIReferenceView from "@/app/(dashboard)/api-reference/APIReferenceView";
 import SidebarProvider from "@/app/(dashboard)/components/SidebarProvider";
 import OldModelDashboard from "@/app/(dashboard)/models-and-endpoints/ModelsAndEndpointsView";
 import PlaygroundPage from "@/app/(dashboard)/playground/page";
@@ -9,7 +8,7 @@ import AgentsPanel from "@/components/agents";
 import BudgetPanel from "@/components/budgets/budget_panel";
 import CacheDashboard from "@/components/cache_dashboard";
 import ClaudeCodePluginsPanel from "@/components/claude_code_plugins";
-import { fetchTeams } from "@/components/common_components/fetch_teams";
+import { teamListCall as v2TeamListCall } from "@/app/(dashboard)/hooks/teams/useTeams";
 import LoadingScreen from "@/components/common_components/LoadingScreen";
 import { CostTrackingSettings } from "@/components/CostTrackingSettings";
 import GeneralSettings from "@/components/general_settings";
@@ -48,7 +47,7 @@ import { buildLoginUrlWithReturn, consumeReturnUrl, normalizeUrlForCompare, stor
 import { formatUserRole, isAdminRole } from "@/utils/roles";
 import { QueryClient, QueryClientProvider } from "@tanstack/react-query";
 import { jwtDecode } from "jwt-decode";
-import { useSearchParams } from "next/navigation";
+import { useRouter, useSearchParams } from "next/navigation";
 import { Suspense, useEffect, useMemo, useRef, useState } from "react";
 import { ConfigProvider, theme } from "antd";
 
@@ -75,6 +74,16 @@ interface ProxySettings {
   LITELLM_UI_API_DOC_BASE_URL?: string | null;
 }
 
+/**
+ * Map of legacy query-param page keys → new path-based route segments.
+ * When a user visits ?page=<key>, they are redirected to /ui/<value>.
+ * Add entries here as pages are migrated from the if/else chain to path-based routes.
+ */
+const LEGACY_REDIRECTS: Record<string, string> = {
+  api_ref: "api-reference",
+  "api-reference": "api-reference",
+};
+
 function CreateKeyPageContent() {
   const [userRole, setUserRole] = useState("");
   const [premiumUser, setPremiumUser] = useState(false);
@@ -90,6 +99,7 @@ function CreateKeyPageContent() {
   });
 
   const [showSSOBanner, setShowSSOBanner] = useState<boolean>(true);
+  const router = useRouter();
   const searchParams = useSearchParams()!;
   const [modelData, setModelData] = useState<any>({ data: [] });
   const [token, setToken] = useState<string | null>(null);
@@ -243,6 +253,15 @@ function CreateKeyPageContent() {
     }
   }, [redirectToLogin]);
 
+  // Redirect legacy query-param pages to their new path-based routes
+  const isLegacyRedirect = page in LEGACY_REDIRECTS;
+  useEffect(() => {
+    if (!authLoading && isLegacyRedirect) {
+      const base = (proxyBaseUrl || "") + "/ui";
+      router.replace(`${base}/${LEGACY_REDIRECTS[page]}`);
+    }
+  }, [authLoading, isLegacyRedirect, page, router]);
+
   // Check for a stored return URL after successful authentication
   // This handles the case where user comes back from SSO and we need to redirect to the original URL
   useEffect(() => {
@@ -339,7 +358,9 @@ function CreateKeyPageContent() {
       fetchUserModels(userID, userRole, accessToken, setUserModels);
     }
     if (accessToken && userID && userRole) {
-      fetchTeams(accessToken, userID, userRole, null, setTeams);
+      v2TeamListCall(accessToken, 1, 100, {
+        userID: userRole !== "Admin" && userRole !== "Admin Viewer" ? userID : null,
+      }).then((response) => setTeams(response.teams ?? [])).catch(console.error);
     }
     if (accessToken) {
       fetchOrganizations(accessToken, setOrganizations);
@@ -427,7 +448,7 @@ function CreateKeyPageContent() {
     setShowClaudeCodePrompt(true);
   };
 
-  if (authLoading || redirectToLogin) {
+  if (authLoading || redirectToLogin || isLegacyRedirect) {
     return <LoadingScreen />;
   }
 
@@ -536,8 +557,6 @@ function CreateKeyPageContent() {
                     <AdminPanel
                       proxySettings={proxySettings}
                     />
-                  ) : page == "api_ref" ? (
-                    <APIReferenceView proxySettings={proxySettings} />
                   ) : page == "logging-and-alerts" ? (
                     <Settings userID={userID} userRole={userRole} accessToken={accessToken} premiumUser={premiumUser} />
                   ) : page == "budgets" ? (

ui/litellm-dashboard/src/components/AccessGroups/AccessGroupsModal/AccessGroupCreateModal.tsx

@@ -1,5 +1,6 @@
 import React from "react";
-import { Modal, Form, message } from "antd";
+import { Modal, Form } from "antd";
+import MessageManager from "@/components/molecules/message_manager";
 import {
   AccessGroupBaseForm,
   AccessGroupFormValues,
@@ -37,7 +38,7 @@ export function AccessGroupCreateModal({
 
         createMutation.mutate(params, {
           onSuccess: () => {
-            message.success("Access group created successfully");
+            MessageManager.success("Access group created successfully");
             form.resetFields();
             onSuccess?.();
             onCancel();

ui/litellm-dashboard/src/components/AccessGroups/AccessGroupsModal/AccessGroupEditModal.tsx

@@ -1,5 +1,6 @@
 import React, { useEffect } from "react";
-import { Modal, Form, message } from "antd";
+import { Modal, Form } from "antd";
+import MessageManager from "@/components/molecules/message_manager";
 import {
   AccessGroupBaseForm,
   AccessGroupFormValues,
@@ -55,7 +56,7 @@ export function AccessGroupEditModal({
           { accessGroupId: accessGroup.access_group_id, params },
           {
             onSuccess: () => {
-              message.success("Access group updated successfully");
+              MessageManager.success("Access group updated successfully");
               onSuccess?.();
               onCancel();
             },

ui/litellm-dashboard/src/components/BulkEditUsers.tsx

@@ -3,7 +3,6 @@ import {
   Modal,
   Typography,
   Divider,
-  message,
   Table,
   Select,
   InputNumber,
@@ -14,6 +13,7 @@ import {
 import { userBulkUpdateUserCall, teamBulkMemberAddCall, Member } from "./networking";
 import { UserEditView } from "./user_edit_view";
 import NotificationsManager from "./molecules/notifications_manager";
+import MessageManager from "@/components/molecules/message_manager";
 
 const { Text, Title } = Typography;
 
@@ -188,7 +188,7 @@ const BulkEditUserModal: React.FC<BulkEditUserModalProps> = ({
         }
 
         if (failedTeams.length > 0) {
-          message.warning(`Failed to add users to ${failedTeams.length} team(s)`);
+          MessageManager.warning(`Failed to add users to ${failedTeams.length} team(s)`);
         }
       }
 

ui/litellm-dashboard/src/components/CloudZeroCostTracking/CloudZeroCreateModal.tsx

@@ -1,4 +1,5 @@
-import { Form, Modal, Input, message } from "antd";
+import { Form, Modal, Input } from "antd";
+import MessageManager from "@/components/molecules/message_manager";
 import { useEffect } from "react";
 import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized";
 import { useCloudZeroCreate } from "@/app/(dashboard)/hooks/cloudzero/useCloudZeroCreate";
@@ -31,7 +32,7 @@ export default function CloudZeroCreationModal({ open, onOk, onCancel }: CloudZe
         },
         {
           onSuccess: () => {
-            message.success("CloudZero integration created successfully");
+            MessageManager.success("CloudZero integration created successfully");
             form.resetFields();
             onOk();
           },
@@ -39,7 +40,7 @@ export default function CloudZeroCreationModal({ open, onOk, onCancel }: CloudZe
             if (error?.errorFields) {
               return;
             }
-            message.error(error?.message || "Failed to create CloudZero integration");
+            MessageManager.error(error?.message || "Failed to create CloudZero integration");
           },
         },
       );
@@ -47,7 +48,7 @@ export default function CloudZeroCreationModal({ open, onOk, onCancel }: CloudZe
       if (error?.errorFields) {
         return;
       }
-      message.error(error?.message || "Failed to create CloudZero integration");
+      MessageManager.error(error?.message || "Failed to create CloudZero integration");
     }
   };
 

ui/litellm-dashboard/src/components/CloudZeroCostTracking/CloudZeroIntegrationSettings.tsx

@@ -3,7 +3,8 @@ import { useCloudZeroExport } from "@/app/(dashboard)/hooks/cloudzero/useCloudZe
 import { useCloudZeroDeleteSettings } from "@/app/(dashboard)/hooks/cloudzero/useCloudZeroSettings";
 import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized";
 import DeleteResourceModal from "@/components/common_components/DeleteResourceModal";
-import { Alert, Button, Card, Descriptions, Divider, message, Popconfirm, Tag } from "antd";
+import { Alert, Button, Card, Descriptions, Divider, Popconfirm, Tag } from "antd";
+import MessageManager from "@/components/molecules/message_manager";
 import { CheckCircle, Edit, Play, Trash2, Upload } from "lucide-react";
 import { useState } from "react";
 import CloudZeroUpdateModal from "./CloudZeroUpdateModal";
@@ -30,10 +31,10 @@ export function CloudZeroIntegrationSettings({ settings, onSettingsUpdated }: Cl
       { limit: 10 },
       {
         onSuccess: (data) => {
-          message.success("Dry run completed successfully");
+          MessageManager.success("Dry run completed successfully");
         },
         onError: (error) => {
-          message.error(error?.message || "Failed to perform dry run");
+          MessageManager.error(error?.message || "Failed to perform dry run");
         },
       },
     );
@@ -48,10 +49,10 @@ export function CloudZeroIntegrationSettings({ settings, onSettingsUpdated }: Cl
       { operation: "replace_hourly" },
       {
         onSuccess: () => {
-          message.success("Data successfully exported to CloudZero");
+          MessageManager.success("Data successfully exported to CloudZero");
         },
         onError: (error) => {
-          message.error(error?.message || "Failed to export data");
+          MessageManager.error(error?.message || "Failed to export data");
         },
       },
     );
@@ -79,12 +80,12 @@ export function CloudZeroIntegrationSettings({ settings, onSettingsUpdated }: Cl
 
     deleteMutation.mutate(undefined, {
       onSuccess: () => {
-        message.success("CloudZero integration deleted successfully");
+        MessageManager.success("CloudZero integration deleted successfully");
         setIsDeleteModalOpen(false);
         onSettingsUpdated();
       },
       onError: (error) => {
-        message.error(error?.message || "Failed to delete CloudZero integration");
+        MessageManager.error(error?.message || "Failed to delete CloudZero integration");
       },
     });
   };

ui/litellm-dashboard/src/components/CloudZeroCostTracking/CloudZeroUpdateModal.tsx

@@ -1,6 +1,7 @@
 import { useCloudZeroUpdateSettings } from "@/app/(dashboard)/hooks/cloudzero/useCloudZeroSettings";
 import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized";
-import { Form, Input, message, Modal } from "antd";
+import { Form, Input, Modal } from "antd";
+import MessageManager from "@/components/molecules/message_manager";
 import { useEffect } from "react";
 import { CloudZeroSettings } from "./types";
 
@@ -39,7 +40,7 @@ export default function CloudZeroUpdateModal({ open, onOk, onCancel, settings }:
         },
         {
           onSuccess: () => {
-            message.success("CloudZero integration updated successfully");
+            MessageManager.success("CloudZero integration updated successfully");
             form.resetFields();
             onOk();
           },
@@ -47,7 +48,7 @@ export default function CloudZeroUpdateModal({ open, onOk, onCancel, settings }:
             if (error?.errorFields) {
               return;
             }
-            message.error(error?.message || "Failed to update CloudZero integration");
+            MessageManager.error(error?.message || "Failed to update CloudZero integration");
           },
         },
       );
@@ -55,7 +56,7 @@ export default function CloudZeroUpdateModal({ open, onOk, onCancel, settings }:
       if (error?.errorFields) {
         return;
       }
-      message.error(error?.message || "Failed to update CloudZero integration");
+      MessageManager.error(error?.message || "Failed to update CloudZero integration");
     }
   };
 

ui/litellm-dashboard/src/components/Navbar/WorkerDropdown/WorkerDropdown.tsx

@@ -0,0 +1,38 @@
+"use client";
+
+import React from "react";
+import { Select } from "antd";
+import { CloudServerOutlined } from "@ant-design/icons";
+import { useWorker } from "@/hooks/useWorker";
+
+interface WorkerDropdownProps {
+  onWorkerSwitch: (workerId: string) => void;
+}
+
+const WorkerDropdown: React.FC<WorkerDropdownProps> = ({ onWorkerSwitch }) => {
+  const { isControlPlane, selectedWorker, workers } = useWorker();
+
+  if (!isControlPlane || !selectedWorker) return null;
+
+  return (
+    <Select
+      showSearch
+      filterOption={(input, option) =>
+        (option?.label as string ?? "").toLowerCase().includes(input.toLowerCase())
+      }
+      value={selectedWorker.worker_id}
+      style={{ minWidth: 180 }}
+      suffixIcon={<CloudServerOutlined />}
+      options={workers.map((w) => ({
+        label: w.name,
+        value: w.worker_id,
+        disabled: w.worker_id === selectedWorker.worker_id,
+      }))}
+      onChange={(newWorkerId) => {
+        onWorkerSwitch(newWorkerId);
+      }}
+    />
+  );
+};
+
+export default WorkerDropdown;

ui/litellm-dashboard/src/components/OldTeams.test.tsx

@@ -18,8 +18,8 @@ vi.mock("./networking", () => ({
   getPoliciesList: vi.fn().mockResolvedValue({ policies: [] }),
 }));
 
-vi.mock("./common_components/fetch_teams", () => ({
-  fetchTeams: vi.fn(),
+vi.mock("@/app/(dashboard)/hooks/teams/useTeams", () => ({
+  teamListCall: vi.fn().mockResolvedValue({ teams: [], total: 0, page: 1, page_size: 100, total_pages: 0 }),
 }));
 
 vi.mock("./molecules/notifications_manager", () => ({
@@ -375,6 +375,9 @@ describe("OldTeams - handleCreate organization handling", () => {
         organizations={[]}
       />,
     );
+    await waitFor(() => {
+      expect(screen.getByTestId("delete-team-button")).toBeInTheDocument();
+    });
     const deleteTeamButton = screen.getByTestId("delete-team-button");
     act(() => {
       fireEvent.click(deleteTeamButton);
@@ -389,7 +392,7 @@ describe("OldTeams - empty state", () => {
     mockUseOrganizations.mockReturnValue({ data: [] });
   });
 
-  it("should display empty state message when teams array is empty", () => {
+  it("should display empty state message when teams array is empty", async () => {
     renderWithQueryClient(
       <OldTeams
         teams={[]}
@@ -402,11 +405,13 @@ describe("OldTeams - empty state", () => {
       />,
     );
 
-    expect(screen.getByText("No teams found")).toBeInTheDocument();
-    expect(screen.getByText("Adjust your filters or create a new team")).toBeInTheDocument();
+    await waitFor(() => {
+      expect(screen.getByText("No teams yet")).toBeInTheDocument();
+    });
+    expect(screen.getByText("Create your first team to organize members and manage access to models.")).toBeInTheDocument();
   });
 
-  it("should display empty state message when teams is null", () => {
+  it("should display empty state message when teams is null", async () => {
     renderWithQueryClient(
       <OldTeams
         teams={null}
@@ -419,11 +424,13 @@ describe("OldTeams - empty state", () => {
       />,
     );
 
-    expect(screen.getByText("No teams found")).toBeInTheDocument();
-    expect(screen.getByText("Adjust your filters or create a new team")).toBeInTheDocument();
+    await waitFor(() => {
+      expect(screen.getByText("No teams yet")).toBeInTheDocument();
+    });
+    expect(screen.getByText("Create your first team to organize members and manage access to models.")).toBeInTheDocument();
   });
 
-  it("should not display empty state when teams array has items", () => {
+  it("should not display empty state when teams array has items", async () => {
     renderWithQueryClient(
       <OldTeams
         teams={[
@@ -451,9 +458,11 @@ describe("OldTeams - empty state", () => {
       />,
     );
 
-    expect(screen.queryByText("No teams found")).not.toBeInTheDocument();
-    expect(screen.queryByText("Adjust your filters or create a new team")).not.toBeInTheDocument();
-    expect(screen.getByText("Test Team")).toBeInTheDocument();
+    await waitFor(() => {
+      expect(screen.getByText("Test Team")).toBeInTheDocument();
+    });
+    expect(screen.queryByText("No teams yet")).not.toBeInTheDocument();
+    expect(screen.queryByText("Create your first team to organize members and manage access to models.")).not.toBeInTheDocument();
   });
 });
 
@@ -621,12 +630,9 @@ describe("OldTeams - premium props", () => {
       />,
     );
 
-    const truncatedTeamId = "team-123456789".slice(0, 7);
-    const teamButton = await screen.findByRole("button", {
-      name: new RegExp(`${truncatedTeamId}\\.\\.\\.`),
-    });
+    const teamIdElement = await screen.findByText("team-123456789");
     act(() => {
-      fireEvent.click(teamButton);
+      fireEvent.click(teamIdElement);
     });
 
     await waitFor(() => expect(mockTeamInfoView).toHaveBeenCalled());
@@ -798,7 +804,7 @@ describe("OldTeams - access_group_ids in team create", () => {
       />,
     );
 
-    const createButton = screen.getByRole("button", { name: /create new team/i });
+    const createButton = screen.getAllByRole("button", { name: /create team/i })[0];
     act(() => {
       fireEvent.click(createButton);
     });
@@ -823,7 +829,8 @@ describe("OldTeams - access_group_ids in team create", () => {
     const accessGroupInput = screen.getByTestId("access-group-selector");
     fireEvent.change(accessGroupInput, { target: { value: "ag-1,ag-2" } });
 
-    const createTeamSubmitButton = screen.getByRole("button", { name: /create team/i });
+    const createTeamSubmitButtons = screen.getAllByRole("button", { name: /create team/i });
+    const createTeamSubmitButton = createTeamSubmitButtons[createTeamSubmitButtons.length - 1];
     fireEvent.click(createTeamSubmitButton);
 
     await waitFor(() => {
@@ -865,7 +872,7 @@ describe("OldTeams - models dropdown options", () => {
       expect(fetchAvailableModelsForTeamOrKey).toHaveBeenCalled();
     });
 
-    const createButton = screen.getByRole("button", { name: /create new team/i });
+    const createButton = screen.getAllByRole("button", { name: /create team/i })[0];
     act(() => {
       fireEvent.click(createButton);
     });
@@ -884,7 +891,7 @@ describe("OldTeams - organization alias display", () => {
     mockUseOrganizations.mockReturnValue({ data: [] });
   });
 
-  it("should display organization alias instead of organization id", () => {
+  it("should display organization alias instead of organization id", async () => {
     const mockOrganizations = [
       {
         organization_id: "org-123",
@@ -934,11 +941,13 @@ describe("OldTeams - organization alias display", () => {
       />,
     );
 
-    expect(screen.getByText("Test Organization")).toBeInTheDocument();
+    await waitFor(() => {
+      expect(screen.getByText("Test Organization")).toBeInTheDocument();
+    });
     expect(screen.queryByText("org-123")).not.toBeInTheDocument();
   });
 
-  it("should display organization id when alias is not found", () => {
+  it("should display organization id when alias is not found", async () => {
     mockUseOrganizations.mockReturnValue({ data: [] });
 
     renderWithQueryClient(
@@ -968,10 +977,12 @@ describe("OldTeams - organization alias display", () => {
       />,
     );
 
-    expect(screen.getByText("org-unknown")).toBeInTheDocument();
+    await waitFor(() => {
+      expect(screen.getByText("org-unknown")).toBeInTheDocument();
+    });
   });
 
-  it("should display N/A when organization_id is null", () => {
+  it("should display N/A when organization_id is null", async () => {
     mockUseOrganizations.mockReturnValue({ data: [] });
 
     renderWithQueryClient(
@@ -1001,6 +1012,9 @@ describe("OldTeams - organization alias display", () => {
       />,
     );
 
-    expect(screen.getByText("N/A")).toBeInTheDocument();
+    await waitFor(() => {
+      // When organization_id is null, the table shows "—" in the Organization column
+      expect(screen.getAllByText("—").length).toBeGreaterThan(0);
+    });
   });
 });

ui/litellm-dashboard/src/components/Projects/ProjectModals/CreateProjectModal.tsx

@@ -1,5 +1,6 @@
-import { Modal, Form, Button, Typography, message } from "antd";
+import { Modal, Form, Button, Typography } from "antd";
 import { FolderAddOutlined } from "@ant-design/icons";
+import MessageManager from "@/components/molecules/message_manager";
 import {
   useCreateProject,
   ProjectCreateParams,
@@ -32,12 +33,12 @@ export function CreateProjectModal({
 
       createMutation.mutate(params, {
         onSuccess: () => {
-          message.success("Project created successfully");
+          MessageManager.success("Project created successfully");
           form.resetFields();
           onClose();
         },
         onError: (error) => {
-          message.error(error.message || "Failed to create project");
+          MessageManager.error(error.message || "Failed to create project");
         },
       });
     } catch (error) {

ui/litellm-dashboard/src/components/Projects/ProjectModals/EditProjectModal.tsx

@@ -1,6 +1,7 @@
 import { useEffect } from "react";
-import { Modal, Form, Button, Typography, message } from "antd";
+import { Modal, Form, Button, Typography } from "antd";
 import { SaveOutlined } from "@ant-design/icons";
+import MessageManager from "@/components/molecules/message_manager";
 import { ProjectResponse } from "@/app/(dashboard)/hooks/projects/useProjects";
 import {
   useUpdateProject,
@@ -80,12 +81,12 @@ export function EditProjectModal({
         { projectId: project.project_id, params },
         {
           onSuccess: () => {
-            message.success("Project updated successfully");
+            MessageManager.success("Project updated successfully");
             onSuccess?.();
             onClose();
           },
           onError: (error) => {
-            message.error(error.message || "Failed to update project");
+            MessageManager.error(error.message || "Failed to update project");
           },
         },
       );

ui/litellm-dashboard/src/components/SearchTools/SearchToolTester.tsx

@@ -1,5 +1,6 @@
 import React, { useState } from "react";
-import { Button, Input, Typography, Spin, message } from "antd";
+import { Button, Input, Typography, Spin } from "antd";
+import MessageManager from "@/components/molecules/message_manager";
 import { SearchOutlined, LoadingOutlined } from "@ant-design/icons";
 import { searchToolQueryCall } from "../networking";
 import NotificationsManager from "../molecules/notifications_manager";
@@ -39,7 +40,7 @@ export const SearchToolTester: React.FC<SearchToolTesterProps> = ({ searchToolNa
 
   const handleSearch = async () => {
     if (!query.trim()) {
-      message.warning("Please enter a search query");
+      MessageManager.warning("Please enter a search query");
       return;
     }
 

ui/litellm-dashboard/src/components/Settings/RouterSettings/Fallbacks/AddFallbacks.tsx

@@ -5,8 +5,9 @@
  */
 
 import { Button as TremorButton } from "@tremor/react";
-import { Button, message } from "antd";
+import { Button } from "antd";
 import React, { useEffect, useState } from "react";
+import MessageManager from "@/components/molecules/message_manager";
 import NotificationManager from "../../../molecules/notifications_manager";
 import { fetchAvailableModels, ModelGroup } from "../../../playground/llm_calls/fetch_models";
 import { AddFallbacksModal } from "./AddFallbacksModal";
@@ -90,7 +91,7 @@ export default function AddFallbacks({
       (g) => !g.primaryModel || g.fallbackModels.length === 0,
     );
     if (invalidGroups.length > 0) {
-      message.error(
+      MessageManager.error(
         `Please complete configuration for all groups. ${invalidGroups.length} group(s) incomplete.`,
       );
       return;

ui/litellm-dashboard/src/components/Settings/RouterSettings/Fallbacks/FallbackSelectionForm.tsx

@@ -5,9 +5,10 @@
  */
 
 import { Button } from "@tremor/react";
-import { message, Tabs } from "antd";
+import { Tabs } from "antd";
 import { Plus } from "lucide-react";
 import React, { useEffect, useState } from "react";
+import MessageManager from "@/components/molecules/message_manager";
 import { FallbackGroup, FallbackGroupConfig } from "./FallbackGroupConfig";
 
 interface FallbackSelectionFormProps {
@@ -60,7 +61,7 @@ export function FallbackSelectionForm({
 
   const handleRemoveGroup = (targetId: string) => {
     if (groups.length === 1) {
-      message.warning("At least one group is required");
+      MessageManager.warning("At least one group is required");
       return;
     }
     const newGroups = groups.filter((g) => g.id !== targetId);

ui/litellm-dashboard/src/components/agents/add_agent_form.tsx

@@ -1,5 +1,6 @@
 import React, { useState, useEffect } from "react";
-import { Modal, Form, message, Select, Input, Steps, Radio, Tag, Divider, Switch, InputNumber, Collapse } from "antd";
+import { Modal, Form, Select, Input, Steps, Radio, Tag, Divider, Switch, InputNumber, Collapse } from "antd";
+import MessageManager from "@/components/molecules/message_manager";
 import { Button } from "@tremor/react";
 import { CheckCircleFilled, KeyOutlined, RobotOutlined, AppstoreOutlined, InfoCircleOutlined } from "@ant-design/icons";
 import CreatedKeyDisplay from "../shared/CreatedKeyDisplay";
@@ -216,7 +217,7 @@ const AddAgentForm: React.FC<AddAgentFormProps> = ({
 
   const handleCreateAgent = async () => {
     if (!accessToken) {
-      message.error("No access token available");
+      MessageManager.error("No access token available");
       return;
     }
 
@@ -226,7 +227,7 @@ const AddAgentForm: React.FC<AddAgentFormProps> = ({
       const values = { ...form.getFieldsValue(true) };
       const agentData = buildAgentData(values);
       if (!agentData) {
-        message.error("Failed to build agent data");
+        MessageManager.error("Failed to build agent data");
         setIsSubmitting(false);
         return;
       }
@@ -301,7 +302,7 @@ const AddAgentForm: React.FC<AddAgentFormProps> = ({
         setCreatedKeyValue(keyResponse.key || null);
       } else if (keyAssignOption === "existing_key") {
         if (!selectedExistingKey) {
-          message.error("Please select an existing key to assign");
+          MessageManager.error("Please select an existing key to assign");
           setIsSubmitting(false);
           return;
         }
@@ -318,7 +319,7 @@ const AddAgentForm: React.FC<AddAgentFormProps> = ({
     } catch (error) {
       console.error("Error creating agent:", error);
       const errorMessage = error instanceof Error ? error.message : String(error);
-      message.error(errorMessage ? `Failed to create agent: ${errorMessage}` : "Failed to create agent");
+      MessageManager.error(errorMessage ? `Failed to create agent: ${errorMessage}` : "Failed to create agent");
     } finally {
       setIsSubmitting(false);
     }

ui/litellm-dashboard/src/components/agents/agent_info.tsx

@@ -1,6 +1,7 @@
 import React, { useState, useEffect } from "react";
 import { Card, Title, Text, Button as TremorButton, Tab, TabGroup, TabList, TabPanel, TabPanels} from "@tremor/react";
-import { Form, Input, InputNumber, Button as AntButton, message, Spin, Descriptions, Divider } from "antd";
+import { Form, Input, InputNumber, Button as AntButton, Spin, Descriptions, Divider } from "antd";
+import MessageManager from "@/components/molecules/message_manager";
 import { ArrowLeftIcon } from "@heroicons/react/outline";
 import { getAgentInfo, patchAgentCall, getAgentCreateMetadata, AgentCreateInfo } from "../networking";
 import { Agent } from "./types";
@@ -72,7 +73,7 @@ const AgentInfoView: React.FC<AgentInfoViewProps> = ({
       }
     } catch (error) {
       console.error("Error fetching agent info:", error);
-      message.error("Failed to load agent information");
+      MessageManager.error("Failed to load agent information");
     } finally {
       setIsLoading(false);
     }
@@ -111,12 +112,12 @@ const AgentInfoView: React.FC<AgentInfoViewProps> = ({
       }
       
       await patchAgentCall(accessToken, agentId, updateData);
-      message.success("Agent updated successfully");
+      MessageManager.success("Agent updated successfully");
       setIsEditing(false);
       fetchAgentInfo();
     } catch (error) {
       console.error("Error updating agent:", error);
-      message.error("Failed to update agent");
+      MessageManager.error("Failed to update agent");
     } finally {
       setIsSaving(false);
     }

ui/litellm-dashboard/src/components/chat/ChatPage.tsx

@@ -1,7 +1,8 @@
 "use client";
 
 import React, { useCallback, useEffect, useRef, useState, useLayoutEffect } from "react";
-import { Tooltip, Skeleton, Popover, message } from "antd";
+import { Tooltip, Skeleton, Popover } from "antd";
+import MessageManager from "@/components/molecules/message_manager";
 import {
   SettingOutlined,
   PlusOutlined,
@@ -212,7 +213,7 @@ const ChatPage: React.FC<ChatPageProps> = ({ accessToken, userRole, userId, user
           localStorage.setItem(LOCALSTORAGE_MODEL_KEY, JSON.stringify([names[0]]));
         }
       })
-      .catch(() => message.error("Could not load models"))
+      .catch(() => MessageManager.error("Could not load models"))
       .finally(() => setIsLoadingModels(false));
   }, [accessToken]);
 

ui/litellm-dashboard/src/components/chat/MCPAppsPanel.tsx

@@ -5,7 +5,7 @@ import { Spin, Input, Button, Skeleton } from "antd";
 import { SearchOutlined, ArrowLeftOutlined, RightOutlined, ToolOutlined, CheckCircleOutlined } from "@ant-design/icons";
 import { deleteMCPOAuthUserCredential, fetchMCPServers, getMCPOAuthUserCredentialStatus, listMCPTools } from "../networking";
 import { AUTH_TYPE, MCPServer, MCPTool, handleTransport } from "../mcp_tools/types";
-import { message } from "antd";
+import MessageManager from "@/components/molecules/message_manager";
 import { useUserMcpOAuthFlow } from "@/hooks/useUserMcpOAuthFlow";
 
 // ── OAuth2 connect button ─────────────────────────────────────────────────────
@@ -198,7 +198,7 @@ const MCPAppsPanel: React.FC<Props> = ({ accessToken, selectedServers, onChange
       const idToFetch = serverId ?? serverName;
       const result = await listMCPTools(accessToken, idToFetch);
       if (result?.error) {
-        message.warning(`Could not load tools for ${serverName}`);
+        MessageManager.warning(`Could not load tools for ${serverName}`);
         return;
       }
       // Use the ref so we read the most up-to-date list; guard against duplicates
@@ -207,7 +207,7 @@ const MCPAppsPanel: React.FC<Props> = ({ accessToken, selectedServers, onChange
         onChange([...selectedServersRef.current, serverName]);
       }
     } catch {
-      message.warning(`Could not load tools for ${serverName}`);
+      MessageManager.warning(`Could not load tools for ${serverName}`);
     } finally {
       setTogglingOn((prev) => {
         const next = new Set(prev);

ui/litellm-dashboard/src/components/chat/MCPConnectPicker.tsx

@@ -1,5 +1,6 @@
 import React, { useEffect, useState } from "react";
-import { Switch, Spin, message } from "antd";
+import { Switch, Spin } from "antd";
+import MessageManager from "@/components/molecules/message_manager";
 import { fetchMCPServers, listMCPTools } from "../networking";
 import { MCPServer } from "../mcp_tools/types";
 
@@ -57,7 +58,7 @@ const MCPConnectPicker: React.FC<Props> = ({ accessToken, selectedServers, onCha
       const result = await listMCPTools(accessToken, serverName);
       // listMCPTools never throws; it returns { tools, error, message } on failure
       if (result?.error) {
-        message.warning(
+        MessageManager.warning(
           `Could not load tools for ${serverName} — it will be excluded from this message.`
         );
         // Do not add to selectedServers
@@ -65,7 +66,7 @@ const MCPConnectPicker: React.FC<Props> = ({ accessToken, selectedServers, onCha
       }
       onChange([...selectedServers, serverName]);
     } catch {
-      message.warning(
+      MessageManager.warning(
         `Could not load tools for ${serverName} — it will be excluded from this message.`
       );
       // Do not add to selectedServers

ui/litellm-dashboard/src/components/chat/MCPCredentialsTab.tsx

@@ -8,7 +8,8 @@
  */
 
 import React, { useCallback, useEffect, useState } from "react";
-import { Spin, message } from "antd";
+import { Spin } from "antd";
+import MessageManager from "@/components/molecules/message_manager";
 import { DeleteOutlined, LinkOutlined } from "@ant-design/icons";
 import { Badge, Table, TableBody, TableCell, TableHead, TableHeaderCell, TableRow } from "@tremor/react";
 import {
@@ -77,7 +78,7 @@ const MCPCredentialsTab: React.FC<Props> = ({ accessToken }) => {
       await deleteMCPOAuthUserCredential(accessToken, serverId);
       setCredentials((prev) => prev.filter((c) => c.server_id !== serverId));
     } catch {
-      message.error("Failed to revoke connection. Please try again.");
+      MessageManager.error("Failed to revoke connection. Please try again.");
     } finally {
       setRevoking((prev) => { const n = new Set(prev); n.delete(serverId); return n; });
     }

ui/litellm-dashboard/src/components/claude_code_plugins/add_plugin_form.tsx

@@ -1,5 +1,6 @@
 import React, { useState } from "react";
-import { Modal, Form, Input, Select, message } from "antd";
+import { Modal, Form, Input, Select } from "antd";
+import MessageManager from "@/components/molecules/message_manager";
 import { Button } from "@tremor/react";
 import { registerClaudeCodePlugin } from "../networking";
 import {
@@ -43,13 +44,13 @@ const AddPluginForm: React.FC<AddPluginFormProps> = ({
 
   const handleSubmit = async (values: any) => {
     if (!accessToken) {
-      message.error("No access token available");
+      MessageManager.error("No access token available");
       return;
     }
 
     // Validate plugin name
     if (!validatePluginName(values.name)) {
-      message.error(
+      MessageManager.error(
         "Plugin name must be kebab-case (lowercase letters, numbers, and hyphens only)"
       );
       return;
@@ -57,7 +58,7 @@ const AddPluginForm: React.FC<AddPluginFormProps> = ({
 
     // Validate semantic version if provided
     if (values.version && !isValidSemanticVersion(values.version)) {
-      message.error(
+      MessageManager.error(
         "Version must be in semantic versioning format (e.g., 1.0.0)"
       );
       return;
@@ -65,13 +66,13 @@ const AddPluginForm: React.FC<AddPluginFormProps> = ({
 
     // Validate email if provided
     if (values.authorEmail && !isValidEmail(values.authorEmail)) {
-      message.error("Invalid email format");
+      MessageManager.error("Invalid email format");
       return;
     }
 
     // Validate homepage URL if provided
     if (values.homepage && !isValidUrl(values.homepage)) {
-      message.error("Invalid homepage URL format");
+      MessageManager.error("Invalid homepage URL format");
       return;
     }
 
@@ -119,14 +120,14 @@ const AddPluginForm: React.FC<AddPluginFormProps> = ({
       }
 
       await registerClaudeCodePlugin(accessToken, pluginData);
-      message.success("Plugin registered successfully");
+      MessageManager.success("Plugin registered successfully");
       form.resetFields();
       setSourceType("github");
       onSuccess();
       onClose();
     } catch (error) {
       console.error("Error registering plugin:", error);
-      message.error("Failed to register plugin");
+      MessageManager.error("Failed to register plugin");
     } finally {
       setIsSubmitting(false);
     }

ui/litellm-dashboard/src/components/common_components/IconActionButton/TableIconActionButtons/TableIconActionButton.tsx

@@ -6,6 +6,7 @@ import {
   ChevronUpIcon,
   ChevronDownIcon,
   ExternalLinkIcon,
+  ClipboardCopyIcon,
 } from "@heroicons/react/outline";
 import { Tooltip } from "antd";
 import BaseActionButton from "../BaseActionButton";
@@ -32,6 +33,7 @@ export const TableIconActionButtonMap: Record<string, TableIconActionButtonBaseP
   Up: { icon: ChevronUpIcon, className: "hover:text-blue-600" },
   Down: { icon: ChevronDownIcon, className: "hover:text-blue-600" },
   Open: { icon: ExternalLinkIcon, className: "hover:text-green-600" },
+  Copy: { icon: ClipboardCopyIcon, className: "hover:text-blue-600" },
 };
 
 export default function TableIconActionButton({

ui/litellm-dashboard/src/components/common_components/OrganizationDropdown.tsx

@@ -1,13 +1,16 @@
 import React from "react";
-import { Select } from "antd";
+import { Select, Typography } from "antd";
 import { Organization } from "../networking";
 
+const { Text } = Typography;
+
 interface OrganizationDropdownProps {
   organizations?: Organization[] | null;
   value?: string;
   onChange?: (value: string) => void;
   disabled?: boolean;
   loading?: boolean;
+  style?: React.CSSProperties;
 }
 
 const OrganizationDropdown: React.FC<OrganizationDropdownProps> = ({
@@ -16,16 +19,18 @@ const OrganizationDropdown: React.FC<OrganizationDropdownProps> = ({
   onChange,
   disabled,
   loading,
+  style,
 }) => {
   return (
     <Select
       showSearch
-      placeholder="Search or select an organization"
+      placeholder="All Organizations"
       value={value}
       onChange={onChange}
       disabled={disabled}
       loading={loading}
       allowClear
+      style={{ minWidth: 280, ...style }}
       filterOption={(input, option) => {
         if (!option) return false;
         const org = organizations?.find((o) => o.organization_id === option.key);
@@ -37,12 +42,11 @@ const OrganizationDropdown: React.FC<OrganizationDropdownProps> = ({
 
         return orgAlias.includes(searchTerm) || orgId.includes(searchTerm);
       }}
-      optionFilterProp="children"
     >
       {organizations?.map((org) => (
         <Select.Option key={org.organization_id} value={org.organization_id}>
           <span className="font-medium">{org.organization_alias}</span>{" "}
-          <span className="text-gray-500">({org.organization_id})</span>
+          <Text type="secondary">({org.organization_id})</Text>
         </Select.Option>
       ))}
     </Select>

ui/litellm-dashboard/src/components/key_team_helpers/filter_helpers.test.ts

@@ -0,0 +1,90 @@
+import { describe, expect, it, vi } from "vitest";
+import { fetchTeamFilterOptions } from "./filter_helpers";
+
+const mockKeyListCall = vi.fn();
+
+vi.mock("@/components/networking", () => ({
+  keyListCall: (...args: unknown[]) => mockKeyListCall(...args),
+  teamListCall: vi.fn(),
+  organizationListCall: vi.fn(),
+}));
+
+describe("fetchTeamFilterOptions", () => {
+  it("should return empty arrays when accessToken is null", async () => {
+    const result = await fetchTeamFilterOptions(null, "team-1");
+
+    expect(result).toEqual({ keyAliases: [], organizationIds: [], userIds: [] });
+    expect(mockKeyListCall).not.toHaveBeenCalled();
+  });
+
+  it("should return empty arrays when teamId is empty", async () => {
+    const result = await fetchTeamFilterOptions("tok-123", "");
+
+    expect(result).toEqual({ keyAliases: [], organizationIds: [], userIds: [] });
+    expect(mockKeyListCall).not.toHaveBeenCalled();
+  });
+
+  it("should return sorted key aliases from fetched keys", async () => {
+    mockKeyListCall.mockResolvedValue({
+      keys: [
+        { key_alias: "zeta-key" },
+        { key_alias: "alpha-key" },
+        { key_alias: "mid-key" },
+      ],
+      total_pages: 1,
+    });
+
+    const result = await fetchTeamFilterOptions("tok-123", "team-1");
+
+    expect(result.keyAliases).toEqual(["alpha-key", "mid-key", "zeta-key"]);
+  });
+
+  it("should deduplicate organization IDs across pages", async () => {
+    mockKeyListCall
+      .mockResolvedValueOnce({
+        keys: [
+          { organization_id: "org-b" },
+          { organization_id: "org-a" },
+        ],
+        total_pages: 2,
+      })
+      .mockResolvedValueOnce({
+        keys: [
+          { organization_id: "org-a" },
+          { organization_id: "org-c" },
+        ],
+        total_pages: 2,
+      });
+
+    const result = await fetchTeamFilterOptions("tok-123", "team-1");
+
+    expect(result.organizationIds).toEqual(["org-a", "org-b", "org-c"]);
+  });
+
+  it("should map user IDs with email addresses", async () => {
+    mockKeyListCall.mockResolvedValue({
+      keys: [
+        { user_id: "u1", user: { user_email: "alice@example.com" } },
+        { user_id: "u2", user: { user_email: "bob@example.com" } },
+      ],
+      total_pages: 1,
+    });
+
+    const result = await fetchTeamFilterOptions("tok-123", "team-1");
+
+    expect(result.userIds).toEqual(
+      expect.arrayContaining([
+        { id: "u1", email: "alice@example.com" },
+        { id: "u2", email: "bob@example.com" },
+      ]),
+    );
+  });
+
+  it("should handle API errors gracefully and return empty arrays", async () => {
+    mockKeyListCall.mockRejectedValue(new Error("Network error"));
+
+    const result = await fetchTeamFilterOptions("tok-123", "team-1");
+
+    expect(result).toEqual({ keyAliases: [], organizationIds: [], userIds: [] });
+  });
+});

ui/litellm-dashboard/src/components/key_team_helpers/transform_key_info.test.ts

@@ -0,0 +1,62 @@
+import { describe, it, expect } from "vitest";
+import { transformKeyInfo } from "./transform_key_info";
+
+describe("transformKeyInfo", () => {
+  it("should combine key and info fields into a single object", () => {
+    const apiResponse = {
+      key: "sk-abc123",
+      info: {
+        token_id: "tok_1",
+        key_name: "my-key",
+        spend: 10.5,
+      },
+    };
+    const result = transformKeyInfo(apiResponse);
+    expect(result).toEqual({
+      token: "sk-abc123",
+      token_id: "tok_1",
+      key_name: "my-key",
+      spend: 10.5,
+    });
+  });
+
+  it("should set the token field from the key property", () => {
+    const apiResponse = {
+      key: "sk-xyz789",
+      info: { key_name: "test" },
+    };
+    const result = transformKeyInfo(apiResponse);
+    expect(result.token).toBe("sk-xyz789");
+  });
+
+  it("should preserve all info fields in the result", () => {
+    const apiResponse = {
+      key: "sk-abc",
+      info: {
+        token_id: "tok_2",
+        key_name: "prod-key",
+        spend: 42,
+        models: ["gpt-4"],
+        team_id: "team-1",
+        metadata: { env: "production" },
+      },
+    };
+    const result = transformKeyInfo(apiResponse);
+    expect(result.token_id).toBe("tok_2");
+    expect(result.key_name).toBe("prod-key");
+    expect(result.spend).toBe(42);
+    expect(result.models).toEqual(["gpt-4"]);
+    expect(result.team_id).toBe("team-1");
+    expect(result.metadata).toEqual({ env: "production" });
+  });
+
+  it("should handle empty info object", () => {
+    const apiResponse = {
+      key: "sk-empty",
+      info: {},
+    };
+    const result = transformKeyInfo(apiResponse);
+    expect(result.token).toBe("sk-empty");
+    expect(Object.keys(result)).toContain("token");
+  });
+});

ui/litellm-dashboard/src/components/leftnav.tsx

@@ -232,8 +232,8 @@ const menuGroups: MenuGroup[] = [
     groupLabel: "DEVELOPER TOOLS",
     items: [
       {
-        key: "api_ref",
-        page: "api_ref",
+        key: "api-reference",
+        page: "api-reference",
         label: "API Reference",
         icon: <ApiOutlined />,
       },

ui/litellm-dashboard/src/components/mcp_tools/ByokCredentialModal.tsx

@@ -1,7 +1,8 @@
 "use client";
 
 import React, { useState } from "react";
-import { Modal, Input, Switch, message } from "antd";
+import { Modal, Input, Switch } from "antd";
+import MessageManager from "@/components/molecules/message_manager";
 import {
   KeyOutlined,
   LockOutlined,
@@ -46,7 +47,7 @@ export const ByokCredentialModal: React.FC<ByokCredentialModalProps> = ({
 
   const handleAuthorize = async () => {
     if (!apiKey.trim()) {
-      message.error("Please enter your API key");
+      MessageManager.error("Please enter your API key");
       return;
     }
     setLoading(true);
@@ -63,11 +64,11 @@ export const ByokCredentialModal: React.FC<ByokCredentialModalProps> = ({
         const err = await response.json();
         throw new Error(err?.detail?.error || "Failed to save credential");
       }
-      message.success(`Connected to ${serverDisplayName}`);
+      MessageManager.success(`Connected to ${serverDisplayName}`);
       onSuccess(server.server_id);
       handleClose();
     } catch (e: any) {
-      message.error(e.message || "Failed to connect");
+      MessageManager.error(e.message || "Failed to connect");
     } finally {
       setLoading(false);
     }

ui/litellm-dashboard/src/components/molecules/message_manager.tsx

@@ -0,0 +1,38 @@
+import { message as staticMessage } from "antd";
+import type { MessageInstance } from "antd/es/message/interface";
+
+let messageInstance: MessageInstance | null = null;
+
+export const setMessageInstance = (instance: MessageInstance) => {
+  messageInstance = instance;
+};
+
+const getMessageApi = () => messageInstance || staticMessage;
+
+const MessageManager = {
+  success(content: string, duration?: number) {
+    getMessageApi().success(content, duration);
+  },
+
+  error(content: string, duration?: number) {
+    getMessageApi().error(content, duration);
+  },
+
+  warning(content: string, duration?: number) {
+    getMessageApi().warning(content, duration);
+  },
+
+  info(content: string, duration?: number) {
+    getMessageApi().info(content, duration);
+  },
+
+  loading(content: string, duration?: number) {
+    return getMessageApi().loading(content, duration);
+  },
+
+  destroy() {
+    getMessageApi().destroy();
+  },
+};
+
+export default MessageManager;

ui/litellm-dashboard/src/components/navbar.tsx

@@ -4,6 +4,7 @@ import { getProxyBaseUrl } from "@/components/networking";
 import { useUIConfig } from "@/app/(dashboard)/hooks/uiConfig/useUIConfig";
 import { useTheme } from "@/contexts/ThemeContext";
 import { clearTokenCookies } from "@/utils/cookieUtils";
+import { clearStoredReturnUrl } from "@/utils/returnUrlUtils";
 import { fetchProxySettings } from "@/utils/proxyUtils";
 import { MenuFoldOutlined, MenuUnfoldOutlined, MessageOutlined, MoonOutlined, SunOutlined } from "@ant-design/icons";
 import { Button, Switch, Tag } from "antd";
@@ -12,6 +13,7 @@ import React, { useEffect, useState } from "react";
 import { BlogDropdown } from "./Navbar/BlogDropdown/BlogDropdown";
 import { CommunityEngagementButtons } from "./Navbar/CommunityEngagementButtons/CommunityEngagementButtons";
 import UserDropdown from "./Navbar/UserDropdown/UserDropdown";
+import WorkerDropdown from "./Navbar/WorkerDropdown/WorkerDropdown";
 
 interface NavbarProps {
   userID: string | null;
@@ -77,9 +79,19 @@ const Navbar: React.FC<NavbarProps> = ({
 
   const handleLogout = () => {
     clearTokenCookies();
+    localStorage.removeItem("litellm_selected_worker_id");
+    localStorage.removeItem("litellm_worker_url");
     window.location.href = logoutUrl;
   };
 
+  const handleWorkerSwitch = (workerId: string) => {
+    clearTokenCookies();
+    clearStoredReturnUrl();
+    localStorage.removeItem("litellm_selected_worker_id");
+    localStorage.removeItem("litellm_worker_url");
+    window.location.href = `/ui/login?worker=${encodeURIComponent(workerId)}`;
+  };
+
   return (
     <nav className="bg-white border-b border-gray-200 sticky top-0 z-10">
       <div className="w-full">
@@ -169,6 +181,7 @@ const Navbar: React.FC<NavbarProps> = ({
                 NEW
               </span>
             </a>
+            <WorkerDropdown onWorkerSwitch={handleWorkerSwitch} />
             <CommunityEngagementButtons />
             {/* Dark mode is currently a work in progress. To test, you can change 'false' to 'true' below.
             Do not set this to true by default until all components are confirmed to support dark mode styles. */}

ui/litellm-dashboard/src/components/networking.tsx

@@ -68,7 +68,7 @@ export const getInProductNudgesCall = async (accessToken: string) => {
 /**
  * Helper file for calls being made to proxy
  */
-import { message } from "antd";
+import MessageManager from "@/components/molecules/message_manager";
 import { clearTokenCookies } from "@/utils/cookieUtils";
 import { TagNewRequest, TagUpdateRequest, TagListResponse, TagInfoResponse } from "./tag_management/types";
 import { Team } from "./key_team_helpers/key_list";
@@ -86,7 +86,23 @@ const defaultProxyBaseUrl =
     : null;
 const defaultServerRootPath = "/";
 export let serverRootPath = defaultServerRootPath;
-export let proxyBaseUrl = defaultProxyBaseUrl;
+const WORKER_URL_KEY = "litellm_worker_url";
+// If a worker URL is in localStorage, use it as the initial proxyBaseUrl.
+// This survives page navigation and the sessionStorage.clear() in user_dashboard.
+const _rawWorkerUrl =
+  typeof window !== "undefined" ? window.localStorage.getItem(WORKER_URL_KEY) : null;
+// Validate stored worker URL — reject non-HTTP schemes to prevent exfiltration
+const _initialWorkerUrl = (() => {
+  if (!_rawWorkerUrl) return null;
+  try {
+    const parsed = new URL(_rawWorkerUrl);
+    if (parsed.protocol === "http:" || parsed.protocol === "https:") return _rawWorkerUrl;
+  } catch { /* invalid URL */ }
+  // Invalid URL in storage — clear it
+  if (typeof window !== "undefined") window.localStorage.removeItem(WORKER_URL_KEY);
+  return null;
+})();
+export let proxyBaseUrl: string | null = _initialWorkerUrl ?? defaultProxyBaseUrl;
 if (isLocal != true) {
   console.log = function () { };
 }
@@ -102,6 +118,10 @@ const updateProxyBaseUrl = (serverRootPath: string, receivedProxyBaseUrl: string
   /**
    * Special function for updating the proxy base url. Should only be called by getUiConfig.
    */
+  // If a worker URL is in localStorage, don't let getUiConfig overwrite it
+  if (typeof window !== "undefined" && window.localStorage.getItem(WORKER_URL_KEY)) {
+    return;
+  }
   const browserLocation = getWindowLocation();
   const resolvedDefaultProxyBaseUrl =
     isLocal && process.env.NEXT_PUBLIC_USE_REWRITES !== "true"
@@ -137,6 +157,36 @@ export const getProxyBaseUrl = (): string => {
   return browserLocation?.origin ?? "";
 };
 
+/**
+ * Switch API calls to point at a worker (or back to the control plane).
+ * Persists to localStorage so it survives page navigation and the
+ * sessionStorage.clear() in user_dashboard. Also updates the module-level
+ * proxyBaseUrl so in-flight code in this JS execution sees the new value
+ * immediately.
+ */
+function isValidHttpUrl(url: string): boolean {
+  try {
+    const parsed = new URL(url);
+    return parsed.protocol === "http:" || parsed.protocol === "https:";
+  } catch {
+    return false;
+  }
+}
+
+export function switchToWorkerUrl(workerUrl: string | null): void {
+  if (workerUrl && !isValidHttpUrl(workerUrl)) {
+    return;
+  }
+  if (typeof window !== "undefined") {
+    if (workerUrl) {
+      window.localStorage.setItem(WORKER_URL_KEY, workerUrl);
+    } else {
+      window.localStorage.removeItem(WORKER_URL_KEY);
+    }
+  }
+  proxyBaseUrl = workerUrl ?? defaultProxyBaseUrl;
+}
+
 const HTTP_REQUEST = {
   GET: "GET",
   POST: "POST",
@@ -262,12 +312,20 @@ interface PublicModelHubInfo {
   useful_links: Record<string, string | { url: string; index: number }>;
 }
 
+export interface WorkerInfo {
+  worker_id: string;
+  name: string;
+  url: string;
+}
+
 export interface LiteLLMWellKnownUiConfig {
   server_root_path: string;
   proxy_base_url: string | null;
   auto_redirect_to_sso: boolean;
   admin_ui_disabled: boolean;
   sso_configured: boolean;
+  is_control_plane?: boolean;
+  workers?: WorkerInfo[];
 }
 
 export interface CredentialsResponse {
@@ -555,7 +613,7 @@ export const modelCreateCall = async (accessToken: string, formValues: Model) =>
     console.log("API Response:", data);
 
     // Close any existing messages before showing new ones
-    message.destroy();
+    MessageManager.destroy();
 
     // Sequential success messages
     NotificationsManager.success(`Model ${formValues.model_name} created successfully`);
@@ -9030,15 +9088,20 @@ export const deriveErrorMessage = (errorData: any): string => {
 export interface LoginRequest {
   username: string;
   password: string;
+  useV3?: boolean;
 }
 
 interface LoginResponse {
   redirect_url: string;
+  token?: string;
+  code?: string;
+  expires_in?: number;
 }
 
-export const loginCall = async (username: string, password: string): Promise<LoginResponse> => {
+export const loginCall = async (username: string, password: string, useV3?: boolean): Promise<LoginResponse> => {
   const proxyBaseUrl = getProxyBaseUrl();
-  const loginUrl = proxyBaseUrl ? `${proxyBaseUrl}/v2/login` : "/v2/login";
+  const loginPath = useV3 ? "/v3/login" : "/v2/login";
+  const loginUrl = proxyBaseUrl ? `${proxyBaseUrl}${loginPath}` : loginPath;
 
   const body = JSON.stringify({
     username,
@@ -9060,10 +9123,65 @@ export const loginCall = async (username: string, password: string): Promise<Log
     throw new Error(errorMessage);
   }
 
-  const data = await response.json();
+  const data: LoginResponse = await response.json();
+
+  // v3 returns an opaque code — exchange it for the real JWT
+  if (useV3 && data.code) {
+    const exchangeUrl = proxyBaseUrl
+      ? `${proxyBaseUrl}/v3/login/exchange`
+      : "/v3/login/exchange";
+
+    const exchangeResponse = await fetch(exchangeUrl, {
+      method: "POST",
+      body: JSON.stringify({ code: data.code }),
+      credentials: "include",
+      headers: { "Content-Type": "application/json" },
+    });
+
+    if (!exchangeResponse.ok) {
+      const errorData = await exchangeResponse.json();
+      throw new Error(deriveErrorMessage(errorData));
+    }
+
+    const exchangeData: LoginResponse = await exchangeResponse.json();
+    if (exchangeData.token) {
+      document.cookie = `token=${exchangeData.token}; path=/; SameSite=Lax`;
+    }
+    return exchangeData;
+  }
+
+  // Backwards compatibility: v2 or old v3 returns token directly
+  if (data.token) {
+    document.cookie = `token=${data.token}; path=/; SameSite=Lax`;
+  }
+
   return data;
 };
 
+/**
+ * Exchange a single-use login code for a JWT token.
+ * Used by the SSO callback when the worker redirects back with ?code=.
+ */
+export const exchangeLoginCode = async (code: string, workerBaseUrl?: string | null): Promise<string> => {
+  const base = workerBaseUrl || getProxyBaseUrl();
+  const response = await fetch(`${base}/v3/login/exchange`, {
+    method: "POST",
+    body: JSON.stringify({ code }),
+    headers: { "Content-Type": "application/json" },
+  });
+
+  if (!response.ok) {
+    const errorData = await response.json();
+    throw new Error(deriveErrorMessage(errorData));
+  }
+
+  const data = await response.json();
+  if (data.token) {
+    document.cookie = `token=${data.token}; path=/; SameSite=Lax`;
+  }
+  return data.token;
+};
+
 export const getUiSettings = async () => {
   const proxyBaseUrl = getProxyBaseUrl();
   const url = proxyBaseUrl ? `${proxyBaseUrl}/get/ui_settings` : `/get/ui_settings`;

ui/litellm-dashboard/src/components/organisms/create_key_button.tsx

@@ -8,7 +8,7 @@ import { formatNumberWithCommas } from "@/utils/dataUtils";
 import { InfoCircleOutlined } from "@ant-design/icons";
 import { useQueryClient } from "@tanstack/react-query";
 import { Accordion, AccordionBody, AccordionHeader, Button, Col, Grid, Text, TextInput, Title } from "@tremor/react";
-import { Button as Button2, Form, Input, message, Modal, Radio, Select, Switch, Tag, Tooltip } from "antd";
+import { Button as Button2, Form, Input, Modal, Radio, Select, Switch, Tag, Tooltip } from "antd";
 import debounce from "lodash/debounce";
 import React, { useCallback, useEffect, useState } from "react";
 import { rolesWithWriteAccess } from "../../utils/roles";

ui/litellm-dashboard/src/components/page_metadata.ts

@@ -24,7 +24,7 @@ export const pageDescriptions: Record<string, string> = {
   projects: "Manage projects within teams",
   "access-groups": "Manage access groups for role-based permissions",
   budgets: "Set and monitor spending budgets",
-  api_ref: "Browse API documentation and endpoints",
+  "api-reference": "Browse API documentation and endpoints",
   "model-hub-table": "Explore available AI models and providers",
   "learning-resources": "Access tutorials and documentation",
   caching: "Configure response caching settings",

ui/litellm-dashboard/src/components/playground/chat_ui/ChatUI.tsx

@@ -28,7 +28,6 @@ import ReactMarkdown from "react-markdown";
 import { Prism as SyntaxHighlighter } from "react-syntax-highlighter";
 import { coy } from "react-syntax-highlighter/dist/esm/styles/prism";
 import { v4 as uuidv4 } from "uuid";
-import { truncateString } from "../../../utils/textUtils";
 import GuardrailSelector from "../../guardrails/GuardrailSelector";
 import PolicySelector from "../../policies/PolicySelector";
 import MCPToolArgumentsForm, { MCPToolArgumentsFormRef } from "../../mcp_tools/MCPToolArgumentsForm";
@@ -61,9 +60,8 @@ import CodeInterpreterTool from "./CodeInterpreterTool";
 import { generateCodeSnippet } from "./CodeSnippets";
 import EndpointSelector from "./EndpointSelector";
 import FilePreviewCard from "./FilePreviewCard";
-import MCPEventsDisplay from "./MCPEventsDisplay";
-import type { MCPEvent } from "../../mcp_tools/types";
 import ChatMessageBubble from "./ChatMessageBubble";
+import MCPEventsDisplay from "./MCPEventsDisplay";
 import { EndpointType, getEndpointType } from "./mode_endpoint_mapping";
 import ReasoningContent from "./ReasoningContent";
 import ResponseMetrics, { TokenUsage } from "./ResponseMetrics";
@@ -75,6 +73,7 @@ import SessionManagement from "./SessionManagement";
 import RealtimePlayground from "./RealtimePlayground";
 import { A2ATaskMetadata, MessageType } from "./types";
 import { useCodeInterpreter } from "./useCodeInterpreter";
+import { useChatHistory } from "./useChatHistory";
 
 const { TextArea } = Input;
 const { Dragger } = Upload;
@@ -135,6 +134,34 @@ const ChatUI: React.FC<ChatUIProps> = ({
       return {};
     }
   });
+  const {
+    chatHistory,
+    setChatHistory,
+    mcpEvents,
+    setMCPEvents,
+    messageTraceId,
+    setMessageTraceId,
+    responsesSessionId,
+    setResponsesSessionId,
+    useApiSessionManagement,
+    setUseApiSessionManagement,
+    updateTextUI,
+    updateReasoningContent,
+    updateTimingData,
+    updateUsageData,
+    updateA2AMetadata,
+    updateTotalLatency,
+    updateSearchResults,
+    handleResponseId,
+    handleToggleSessionManagement,
+    handleMCPEvent,
+    updateImageUI,
+    updateEmbeddingsUI,
+    updateAudioUI,
+    updateChatImageUI,
+    clearChatHistory: clearChatHistoryHook,
+    clearMCPEvents,
+  } = useChatHistory({ simplified });
   const [apiKeySource, setApiKeySource] = useState<"session" | "custom">(() => {
     const saved = sessionStorage.getItem("apiKeySource");
     if (saved) {
@@ -151,16 +178,6 @@ const ChatUI: React.FC<ChatUIProps> = ({
     () => sessionStorage.getItem("customProxyBaseUrl") || "",
   );
   const [inputMessage, setInputMessage] = useState("");
-  const [chatHistory, setChatHistory] = useState<MessageType[]>(() => {
-    if (simplified) return [];
-    try {
-      const saved = sessionStorage.getItem("chatHistory");
-      return saved ? JSON.parse(saved) : [];
-    } catch (error) {
-      console.error("Error parsing chatHistory from sessionStorage", error);
-      return [];
-    }
-  });
   const [selectedModel, setSelectedModel] = useState<string | undefined>(simplified ? fixedModel : undefined);
   const [showCustomModelInput, setShowCustomModelInput] = useState<boolean>(false);
   const [modelInfo, setModelInfo] = useState<ModelGroup[]>([]);
@@ -218,16 +235,6 @@ const ChatUI: React.FC<ChatUIProps> = ({
       return [];
     }
   });
-  const [messageTraceId, setMessageTraceId] = useState<string | null>(
-    () => sessionStorage.getItem("messageTraceId") || null,
-  );
-  const [responsesSessionId, setResponsesSessionId] = useState<string | null>(
-    () => sessionStorage.getItem("responsesSessionId") || null,
-  );
-  const [useApiSessionManagement, setUseApiSessionManagement] = useState<boolean>(() => {
-    const saved = sessionStorage.getItem("useApiSessionManagement");
-    return saved ? JSON.parse(saved) : true; // Default to API session management
-  });
   const [uploadedImages, setUploadedImages] = useState<File[]>([]);
   const [imagePreviewUrls, setImagePreviewUrls] = useState<string[]>([]);
   const [responsesUploadedImage, setResponsesUploadedImage] = useState<File | null>(null);
@@ -238,7 +245,6 @@ const ChatUI: React.FC<ChatUIProps> = ({
   const [isGetCodeModalVisible, setIsGetCodeModalVisible] = useState(false);
   const [generatedCode, setGeneratedCode] = useState("");
   const [selectedSdk, setSelectedSdk] = useState<"openai" | "azure">("openai");
-  const [mcpEvents, setMCPEvents] = useState<MCPEvent[]>([]);
   const [temperature, setTemperature] = useState<number>(1.0);
   const [maxTokens, setMaxTokens] = useState<number>(2048);
   const [useAdvancedParams, setUseAdvancedParams] = useState<boolean>(false);
@@ -332,17 +338,6 @@ const ChatUI: React.FC<ChatUIProps> = ({
     proxySettings,
   ]);
 
-  useEffect(() => {
-    if (simplified) return; // Do not persist chat history in simplified (embedded) mode
-    const handler = setTimeout(() => {
-      sessionStorage.setItem("chatHistory", JSON.stringify(chatHistory));
-    }, 500); // Debounce by 500ms
-
-    return () => {
-      clearTimeout(handler);
-    };
-  }, [chatHistory, simplified]);
-
   useEffect(() => {
     sessionStorage.setItem("apiKeySource", JSON.stringify(apiKeySource));
     sessionStorage.setItem("apiKey", apiKey);
@@ -363,17 +358,6 @@ const ChatUI: React.FC<ChatUIProps> = ({
         sessionStorage.removeItem("selectedModel");
       }
     }
-    if (messageTraceId) {
-      sessionStorage.setItem("messageTraceId", messageTraceId);
-    } else {
-      sessionStorage.removeItem("messageTraceId");
-    }
-    if (responsesSessionId) {
-      sessionStorage.setItem("responsesSessionId", responsesSessionId);
-    } else {
-      sessionStorage.removeItem("responsesSessionId");
-    }
-    sessionStorage.setItem("useApiSessionManagement", JSON.stringify(useApiSessionManagement));
     // Note: codeInterpreterEnabled and selectedContainerId are persisted by useCodeInterpreter hook
   }, [
     simplified,
@@ -385,9 +369,6 @@ const ChatUI: React.FC<ChatUIProps> = ({
     selectedVectorStores,
     selectedGuardrails,
     selectedPolicies,
-    messageTraceId,
-    responsesSessionId,
-    useApiSessionManagement,
     selectedMCPServers,
     mcpServerToolRestrictions,
     selectedVoice,
@@ -479,264 +460,6 @@ const ChatUI: React.FC<ChatUIProps> = ({
     }
   }, [chatHistory]);
 
-  const updateTextUI = (role: string, chunk: string, model?: string) => {
-    console.log("updateTextUI called with:", role, chunk, model);
-    setChatHistory((prev) => {
-      const last = prev[prev.length - 1];
-      // if the last message is already from this same role, append
-      if (last && last.role === role && !last.isImage && !last.isAudio) {
-        // build a new object, but only set `model` if it wasn't there already
-        const updated: MessageType = {
-          ...last,
-          content: last.content + chunk,
-          model: last.model ?? model, // ← only use the passed‐in model on the first chunk
-        };
-        return [...prev.slice(0, -1), updated];
-      } else {
-        // otherwise start a brand new assistant bubble
-        return [
-          ...prev,
-          {
-            role,
-            content: chunk,
-            model, // model set exactly once here
-          },
-        ];
-      }
-    });
-  };
-
-  const updateReasoningContent = (chunk: string) => {
-    setChatHistory((prevHistory) => {
-      const lastMessage = prevHistory[prevHistory.length - 1];
-
-      if (lastMessage && lastMessage.role === "assistant" && !lastMessage.isImage && !lastMessage.isAudio) {
-        return [
-          ...prevHistory.slice(0, prevHistory.length - 1),
-          {
-            ...lastMessage,
-            reasoningContent: (lastMessage.reasoningContent || "") + chunk,
-          },
-        ];
-      } else {
-        // If there's no assistant message yet, we'll create one with empty content
-        // but with reasoning content
-        if (prevHistory.length > 0 && prevHistory[prevHistory.length - 1].role === "user") {
-          return [
-            ...prevHistory,
-            {
-              role: "assistant",
-              content: "",
-              reasoningContent: chunk,
-            },
-          ];
-        }
-
-        return prevHistory;
-      }
-    });
-  };
-
-  const updateTimingData = (timeToFirstToken: number) => {
-    console.log("updateTimingData called with:", timeToFirstToken);
-    setChatHistory((prevHistory) => {
-      const lastMessage = prevHistory[prevHistory.length - 1];
-      console.log("Current last message:", lastMessage);
-
-      if (lastMessage && lastMessage.role === "assistant") {
-        console.log("Updating assistant message with timeToFirstToken:", timeToFirstToken);
-        const updatedHistory = [
-          ...prevHistory.slice(0, prevHistory.length - 1),
-          {
-            ...lastMessage,
-            timeToFirstToken,
-          },
-        ];
-        console.log("Updated chat history:", updatedHistory);
-        return updatedHistory;
-      }
-      // If the last message is a user message and no assistant message exists yet,
-      // create a new assistant message with empty content
-      else if (lastMessage && lastMessage.role === "user") {
-        console.log("Creating new assistant message with timeToFirstToken:", timeToFirstToken);
-        return [
-          ...prevHistory,
-          {
-            role: "assistant",
-            content: "",
-            timeToFirstToken,
-          },
-        ];
-      }
-
-      console.log("No appropriate message found to update timing");
-      return prevHistory;
-    });
-  };
-
-  const updateUsageData = (usage: TokenUsage, toolName?: string) => {
-    console.log("Received usage data:", usage);
-    setChatHistory((prevHistory) => {
-      const lastMessage = prevHistory[prevHistory.length - 1];
-
-      if (lastMessage && lastMessage.role === "assistant") {
-        console.log("Updating message with usage data:", usage);
-        const updatedMessage = {
-          ...lastMessage,
-          usage,
-          toolName,
-        };
-        console.log("Updated message:", updatedMessage);
-
-        return [...prevHistory.slice(0, prevHistory.length - 1), updatedMessage];
-      }
-
-      return prevHistory;
-    });
-  };
-
-  const updateA2AMetadata = (a2aMetadata: A2ATaskMetadata) => {
-    console.log("Received A2A metadata:", a2aMetadata);
-    setChatHistory((prevHistory) => {
-      const lastMessage = prevHistory[prevHistory.length - 1];
-
-      if (lastMessage && lastMessage.role === "assistant") {
-        const updatedMessage = {
-          ...lastMessage,
-          a2aMetadata,
-        };
-        return [...prevHistory.slice(0, prevHistory.length - 1), updatedMessage];
-      }
-
-      return prevHistory;
-    });
-  };
-
-  const updateTotalLatency = (totalLatency: number) => {
-    setChatHistory((prevHistory) => {
-      const lastMessage = prevHistory[prevHistory.length - 1];
-
-      if (lastMessage && lastMessage.role === "assistant") {
-        return [
-          ...prevHistory.slice(0, prevHistory.length - 1),
-          {
-            ...lastMessage,
-            totalLatency,
-          },
-        ];
-      }
-
-      return prevHistory;
-    });
-  };
-
-  const updateSearchResults = (searchResults: any[]) => {
-    console.log("Received search results:", searchResults);
-    setChatHistory((prevHistory) => {
-      const lastMessage = prevHistory[prevHistory.length - 1];
-
-      if (lastMessage && lastMessage.role === "assistant") {
-        console.log("Updating message with search results");
-        const updatedMessage = {
-          ...lastMessage,
-          searchResults,
-        };
-
-        return [...prevHistory.slice(0, prevHistory.length - 1), updatedMessage];
-      }
-
-      return prevHistory;
-    });
-  };
-
-  const handleResponseId = (responseId: string) => {
-    console.log("Received response ID for session management:", responseId);
-    if (useApiSessionManagement) {
-      setResponsesSessionId(responseId);
-    }
-  };
-
-  const handleToggleSessionManagement = (useApi: boolean) => {
-    setUseApiSessionManagement(useApi);
-    if (!useApi) {
-      // Clear API session when switching to UI mode
-      setResponsesSessionId(null);
-    }
-  };
-
-  const handleMCPEvent = (event: MCPEvent) => {
-    console.log("ChatUI: Received MCP event:", event);
-    setMCPEvents((prev) => {
-      // Check if this is a duplicate event (same item_id and type)
-      // Only check for duplicates if item_id is defined (for mcp_list_tools, item_id is "mcp_list_tools")
-      const isDuplicate = event.item_id
-        ? prev.some(
-            (existingEvent) =>
-              existingEvent.item_id === event.item_id &&
-              existingEvent.type === event.type &&
-              (existingEvent.sequence_number === event.sequence_number ||
-                (existingEvent.sequence_number === undefined && event.sequence_number === undefined)),
-          )
-        : false;
-
-      if (isDuplicate) {
-        console.log("ChatUI: Duplicate MCP event, skipping");
-        return prev;
-      }
-
-      const newEvents = [...prev, event];
-      console.log("ChatUI: Updated MCP events:", newEvents);
-      return newEvents;
-    });
-  };
-
-  const updateImageUI = (imageUrl: string, model: string) => {
-    setChatHistory((prevHistory) => [...prevHistory, { role: "assistant", content: imageUrl, model, isImage: true }]);
-  };
-
-  const updateEmbeddingsUI = (embeddings: string, model?: string) => {
-    setChatHistory((prevHistory) => [
-      ...prevHistory,
-      { role: "assistant", content: truncateString(embeddings, 100), model, isEmbeddings: true },
-    ]);
-  };
-
-  const updateAudioUI = (audioUrl: string, model: string) => {
-    setChatHistory((prevHistory) => [...prevHistory, { role: "assistant", content: audioUrl, model, isAudio: true }]);
-  };
-
-  const updateChatImageUI = (imageUrl: string, model?: string) => {
-    setChatHistory((prev) => {
-      const last = prev[prev.length - 1];
-      // If the last message is from assistant and has content, add image to it
-      if (last && last.role === "assistant" && !last.isImage && !last.isAudio) {
-        const updated = {
-          ...last,
-          image: {
-            url: imageUrl,
-            detail: "auto",
-          },
-          model: last.model ?? model,
-        };
-        return [...prev.slice(0, -1), updated];
-      } else {
-        // Otherwise create a new assistant message with just the image
-        return [
-          ...prev,
-          {
-            role: "assistant",
-            content: "",
-            model,
-            image: {
-              url: imageUrl,
-              detail: "auto",
-            },
-          },
-        ];
-      }
-    });
-  };
-
   const handleKeyDown = (event: React.KeyboardEvent<HTMLTextAreaElement>) => {
     if (event.key === "Enter" && !event.shiftKey) {
       event.preventDefault(); // Prevent default to avoid newline
@@ -967,7 +690,7 @@ const ChatUI: React.FC<ChatUIProps> = ({
     }
 
     setChatHistory([...chatHistory, displayMessage]);
-    setMCPEvents([]); // Clear previous MCP events for new conversation turn
+    clearMCPEvents(); // Clear previous MCP events for new conversation turn
     codeInterpreter.clearResult(); // Clear previous code interpreter results
     setIsLoading(true);
 
@@ -1223,26 +946,11 @@ const ChatUI: React.FC<ChatUIProps> = ({
   };
 
   const clearChatHistory = () => {
-    // Clean up audio object URLs before clearing history
-    chatHistory.forEach((message) => {
-      if (message.isAudio && typeof message.content === "string") {
-        URL.revokeObjectURL(message.content);
-      }
-    });
-
-    setChatHistory([]);
-    setMessageTraceId(null);
-    setResponsesSessionId(null); // Clear responses session ID
-    setMCPEvents([]); // Clear MCP events
-    handleRemoveAllImages(); // Clear any uploaded images for image edits
-    handleRemoveResponsesImage(); // Clear any uploaded images for responses
-    handleRemoveChatImage(); // Clear any uploaded images for chat completions
-    handleRemoveAudio(); // Clear any uploaded audio for transcription
-    if (!simplified) {
-      sessionStorage.removeItem("chatHistory");
-      sessionStorage.removeItem("messageTraceId");
-      sessionStorage.removeItem("responsesSessionId");
-    }
+    clearChatHistoryHook();
+    handleRemoveAllImages();
+    handleRemoveResponsesImage();
+    handleRemoveChatImage();
+    handleRemoveAudio();
     NotificationsManager.success("Chat history cleared.");
   };
 

ui/litellm-dashboard/src/components/playground/chat_ui/CodeInterpreterTool.tsx

@@ -1,5 +1,6 @@
 import React from "react";
-import { Switch, Tooltip, message } from "antd";
+import { Switch, Tooltip } from "antd";
+import MessageManager from "@/components/molecules/message_manager";
 import { CodeOutlined, InfoCircleOutlined, ExclamationCircleOutlined } from "@ant-design/icons";
 import { Text } from "@tremor/react";
 
@@ -38,7 +39,7 @@ const CodeInterpreterTool: React.FC<CodeInterpreterToolProps> = ({
 
   const handleToggle = (checked: boolean) => {
     if (checked && !isOpenAI) {
-      message.warning("Code Interpreter is only available for OpenAI models");
+      MessageManager.warning("Code Interpreter is only available for OpenAI models");
       return;
     }
     onEnabledChange(checked);

ui/litellm-dashboard/src/components/playground/chat_ui/useChatHistory.test.ts

@@ -0,0 +1,591 @@
+import { renderHook, act } from "@testing-library/react";
+import { describe, it, expect, beforeEach, vi } from "vitest";
+import { useChatHistory } from "./useChatHistory";
+
+describe("useChatHistory", () => {
+  beforeEach(() => {
+    sessionStorage.clear();
+  });
+
+  describe("updateTextUI", () => {
+    it("should create a new assistant message when chat is empty", () => {
+      const { result } = renderHook(() => useChatHistory({ simplified: false }));
+
+      act(() => {
+        result.current.updateTextUI("assistant", "Hello", "gpt-4");
+      });
+
+      expect(result.current.chatHistory).toEqual([
+        { role: "assistant", content: "Hello", model: "gpt-4" },
+      ]);
+    });
+
+    it("should append to the last assistant message", () => {
+      const { result } = renderHook(() => useChatHistory({ simplified: false }));
+
+      act(() => {
+        result.current.updateTextUI("assistant", "Hello", "gpt-4");
+      });
+      act(() => {
+        result.current.updateTextUI("assistant", " world");
+      });
+
+      expect(result.current.chatHistory).toEqual([
+        { role: "assistant", content: "Hello world", model: "gpt-4" },
+      ]);
+    });
+
+    it("should not overwrite model on subsequent chunks", () => {
+      const { result } = renderHook(() => useChatHistory({ simplified: false }));
+
+      act(() => {
+        result.current.updateTextUI("assistant", "Hello", "gpt-4");
+      });
+      act(() => {
+        result.current.updateTextUI("assistant", " world", "gpt-3.5");
+      });
+
+      expect(result.current.chatHistory[0].model).toBe("gpt-4");
+    });
+
+    it("should create a new message when role changes", () => {
+      const { result } = renderHook(() => useChatHistory({ simplified: false }));
+
+      act(() => {
+        result.current.updateTextUI("user", "Hi");
+      });
+      act(() => {
+        result.current.updateTextUI("assistant", "Hello", "gpt-4");
+      });
+
+      expect(result.current.chatHistory).toHaveLength(2);
+      expect(result.current.chatHistory[0].role).toBe("user");
+      expect(result.current.chatHistory[1].role).toBe("assistant");
+    });
+
+    it("should not append to image messages", () => {
+      const { result } = renderHook(() => useChatHistory({ simplified: false }));
+
+      act(() => {
+        result.current.updateImageUI("http://img.png", "dall-e");
+      });
+      act(() => {
+        result.current.updateTextUI("assistant", "description", "gpt-4");
+      });
+
+      expect(result.current.chatHistory).toHaveLength(2);
+    });
+
+    it("should not append to audio messages", () => {
+      const { result } = renderHook(() => useChatHistory({ simplified: false }));
+
+      act(() => {
+        result.current.updateAudioUI("http://audio.mp3", "tts-1");
+      });
+      act(() => {
+        result.current.updateTextUI("assistant", "text", "gpt-4");
+      });
+
+      expect(result.current.chatHistory).toHaveLength(2);
+    });
+  });
+
+  describe("updateReasoningContent", () => {
+    it("should add reasoning content to existing assistant message", () => {
+      const { result } = renderHook(() => useChatHistory({ simplified: false }));
+
+      act(() => {
+        result.current.updateTextUI("assistant", "Answer", "gpt-4");
+      });
+      act(() => {
+        result.current.updateReasoningContent("thinking...");
+      });
+
+      expect(result.current.chatHistory[0].reasoningContent).toBe("thinking...");
+    });
+
+    it("should append reasoning content across chunks", () => {
+      const { result } = renderHook(() => useChatHistory({ simplified: false }));
+
+      act(() => {
+        result.current.updateTextUI("assistant", "", "gpt-4");
+      });
+      act(() => {
+        result.current.updateReasoningContent("step 1");
+      });
+      act(() => {
+        result.current.updateReasoningContent(" step 2");
+      });
+
+      expect(result.current.chatHistory[0].reasoningContent).toBe("step 1 step 2");
+    });
+
+    it("should create assistant message with reasoning when last message is user", () => {
+      const { result } = renderHook(() => useChatHistory({ simplified: false }));
+
+      act(() => {
+        result.current.setChatHistory([{ role: "user", content: "question" }]);
+      });
+      act(() => {
+        result.current.updateReasoningContent("thinking...");
+      });
+
+      expect(result.current.chatHistory).toHaveLength(2);
+      expect(result.current.chatHistory[1]).toEqual({
+        role: "assistant",
+        content: "",
+        reasoningContent: "thinking...",
+      });
+    });
+
+    it("should not update when chat is empty", () => {
+      const { result } = renderHook(() => useChatHistory({ simplified: false }));
+
+      act(() => {
+        result.current.updateReasoningContent("thinking...");
+      });
+
+      expect(result.current.chatHistory).toHaveLength(0);
+    });
+  });
+
+  describe("updateTimingData", () => {
+    it("should add timeToFirstToken to existing assistant message", () => {
+      const { result } = renderHook(() => useChatHistory({ simplified: false }));
+
+      act(() => {
+        result.current.updateTextUI("assistant", "Hello", "gpt-4");
+      });
+      act(() => {
+        result.current.updateTimingData(150);
+      });
+
+      expect(result.current.chatHistory[0].timeToFirstToken).toBe(150);
+    });
+
+    it("should create assistant message when last is user", () => {
+      const { result } = renderHook(() => useChatHistory({ simplified: false }));
+
+      act(() => {
+        result.current.setChatHistory([{ role: "user", content: "hi" }]);
+      });
+      act(() => {
+        result.current.updateTimingData(200);
+      });
+
+      expect(result.current.chatHistory).toHaveLength(2);
+      expect(result.current.chatHistory[1].timeToFirstToken).toBe(200);
+    });
+  });
+
+  describe("updateUsageData", () => {
+    it("should add usage data to assistant message", () => {
+      const { result } = renderHook(() => useChatHistory({ simplified: false }));
+
+      act(() => {
+        result.current.updateTextUI("assistant", "Hello", "gpt-4");
+      });
+
+      const usage = { completionTokens: 10, promptTokens: 5, totalTokens: 15 };
+      act(() => {
+        result.current.updateUsageData(usage);
+      });
+
+      expect(result.current.chatHistory[0].usage).toEqual(usage);
+    });
+
+    it("should add toolName when provided", () => {
+      const { result } = renderHook(() => useChatHistory({ simplified: false }));
+
+      act(() => {
+        result.current.updateTextUI("assistant", "Hello", "gpt-4");
+      });
+
+      const usage = { completionTokens: 10, promptTokens: 5, totalTokens: 15 };
+      act(() => {
+        result.current.updateUsageData(usage, "search_tool");
+      });
+
+      expect(result.current.chatHistory[0].toolName).toBe("search_tool");
+    });
+  });
+
+  describe("updateTotalLatency", () => {
+    it("should add totalLatency to assistant message", () => {
+      const { result } = renderHook(() => useChatHistory({ simplified: false }));
+
+      act(() => {
+        result.current.updateTextUI("assistant", "Hello", "gpt-4");
+      });
+      act(() => {
+        result.current.updateTotalLatency(500);
+      });
+
+      expect(result.current.chatHistory[0].totalLatency).toBe(500);
+    });
+  });
+
+  describe("updateA2AMetadata", () => {
+    it("should add A2A metadata to assistant message", () => {
+      const { result } = renderHook(() => useChatHistory({ simplified: false }));
+
+      act(() => {
+        result.current.updateTextUI("assistant", "Hello", "gpt-4");
+      });
+
+      const metadata = { taskId: "task-1", contextId: "ctx-1" };
+      act(() => {
+        result.current.updateA2AMetadata(metadata);
+      });
+
+      expect(result.current.chatHistory[0].a2aMetadata).toEqual(metadata);
+    });
+  });
+
+  describe("updateSearchResults", () => {
+    it("should add search results to assistant message", () => {
+      const { result } = renderHook(() => useChatHistory({ simplified: false }));
+
+      act(() => {
+        result.current.updateTextUI("assistant", "Hello", "gpt-4");
+      });
+
+      const searchResults = [{ object: "search", search_query: "test", data: [] }];
+      act(() => {
+        result.current.updateSearchResults(searchResults);
+      });
+
+      expect(result.current.chatHistory[0].searchResults).toEqual(searchResults);
+    });
+  });
+
+  describe("updateImageUI", () => {
+    it("should add image message to history", () => {
+      const { result } = renderHook(() => useChatHistory({ simplified: false }));
+
+      act(() => {
+        result.current.updateImageUI("http://img.png", "dall-e-3");
+      });
+
+      expect(result.current.chatHistory).toEqual([
+        { role: "assistant", content: "http://img.png", model: "dall-e-3", isImage: true },
+      ]);
+    });
+  });
+
+  describe("updateEmbeddingsUI", () => {
+    it("should add truncated embeddings message", () => {
+      const { result } = renderHook(() => useChatHistory({ simplified: false }));
+
+      act(() => {
+        result.current.updateEmbeddingsUI("[0.1, 0.2, 0.3]", "text-embedding-ada");
+      });
+
+      expect(result.current.chatHistory[0].isEmbeddings).toBe(true);
+      expect(result.current.chatHistory[0].model).toBe("text-embedding-ada");
+    });
+  });
+
+  describe("updateAudioUI", () => {
+    it("should add audio message to history", () => {
+      const { result } = renderHook(() => useChatHistory({ simplified: false }));
+
+      act(() => {
+        result.current.updateAudioUI("http://audio.mp3", "tts-1");
+      });
+
+      expect(result.current.chatHistory).toEqual([
+        { role: "assistant", content: "http://audio.mp3", model: "tts-1", isAudio: true },
+      ]);
+    });
+  });
+
+  describe("updateChatImageUI", () => {
+    it("should add image to existing assistant message", () => {
+      const { result } = renderHook(() => useChatHistory({ simplified: false }));
+
+      act(() => {
+        result.current.updateTextUI("assistant", "Here is the image", "gpt-4");
+      });
+      act(() => {
+        result.current.updateChatImageUI("http://img.png", "gpt-4");
+      });
+
+      expect(result.current.chatHistory[0].image).toEqual({
+        url: "http://img.png",
+        detail: "auto",
+      });
+    });
+
+    it("should create new assistant message with image when no assistant message exists", () => {
+      const { result } = renderHook(() => useChatHistory({ simplified: false }));
+
+      act(() => {
+        result.current.updateChatImageUI("http://img.png", "gpt-4");
+      });
+
+      expect(result.current.chatHistory[0]).toEqual({
+        role: "assistant",
+        content: "",
+        model: "gpt-4",
+        image: { url: "http://img.png", detail: "auto" },
+      });
+    });
+  });
+
+  describe("handleMCPEvent", () => {
+    it("should add MCP event", () => {
+      const { result } = renderHook(() => useChatHistory({ simplified: false }));
+
+      act(() => {
+        result.current.handleMCPEvent({ type: "tool_call", item_id: "1" });
+      });
+
+      expect(result.current.mcpEvents).toHaveLength(1);
+    });
+
+    it("should deduplicate events by item_id and type", () => {
+      const { result } = renderHook(() => useChatHistory({ simplified: false }));
+
+      const event = { type: "tool_call", item_id: "1" };
+      act(() => {
+        result.current.handleMCPEvent(event);
+      });
+      act(() => {
+        result.current.handleMCPEvent(event);
+      });
+
+      expect(result.current.mcpEvents).toHaveLength(1);
+    });
+
+    it("should allow events without item_id (no dedup)", () => {
+      const { result } = renderHook(() => useChatHistory({ simplified: false }));
+
+      act(() => {
+        result.current.handleMCPEvent({ type: "tool_call" });
+      });
+      act(() => {
+        result.current.handleMCPEvent({ type: "tool_call" });
+      });
+
+      expect(result.current.mcpEvents).toHaveLength(2);
+    });
+
+    it("should allow events with same item_id/type but different sequence_number", () => {
+      const { result } = renderHook(() => useChatHistory({ simplified: false }));
+
+      act(() => {
+        result.current.handleMCPEvent({ type: "tool_call", item_id: "1", sequence_number: 1 });
+      });
+      act(() => {
+        result.current.handleMCPEvent({ type: "tool_call", item_id: "1", sequence_number: 2 });
+      });
+
+      expect(result.current.mcpEvents).toHaveLength(2);
+    });
+  });
+
+  describe("clearMCPEvents", () => {
+    it("should clear MCP events without affecting chat history", () => {
+      const { result } = renderHook(() => useChatHistory({ simplified: false }));
+
+      act(() => {
+        result.current.updateTextUI("assistant", "Hello", "gpt-4");
+        result.current.handleMCPEvent({ type: "tool_call", item_id: "1" });
+      });
+      act(() => {
+        result.current.clearMCPEvents();
+      });
+
+      expect(result.current.mcpEvents).toEqual([]);
+      expect(result.current.chatHistory).toHaveLength(1);
+    });
+  });
+
+  describe("clearChatHistory", () => {
+    it("should clear all state", () => {
+      const { result } = renderHook(() => useChatHistory({ simplified: false }));
+
+      act(() => {
+        result.current.updateTextUI("assistant", "Hello", "gpt-4");
+        result.current.handleMCPEvent({ type: "tool_call", item_id: "1" });
+      });
+      act(() => {
+        result.current.clearChatHistory();
+      });
+
+      expect(result.current.chatHistory).toEqual([]);
+      expect(result.current.mcpEvents).toEqual([]);
+      expect(result.current.messageTraceId).toBeNull();
+      expect(result.current.responsesSessionId).toBeNull();
+    });
+
+    it("should revoke audio object URLs when clearing", () => {
+      const revokeSpy = vi.spyOn(URL, "revokeObjectURL").mockImplementation(() => {});
+      const { result } = renderHook(() => useChatHistory({ simplified: false }));
+
+      act(() => {
+        result.current.updateAudioUI("blob:http://localhost/audio-1", "tts-1");
+      });
+      act(() => {
+        result.current.clearChatHistory();
+      });
+
+      expect(revokeSpy).toHaveBeenCalledWith("blob:http://localhost/audio-1");
+      revokeSpy.mockRestore();
+    });
+
+    it("should clear sessionStorage when not simplified", () => {
+      vi.useFakeTimers();
+      const { result } = renderHook(() => useChatHistory({ simplified: false }));
+
+      sessionStorage.setItem("chatHistory", "[]");
+      sessionStorage.setItem("messageTraceId", "trace-1");
+      sessionStorage.setItem("responsesSessionId", "resp-1");
+
+      act(() => {
+        result.current.clearChatHistory();
+      });
+
+      // Advance past the 500ms debounce to verify it does not re-write the key
+      act(() => {
+        vi.advanceTimersByTime(600);
+      });
+
+      expect(sessionStorage.getItem("chatHistory")).toBeNull();
+      expect(sessionStorage.getItem("messageTraceId")).toBeNull();
+      expect(sessionStorage.getItem("responsesSessionId")).toBeNull();
+
+      vi.useRealTimers();
+    });
+
+    it("should NOT clear sessionStorage when simplified", () => {
+      sessionStorage.setItem("chatHistory", '[{"role":"user","content":"hi"}]');
+
+      const { result } = renderHook(() => useChatHistory({ simplified: true }));
+
+      act(() => {
+        result.current.clearChatHistory();
+      });
+
+      // simplified mode should not touch sessionStorage
+      expect(sessionStorage.getItem("chatHistory")).toBe('[{"role":"user","content":"hi"}]');
+    });
+
+    it("should not re-write chatHistory to sessionStorage after clear via debounce", () => {
+      vi.useFakeTimers();
+      const { result } = renderHook(() => useChatHistory({ simplified: false }));
+
+      // Add a message so the debounce has something to persist
+      act(() => {
+        result.current.updateTextUI("assistant", "Hello", "gpt-4");
+      });
+
+      // Let the debounce fire so the message is persisted
+      act(() => {
+        vi.advanceTimersByTime(600);
+      });
+      expect(sessionStorage.getItem("chatHistory")).not.toBeNull();
+
+      // Now clear
+      act(() => {
+        result.current.clearChatHistory();
+      });
+
+      // Advance past the debounce — the key should stay removed
+      act(() => {
+        vi.advanceTimersByTime(600);
+      });
+
+      expect(sessionStorage.getItem("chatHistory")).toBeNull();
+
+      vi.useRealTimers();
+    });
+  });
+
+  describe("simplified mode session isolation", () => {
+    it("should not hydrate messageTraceId from sessionStorage in simplified mode", () => {
+      sessionStorage.setItem("messageTraceId", "trace-from-playground");
+
+      const { result } = renderHook(() => useChatHistory({ simplified: true }));
+
+      expect(result.current.messageTraceId).toBeNull();
+    });
+
+    it("should not hydrate responsesSessionId from sessionStorage in simplified mode", () => {
+      sessionStorage.setItem("responsesSessionId", "resp-from-playground");
+
+      const { result } = renderHook(() => useChatHistory({ simplified: true }));
+
+      expect(result.current.responsesSessionId).toBeNull();
+    });
+
+    it("should not hydrate useApiSessionManagement from sessionStorage in simplified mode", () => {
+      sessionStorage.setItem("useApiSessionManagement", "false");
+
+      const { result } = renderHook(() => useChatHistory({ simplified: true }));
+
+      // Should get the default (true), not the stored value
+      expect(result.current.useApiSessionManagement).toBe(true);
+    });
+
+    it("should not persist session state to sessionStorage in simplified mode", () => {
+      vi.useFakeTimers();
+      const { result } = renderHook(() => useChatHistory({ simplified: true }));
+
+      act(() => {
+        result.current.setMessageTraceId("trace-embedded");
+        result.current.setResponsesSessionId("resp-embedded");
+      });
+
+      // Flush effects
+      act(() => {
+        vi.advanceTimersByTime(0);
+      });
+
+      expect(sessionStorage.getItem("messageTraceId")).toBeNull();
+      expect(sessionStorage.getItem("responsesSessionId")).toBeNull();
+
+      vi.useRealTimers();
+    });
+  });
+
+  describe("session management", () => {
+    it("handleResponseId should set responsesSessionId when useApiSessionManagement is true", () => {
+      const { result } = renderHook(() => useChatHistory({ simplified: false }));
+
+      act(() => {
+        result.current.handleResponseId("resp-123");
+      });
+
+      expect(result.current.responsesSessionId).toBe("resp-123");
+    });
+
+    it("handleResponseId should NOT set responsesSessionId when useApiSessionManagement is false", () => {
+      const { result } = renderHook(() => useChatHistory({ simplified: false }));
+
+      act(() => {
+        result.current.handleToggleSessionManagement(false);
+      });
+      act(() => {
+        result.current.handleResponseId("resp-123");
+      });
+
+      expect(result.current.responsesSessionId).toBeNull();
+    });
+
+    it("handleToggleSessionManagement should clear session when switching to UI mode", () => {
+      const { result } = renderHook(() => useChatHistory({ simplified: false }));
+
+      act(() => {
+        result.current.handleResponseId("resp-123");
+      });
+      act(() => {
+        result.current.handleToggleSessionManagement(false);
+      });
+
+      expect(result.current.useApiSessionManagement).toBe(false);
+      expect(result.current.responsesSessionId).toBeNull();
+    });
+  });
+});

ui/litellm-dashboard/src/components/playground/chat_ui/useChatHistory.ts

@@ -0,0 +1,392 @@
+import React, { useState, useEffect } from "react";
+import { MessageType, A2ATaskMetadata } from "./types";
+import { TokenUsage } from "./ResponseMetrics";
+import { MCPEvent } from "../../mcp_tools/types";
+import { truncateString } from "../../../utils/textUtils";
+
+export interface UseChatHistoryReturn {
+  // State
+  chatHistory: MessageType[];
+  setChatHistory: React.Dispatch<React.SetStateAction<MessageType[]>>;
+  mcpEvents: MCPEvent[];
+  setMCPEvents: React.Dispatch<React.SetStateAction<MCPEvent[]>>;
+  messageTraceId: string | null;
+  setMessageTraceId: React.Dispatch<React.SetStateAction<string | null>>;
+  responsesSessionId: string | null;
+  setResponsesSessionId: React.Dispatch<React.SetStateAction<string | null>>;
+  useApiSessionManagement: boolean;
+  setUseApiSessionManagement: React.Dispatch<React.SetStateAction<boolean>>;
+
+  // Actions
+  updateTextUI: (role: string, chunk: string, model?: string) => void;
+  updateReasoningContent: (chunk: string) => void;
+  updateTimingData: (timeToFirstToken: number) => void;
+  updateUsageData: (usage: TokenUsage, toolName?: string) => void;
+  updateA2AMetadata: (a2aMetadata: A2ATaskMetadata) => void;
+  updateTotalLatency: (totalLatency: number) => void;
+  updateSearchResults: (searchResults: any[]) => void;
+  handleResponseId: (responseId: string) => void;
+  handleToggleSessionManagement: (useApi: boolean) => void;
+  handleMCPEvent: (event: MCPEvent) => void;
+  updateImageUI: (imageUrl: string, model: string) => void;
+  updateEmbeddingsUI: (embeddings: string, model?: string) => void;
+  updateAudioUI: (audioUrl: string, model: string) => void;
+  updateChatImageUI: (imageUrl: string, model?: string) => void;
+  clearChatHistory: () => void;
+  clearMCPEvents: () => void;
+}
+
+export function useChatHistory({ simplified }: { simplified: boolean }): UseChatHistoryReturn {
+  const [chatHistory, setChatHistory] = useState<MessageType[]>(() => {
+    if (simplified) return [];
+    try {
+      const saved = sessionStorage.getItem("chatHistory");
+      return saved ? JSON.parse(saved) : [];
+    } catch (error) {
+      console.error("Error parsing chatHistory from sessionStorage", error);
+      return [];
+    }
+  });
+
+  const [mcpEvents, setMCPEvents] = useState<MCPEvent[]>([]);
+
+  const [messageTraceId, setMessageTraceId] = useState<string | null>(
+    () => (simplified ? null : sessionStorage.getItem("messageTraceId") || null),
+  );
+
+  const [responsesSessionId, setResponsesSessionId] = useState<string | null>(
+    () => (simplified ? null : sessionStorage.getItem("responsesSessionId") || null),
+  );
+
+  const [useApiSessionManagement, setUseApiSessionManagement] = useState<boolean>(() => {
+    if (simplified) return true;
+    const saved = sessionStorage.getItem("useApiSessionManagement");
+    return saved ? JSON.parse(saved) : true; // Default to API session management
+  });
+
+  // Debounced chatHistory persistence
+  useEffect(() => {
+    if (simplified) return; // Do not persist chat history in simplified (embedded) mode
+    // When chatHistory is empty (e.g. after clearChatHistory removed the key),
+    // don't re-write an empty array back into sessionStorage.
+    if (chatHistory.length === 0) return;
+    const handler = setTimeout(() => {
+      sessionStorage.setItem("chatHistory", JSON.stringify(chatHistory));
+    }, 500); // Debounce by 500ms
+
+    return () => {
+      clearTimeout(handler);
+    };
+  }, [chatHistory, simplified]);
+
+  // messageTraceId/responsesSessionId/useApiSessionManagement persistence
+  useEffect(() => {
+    if (simplified) return;
+    if (messageTraceId) {
+      sessionStorage.setItem("messageTraceId", messageTraceId);
+    } else {
+      sessionStorage.removeItem("messageTraceId");
+    }
+    if (responsesSessionId) {
+      sessionStorage.setItem("responsesSessionId", responsesSessionId);
+    } else {
+      sessionStorage.removeItem("responsesSessionId");
+    }
+    sessionStorage.setItem("useApiSessionManagement", JSON.stringify(useApiSessionManagement));
+  }, [messageTraceId, responsesSessionId, useApiSessionManagement, simplified]);
+
+  const updateTextUI = (role: string, chunk: string, model?: string) => {
+    setChatHistory((prev) => {
+      const last = prev[prev.length - 1];
+      // if the last message is already from this same role, append
+      if (last && last.role === role && !last.isImage && !last.isAudio) {
+        // build a new object, but only set `model` if it wasn't there already
+        const updated: MessageType = {
+          ...last,
+          content: last.content + chunk,
+          model: last.model ?? model, // ← only use the passed‐in model on the first chunk
+        };
+        return [...prev.slice(0, -1), updated];
+      } else {
+        // otherwise start a brand new assistant bubble
+        return [
+          ...prev,
+          {
+            role,
+            content: chunk,
+            model, // model set exactly once here
+          },
+        ];
+      }
+    });
+  };
+
+  const updateReasoningContent = (chunk: string) => {
+    setChatHistory((prevHistory) => {
+      const lastMessage = prevHistory[prevHistory.length - 1];
+
+      if (lastMessage && lastMessage.role === "assistant" && !lastMessage.isImage && !lastMessage.isAudio) {
+        return [
+          ...prevHistory.slice(0, prevHistory.length - 1),
+          {
+            ...lastMessage,
+            reasoningContent: (lastMessage.reasoningContent || "") + chunk,
+          },
+        ];
+      } else {
+        // If there's no assistant message yet, we'll create one with empty content
+        // but with reasoning content
+        if (prevHistory.length > 0 && prevHistory[prevHistory.length - 1].role === "user") {
+          return [
+            ...prevHistory,
+            {
+              role: "assistant",
+              content: "",
+              reasoningContent: chunk,
+            },
+          ];
+        }
+
+        return prevHistory;
+      }
+    });
+  };
+
+  const updateTimingData = (timeToFirstToken: number) => {
+    setChatHistory((prevHistory) => {
+      const lastMessage = prevHistory[prevHistory.length - 1];
+
+      if (lastMessage && lastMessage.role === "assistant") {
+        return [
+          ...prevHistory.slice(0, prevHistory.length - 1),
+          {
+            ...lastMessage,
+            timeToFirstToken,
+          },
+        ];
+      }
+      // If the last message is a user message and no assistant message exists yet,
+      // create a new assistant message with empty content
+      else if (lastMessage && lastMessage.role === "user") {
+        return [
+          ...prevHistory,
+          {
+            role: "assistant",
+            content: "",
+            timeToFirstToken,
+          },
+        ];
+      }
+
+      return prevHistory;
+    });
+  };
+
+  const updateUsageData = (usage: TokenUsage, toolName?: string) => {
+    setChatHistory((prevHistory) => {
+      const lastMessage = prevHistory[prevHistory.length - 1];
+
+      if (lastMessage && lastMessage.role === "assistant") {
+        const updatedMessage = {
+          ...lastMessage,
+          usage,
+          toolName,
+        };
+
+        return [...prevHistory.slice(0, prevHistory.length - 1), updatedMessage];
+      }
+
+      return prevHistory;
+    });
+  };
+
+  const updateA2AMetadata = (a2aMetadata: A2ATaskMetadata) => {
+    setChatHistory((prevHistory) => {
+      const lastMessage = prevHistory[prevHistory.length - 1];
+
+      if (lastMessage && lastMessage.role === "assistant") {
+        const updatedMessage = {
+          ...lastMessage,
+          a2aMetadata,
+        };
+        return [...prevHistory.slice(0, prevHistory.length - 1), updatedMessage];
+      }
+
+      return prevHistory;
+    });
+  };
+
+  const updateTotalLatency = (totalLatency: number) => {
+    setChatHistory((prevHistory) => {
+      const lastMessage = prevHistory[prevHistory.length - 1];
+
+      if (lastMessage && lastMessage.role === "assistant") {
+        return [
+          ...prevHistory.slice(0, prevHistory.length - 1),
+          {
+            ...lastMessage,
+            totalLatency,
+          },
+        ];
+      }
+
+      return prevHistory;
+    });
+  };
+
+  const updateSearchResults = (searchResults: any[]) => {
+    setChatHistory((prevHistory) => {
+      const lastMessage = prevHistory[prevHistory.length - 1];
+
+      if (lastMessage && lastMessage.role === "assistant") {
+        const updatedMessage = {
+          ...lastMessage,
+          searchResults,
+        };
+
+        return [...prevHistory.slice(0, prevHistory.length - 1), updatedMessage];
+      }
+
+      return prevHistory;
+    });
+  };
+
+  const handleResponseId = (responseId: string) => {
+    if (useApiSessionManagement) {
+      setResponsesSessionId(responseId);
+    }
+  };
+
+  const handleToggleSessionManagement = (useApi: boolean) => {
+    setUseApiSessionManagement(useApi);
+    if (!useApi) {
+      // Clear API session when switching to UI mode
+      setResponsesSessionId(null);
+    }
+  };
+
+  const handleMCPEvent = (event: MCPEvent) => {
+    setMCPEvents((prev) => {
+      // Check if this is a duplicate event (same item_id and type)
+      // Only check for duplicates if item_id is defined (for mcp_list_tools, item_id is "mcp_list_tools")
+      const isDuplicate = event.item_id
+        ? prev.some(
+            (existingEvent) =>
+              existingEvent.item_id === event.item_id &&
+              existingEvent.type === event.type &&
+              (existingEvent.sequence_number === event.sequence_number ||
+                (existingEvent.sequence_number === undefined && event.sequence_number === undefined)),
+          )
+        : false;
+
+      if (isDuplicate) {
+        return prev;
+      }
+
+      return [...prev, event];
+    });
+  };
+
+  const updateImageUI = (imageUrl: string, model: string) => {
+    setChatHistory((prevHistory) => [...prevHistory, { role: "assistant", content: imageUrl, model, isImage: true }]);
+  };
+
+  const updateEmbeddingsUI = (embeddings: string, model?: string) => {
+    setChatHistory((prevHistory) => [
+      ...prevHistory,
+      { role: "assistant", content: truncateString(embeddings, 100), model, isEmbeddings: true },
+    ]);
+  };
+
+  const updateAudioUI = (audioUrl: string, model: string) => {
+    setChatHistory((prevHistory) => [...prevHistory, { role: "assistant", content: audioUrl, model, isAudio: true }]);
+  };
+
+  const updateChatImageUI = (imageUrl: string, model?: string) => {
+    setChatHistory((prev) => {
+      const last = prev[prev.length - 1];
+      // If the last message is from assistant and has content, add image to it
+      if (last && last.role === "assistant" && !last.isImage && !last.isAudio) {
+        const updated = {
+          ...last,
+          image: {
+            url: imageUrl,
+            detail: "auto",
+          },
+          model: last.model ?? model,
+        };
+        return [...prev.slice(0, -1), updated];
+      } else {
+        // Otherwise create a new assistant message with just the image
+        return [
+          ...prev,
+          {
+            role: "assistant",
+            content: "",
+            model,
+            image: {
+              url: imageUrl,
+              detail: "auto",
+            },
+          },
+        ];
+      }
+    });
+  };
+
+  const clearChatHistory = () => {
+    // Use functional updater to get the latest snapshot — avoids stale-closure
+    // bugs where audio messages added between the last render and the click
+    // would leak their blob URLs.
+    setChatHistory((prev) => {
+      prev.forEach((message) => {
+        if (message.isAudio && typeof message.content === "string") {
+          URL.revokeObjectURL(message.content);
+        }
+      });
+      return [];
+    });
+
+    setMessageTraceId(null);
+    setResponsesSessionId(null); // Clear responses session ID
+    setMCPEvents([]); // Clear MCP events
+    if (!simplified) {
+      sessionStorage.removeItem("chatHistory");
+      sessionStorage.removeItem("messageTraceId");
+      sessionStorage.removeItem("responsesSessionId");
+    }
+  };
+
+  const clearMCPEvents = () => {
+    setMCPEvents([]);
+  };
+
+  return {
+    chatHistory,
+    setChatHistory,
+    mcpEvents,
+    setMCPEvents,
+    messageTraceId,
+    setMessageTraceId,
+    responsesSessionId,
+    setResponsesSessionId,
+    useApiSessionManagement,
+    setUseApiSessionManagement,
+    updateTextUI,
+    updateReasoningContent,
+    updateTimingData,
+    updateUsageData,
+    updateA2AMetadata,
+    updateTotalLatency,
+    updateSearchResults,
+    handleResponseId,
+    handleToggleSessionManagement,
+    handleMCPEvent,
+    updateImageUI,
+    updateEmbeddingsUI,
+    updateAudioUI,
+    updateChatImageUI,
+    clearChatHistory,
+    clearMCPEvents,
+  };
+}

ui/litellm-dashboard/src/components/policies/index.tsx

@@ -1,6 +1,7 @@
 import React, { useState, useEffect, useCallback } from "react";
 import { Button, TabGroup, TabList, Tab, TabPanels, TabPanel } from "@tremor/react";
-import { Modal, message, Alert } from "antd";
+import { Modal, Alert } from "antd";
+import MessageManager from "@/components/molecules/message_manager";
 import { ExclamationCircleOutlined, InfoCircleOutlined } from "@ant-design/icons";
 import { isAdminRole } from "@/utils/roles";
 import PolicyTable from "./policy_table";
@@ -80,7 +81,7 @@ const PoliciesPanel: React.FC<PoliciesPanelProps> = ({
       setPoliciesList(response.policies || []);
     } catch (error) {
       console.error("Error fetching policies:", error);
-      message.error("Failed to fetch policies");
+      MessageManager.error("Failed to fetch policies");
     } finally {
       setIsLoading(false);
     }
@@ -95,7 +96,7 @@ const PoliciesPanel: React.FC<PoliciesPanelProps> = ({
       setAttachmentsList(response.attachments || []);
     } catch (error) {
       console.error("Error fetching attachments:", error);
-      message.error("Failed to fetch attachments");
+      MessageManager.error("Failed to fetch attachments");
     } finally {
       setIsAttachmentsLoading(false);
     }
@@ -148,11 +149,11 @@ const PoliciesPanel: React.FC<PoliciesPanelProps> = ({
     setIsDeleting(true);
     try {
       await deletePolicyCall(accessToken, policyToDelete.policy_id);
-      message.success(`Policy "${policyToDelete.policy_name}" deleted successfully`);
+      MessageManager.success(`Policy "${policyToDelete.policy_name}" deleted successfully`);
       await fetchPolicies();
     } catch (error) {
       console.error("Error deleting policy:", error);
-      message.error("Failed to delete policy");
+      MessageManager.error("Failed to delete policy");
     } finally {
       setIsDeleting(false);
       setIsDeleteModalOpen(false);
@@ -177,11 +178,11 @@ const PoliciesPanel: React.FC<PoliciesPanelProps> = ({
         if (!accessToken) return;
         try {
           await deletePolicyAttachmentCall(accessToken, attachmentId);
-          message.success("Attachment deleted successfully");
+          MessageManager.success("Attachment deleted successfully");
           fetchAttachments();
         } catch (error) {
           console.error("Error deleting attachment:", error);
-          message.error("Failed to delete attachment");
+          MessageManager.error("Failed to delete attachment");
         }
       },
     });
@@ -193,7 +194,7 @@ const PoliciesPanel: React.FC<PoliciesPanelProps> = ({
 
   const handleUseTemplate = async (template: any) => {
     if (!accessToken) {
-      message.error("Authentication required");
+      MessageManager.error("Authentication required");
       return;
     }
 
@@ -221,7 +222,7 @@ const PoliciesPanel: React.FC<PoliciesPanelProps> = ({
       setIsGuardrailSelectionModalOpen(true);
     } catch (error) {
       console.error("Error fetching guardrails:", error);
-      message.error("Failed to load guardrails. Please try again.");
+      MessageManager.error("Failed to load guardrails. Please try again.");
     }
   };
 
@@ -271,7 +272,7 @@ const PoliciesPanel: React.FC<PoliciesPanelProps> = ({
       await proceedWithTemplate(enrichedTemplate);
     } catch (error) {
       console.error("Error enriching template:", error);
-      message.error("Failed to configure template. Please try again.");
+      MessageManager.error("Failed to configure template. Please try again.");
       setIsEnrichingTemplate(false);
     }
   };
@@ -318,15 +319,15 @@ const PoliciesPanel: React.FC<PoliciesPanelProps> = ({
 
       // Show success message
       if (createdGuardrails.length > 0) {
-        message.success(
+        MessageManager.success(
           `Created ${createdGuardrails.length} guardrail${createdGuardrails.length > 1 ? "s" : ""}! Complete the policy form to save.`
         );
       } else {
-        message.success("Template ready! Complete the policy form to save.");
+        MessageManager.success("Template ready! Complete the policy form to save.");
       }
 
       if (failedGuardrails.length > 0) {
-        message.warning(
+        MessageManager.warning(
           `Failed to create ${failedGuardrails.length} guardrail(s): ${failedGuardrails.join(", ")}. You may need to create them manually.`
         );
       }
@@ -348,7 +349,7 @@ const PoliciesPanel: React.FC<PoliciesPanelProps> = ({
       setTemplateQueue([]);
       setTemplateQueueProgress(null);
       console.error("Error creating guardrails:", error);
-      message.error("Failed to create guardrails. Please try again.");
+      MessageManager.error("Failed to create guardrails. Please try again.");
     }
   };
 

ui/litellm-dashboard/src/components/policies/pipeline_flow_builder.tsx

@@ -1,5 +1,6 @@
 import React, { useState } from "react";
-import { Select, Typography, message, Spin } from "antd";
+import { Select, Typography, Spin } from "antd";
+import MessageManager from "@/components/molecules/message_manager";
 import { Button, TextInput } from "@tremor/react";
 import { ArrowLeftIcon, PlusIcon } from "@heroicons/react/outline";
 import { DotsVerticalIcon } from "@heroicons/react/solid";
@@ -1385,17 +1386,17 @@ export const FlowBuilderPage: React.FC<FlowBuilderPageProps> = ({
 
   const handleSave = async () => {
     if (!policyName.trim()) {
-      message.error("Please enter a policy name");
+      MessageManager.error("Please enter a policy name");
       return;
     }
     if (!accessToken) {
-      message.error("No access token available");
+      MessageManager.error("No access token available");
       return;
     }
 
     const emptySteps = pipeline.steps.filter((s) => !s.guardrail);
     if (emptySteps.length > 0) {
-      message.error("Please select a guardrail for all steps");
+      MessageManager.error("Please select a guardrail for all steps");
       return;
     }
 

ui/litellm-dashboard/src/components/policies/policy_templates.tsx

@@ -1,5 +1,6 @@
 import React, { useState, useEffect, useMemo } from "react";
-import { Card, Button, Spin, message, Checkbox } from "antd";
+import { Card, Button, Spin, Checkbox } from "antd";
+import MessageManager from "@/components/molecules/message_manager";
 import {
   ShieldCheckIcon,
   ShieldExclamationIcon,
@@ -184,7 +185,7 @@ const PolicyTemplates: React.FC<PolicyTemplatesProps> = ({ onUseTemplate, onOpen
         onTemplatesLoaded?.(data);
       } catch (error) {
         console.error("Error fetching policy templates:", error);
-        message.error("Failed to fetch policy templates");
+        MessageManager.error("Failed to fetch policy templates");
       } finally {
         setIsLoading(false);
       }

ui/litellm-dashboard/src/components/shared/CreatedKeyDisplay.test.tsx

@@ -3,15 +3,11 @@ import { render, screen, act } from "@testing-library/react";
 import userEvent from "@testing-library/user-event";
 import CreatedKeyDisplay from "./CreatedKeyDisplay";
 
-vi.mock("antd", async (importOriginal) => {
-  const actual = await importOriginal<typeof import("antd")>();
-  return {
-    ...actual,
-    message: { success: vi.fn() },
-  };
-});
+vi.mock("@/components/molecules/message_manager", () => ({
+  default: { success: vi.fn(), error: vi.fn(), warning: vi.fn(), info: vi.fn(), loading: vi.fn(), destroy: vi.fn() },
+}));
 
-import { message } from "antd";
+import MessageManager from "@/components/molecules/message_manager";
 
 describe("CreatedKeyDisplay", () => {
   beforeEach(() => {
@@ -52,7 +48,7 @@ describe("CreatedKeyDisplay", () => {
 
     await user.click(screen.getByRole("button", { name: /copy virtual key/i }));
 
-    expect(message.success).toHaveBeenCalledWith("Key copied to clipboard");
+    expect(MessageManager.success).toHaveBeenCalledWith("Key copied to clipboard");
   });
 
   it("should revert button text back after 2 seconds", async () => {

ui/litellm-dashboard/src/components/shared/CreatedKeyDisplay.tsx

@@ -1,6 +1,7 @@
 import React, { useState } from "react";
 import { CopyToClipboard } from "react-copy-to-clipboard";
-import { Button, message } from "antd";
+import { Button } from "antd";
+import MessageManager from "@/components/molecules/message_manager";
 
 interface CreatedKeyDisplayProps {
   apiKey: string;
@@ -15,7 +16,7 @@ const CreatedKeyDisplay: React.FC<CreatedKeyDisplayProps> = ({ apiKey }) => {
 
   const handleCopy = () => {
     setCopied(true);
-    message.success("Key copied to clipboard");
+    MessageManager.success("Key copied to clipboard");
     setTimeout(() => setCopied(false), 2000);
   };
 

ui/litellm-dashboard/src/components/team/TeamInfo.tsx

@@ -20,7 +20,8 @@ import { isProxyAdminRole } from "@/utils/roles";
 import { EditOutlined, InfoCircleOutlined, SaveOutlined } from "@ant-design/icons";
 import { ArrowLeftIcon } from "@heroicons/react/outline";
 import { Badge, Card, Grid, Text, TextInput, Title } from "@tremor/react";
-import { Button, Form, Input, message, Select, Switch, Tabs, Tooltip } from "antd";
+import { Button, Form, Input, Select, Switch, Tabs, Tooltip } from "antd";
+import MessageManager from "@/components/molecules/message_manager";
 import { CheckIcon, CopyIcon } from "lucide-react";
 import React, { useEffect, useMemo, useState } from "react";
 import { copyToClipboard as utilCopyToClipboard } from "../../utils/dataUtils";
@@ -366,7 +367,7 @@ const TeamInfoView: React.FC<TeamInfoProps> = ({
         tpm_limit: values.tpm_limit,
         rpm_limit: values.rpm_limit,
       };
-      message.destroy(); // Remove all existing toasts
+      MessageManager.destroy(); // Remove all existing toasts
 
       await teamMemberUpdateCall(accessToken, teamId, member);
 
@@ -388,7 +389,7 @@ const TeamInfoView: React.FC<TeamInfoProps> = ({
       }
       setIsEditMemberModalVisible(false);
 
-      message.destroy(); // Remove all existing toasts
+      MessageManager.destroy(); // Remove all existing toasts
 
       NotificationsManager.fromBackend(errMsg);
       console.error("Error updating team member:", error);

ui/litellm-dashboard/src/components/ui/AntDLoadingSpinner.tsx

@@ -0,0 +1,12 @@
+import { Spin } from "antd";
+import { LoadingOutlined } from "@ant-design/icons";
+
+interface AntDLoadingSpinnerProps {
+  size?: "small" | "default" | "large";
+  fontSize?: number;
+}
+
+export function AntDLoadingSpinner({ size, fontSize }: AntDLoadingSpinnerProps) {
+  const indicator = <LoadingOutlined style={fontSize ? { fontSize } : undefined} spin />;
+  return <Spin indicator={indicator} size={size} />;
+}

ui/litellm-dashboard/src/components/vector_store_management/CreateVectorStore.tsx

@@ -1,6 +1,7 @@
 import React, { useState } from "react";
 import { Card, Title, Text } from "@tremor/react";
-import { Upload, Button, Select, Form, message, Alert, Tooltip, Input } from "antd";
+import { Upload, Button, Select, Form, Alert, Tooltip, Input } from "antd";
+import MessageManager from "@/components/molecules/message_manager";
 import { InboxOutlined, InfoCircleOutlined } from "@ant-design/icons";
 import type { UploadProps } from "antd";
 import { ragIngestCall } from "../networking";
@@ -47,13 +48,13 @@ const CreateVectorStore: React.FC<CreateVectorStoreProps> = ({ accessToken, onSu
       ].includes(file.type);
 
       if (!isValidType) {
-        message.error(`${file.name} is not a supported file type. Please upload PDF, TXT, DOCX, or MD files.`);
+        MessageManager.error(`${file.name} is not a supported file type. Please upload PDF, TXT, DOCX, or MD files.`);
         return Upload.LIST_IGNORE;
       }
 
       const isLt50M = file.size / 1024 / 1024 < 50;
       if (!isLt50M) {
-        message.error(`${file.name} must be smaller than 50MB!`);
+        MessageManager.error(`${file.name} must be smaller than 50MB!`);
         return Upload.LIST_IGNORE;
       }
 
@@ -87,12 +88,12 @@ const CreateVectorStore: React.FC<CreateVectorStoreProps> = ({ accessToken, onSu
 
   const handleCreateVectorStore = async () => {
     if (documents.length === 0) {
-      message.warning("Please upload at least one document");
+      MessageManager.warning("Please upload at least one document");
       return;
     }
 
     if (!selectedProvider) {
-      message.warning("Please select a provider");
+      MessageManager.warning("Please select a provider");
       return;
     }
 
@@ -100,7 +101,7 @@ const CreateVectorStore: React.FC<CreateVectorStoreProps> = ({ accessToken, onSu
     const requiredFields = getProviderSpecificFields(selectedProvider).filter((field) => field.required);
     for (const field of requiredFields) {
       if (!providerParams[field.name]) {
-        message.warning(`Please provide ${field.label}`);
+        MessageManager.warning(`Please provide ${field.label}`);
         return;
       }
     }
@@ -108,17 +109,17 @@ const CreateVectorStore: React.FC<CreateVectorStoreProps> = ({ accessToken, onSu
     // S3 Vectors specific validation
     if (selectedProvider === "s3_vectors") {
       if (providerParams.vector_bucket_name && providerParams.vector_bucket_name.length < 3) {
-        message.warning("Vector bucket name must be at least 3 characters");
+        MessageManager.warning("Vector bucket name must be at least 3 characters");
         return;
       }
       if (providerParams.index_name && providerParams.index_name.length > 0 && providerParams.index_name.length < 3) {
-        message.warning("Index name must be at least 3 characters if provided");
+        MessageManager.warning("Index name must be at least 3 characters if provided");
         return;
       }
     }
 
     if (!accessToken) {
-      message.error("No access token available");
+      MessageManager.error("No access token available");
       return;
     }
 

ui/litellm-dashboard/src/components/vector_store_management/DocumentsTable.tsx

@@ -1,5 +1,6 @@
 import React from "react";
-import { Table, Badge, Tooltip, message } from "antd";
+import { Table, Badge, Tooltip } from "antd";
+import MessageManager from "@/components/molecules/message_manager";
 import { EyeOutlined, CopyOutlined, DeleteOutlined } from "@ant-design/icons";
 import { DocumentUpload } from "./types";
 
@@ -11,7 +12,7 @@ interface DocumentsTableProps {
 const DocumentsTable: React.FC<DocumentsTableProps> = ({ documents, onRemove }) => {
   const handleCopyId = (uid: string) => {
     navigator.clipboard.writeText(uid);
-    message.success("Document ID copied to clipboard");
+    MessageManager.success("Document ID copied to clipboard");
   };
 
   const getStatusBadge = (status: DocumentUpload["status"]) => {

ui/litellm-dashboard/src/components/vector_store_management/VectorStoreTester.tsx

@@ -1,5 +1,6 @@
 import React, { useState } from "react";
-import { Button, Input, Card, Typography, Spin, message, Divider } from "antd";
+import { Button, Input, Card, Typography, Spin, Divider } from "antd";
+import MessageManager from "@/components/molecules/message_manager";
 import { SendOutlined, DatabaseOutlined, LoadingOutlined, DownOutlined, RightOutlined } from "@ant-design/icons";
 import { vectorStoreSearchCall } from "../networking";
 import NotificationsManager from "../molecules/notifications_manager";
@@ -46,7 +47,7 @@ export const VectorStoreTester: React.FC<VectorStoreTesterProps> = ({ vectorStor
 
   const handleSearch = async () => {
     if (!query.trim()) {
-      message.warning("Please enter a search query");
+      MessageManager.warning("Please enter a search query");
       return;
     }
 

ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/CollapsibleMessage.test.tsx

@@ -0,0 +1,54 @@
+import React from "react";
+import { render, screen } from "@testing-library/react";
+import userEvent from "@testing-library/user-event";
+import { describe, it, expect } from "vitest";
+import { CollapsibleMessage } from "./CollapsibleMessage";
+
+describe("CollapsibleMessage", () => {
+  it("should return null when content is empty", () => {
+    const { container } = render(
+      <CollapsibleMessage label="SYSTEM" content="" />
+    );
+    expect(container.innerHTML).toBe("");
+  });
+
+  it("should return null when content is undefined", () => {
+    const { container } = render(<CollapsibleMessage label="SYSTEM" />);
+    expect(container.innerHTML).toBe("");
+  });
+
+  it("should render the label and char count", () => {
+    render(<CollapsibleMessage label="SYSTEM" content="Hello" />);
+    expect(screen.getByText("SYSTEM")).toBeInTheDocument();
+    expect(screen.getByText("(5 chars)")).toBeInTheDocument();
+  });
+
+  it("should show content when defaultExpanded is true", () => {
+    render(
+      <CollapsibleMessage
+        label="SYSTEM"
+        content="Visible text"
+        defaultExpanded={true}
+      />
+    );
+    expect(screen.getByText("Visible text")).toBeInTheDocument();
+  });
+
+  it("should toggle expanded state when header is clicked", async () => {
+    const user = userEvent.setup();
+    render(
+      <CollapsibleMessage
+        label="SYSTEM"
+        content="Toggle me"
+        defaultExpanded={false}
+      />
+    );
+
+    // Content is rendered in DOM but collapsed by default
+    expect(screen.getByText("Toggle me")).toBeInTheDocument();
+
+    // Click the header to expand - should still show content
+    await user.click(screen.getByText("SYSTEM"));
+    expect(screen.getByText("Toggle me")).toBeInTheDocument();
+  });
+});

ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/HistoryTree.test.tsx

@@ -0,0 +1,50 @@
+import React from "react";
+import { render, screen } from "@testing-library/react";
+import userEvent from "@testing-library/user-event";
+import { describe, it, expect } from "vitest";
+import { HistoryTree } from "./HistoryTree";
+import { ParsedMessage } from "./prettyMessagesTypes";
+
+describe("HistoryTree", () => {
+  it("should return null when messages array is empty", () => {
+    const { container } = render(<HistoryTree messages={[]} />);
+    expect(container.innerHTML).toBe("");
+  });
+
+  it('should render message count with plural "messages" for multiple messages', () => {
+    const messages: ParsedMessage[] = [
+      { role: "user", content: "Hello" },
+      { role: "assistant", content: "Hi there" },
+      { role: "user", content: "How are you?" },
+    ];
+    render(<HistoryTree messages={messages} />);
+    expect(
+      screen.getByText("HISTORY (3 messages)")
+    ).toBeInTheDocument();
+  });
+
+  it('should render message count with singular "message" for one message', () => {
+    const messages: ParsedMessage[] = [
+      { role: "user", content: "Hello" },
+    ];
+    render(<HistoryTree messages={messages} />);
+    expect(
+      screen.getByText("HISTORY (1 message)")
+    ).toBeInTheDocument();
+  });
+
+  it("should expand and show messages when header is clicked", async () => {
+    const user = userEvent.setup();
+    const messages: ParsedMessage[] = [
+      { role: "user", content: "Hello" },
+      { role: "assistant", content: "Hi there" },
+    ];
+    render(<HistoryTree messages={messages} />);
+
+    // Click to expand
+    await user.click(screen.getByText("HISTORY (2 messages)"));
+
+    expect(screen.getByText("Hello")).toBeInTheDocument();
+    expect(screen.getByText("Hi there")).toBeInTheDocument();
+  });
+});

ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/InputCard.tsx

@@ -4,7 +4,7 @@
  */
 
 import { useState } from 'react';
-import { message } from 'antd';
+import MessageManager from "@/components/molecules/message_manager";
 import { ParsedMessage } from './prettyMessagesTypes';
 import { SectionHeader } from './SectionHeader';
 import { CollapsibleMessage } from './CollapsibleMessage';
@@ -33,7 +33,7 @@ export function InputCard({ messages, promptTokens, inputCost }: InputCardProps)
   const handleCopy = () => {
     const content = lastMessage?.content || '';
     navigator.clipboard.writeText(content);
-    message.success('Input copied');
+    MessageManager.success('Input copied');
   };
 
   return (

ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/OutputCard.tsx

@@ -4,7 +4,8 @@
  */
 
 import { useState } from 'react';
-import { Typography, message as antdMessage } from 'antd';
+import { Typography } from 'antd';
+import MessageManager from "@/components/molecules/message_manager";
 import { ParsedMessage } from './prettyMessagesTypes';
 import { SectionHeader } from './SectionHeader';
 import { SimpleMessageBlock } from './SimpleMessageBlock';
@@ -25,7 +26,7 @@ export function OutputCard({ message, completionTokens, outputCost }: OutputCard
     
     const content = message.content || '';
     navigator.clipboard.writeText(content);
-    antdMessage.success('Output copied');
+    MessageManager.success('Output copied');
   };
 
   if (!message) {

ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/SimpleMessageBlock.test.tsx

@@ -0,0 +1,55 @@
+import React from "react";
+import { render, screen } from "@testing-library/react";
+import { describe, it, expect } from "vitest";
+import { SimpleMessageBlock } from "./SimpleMessageBlock";
+
+describe("SimpleMessageBlock", () => {
+  it("should render the label and content", () => {
+    render(<SimpleMessageBlock label="USER" content="Hello world" />);
+    expect(screen.getByText("USER")).toBeInTheDocument();
+    expect(screen.getByText("Hello world")).toBeInTheDocument();
+  });
+
+  it("should return null when content is empty and no tool calls", () => {
+    const { container } = render(
+      <SimpleMessageBlock label="USER" content="" />
+    );
+    expect(container.innerHTML).toBe("");
+  });
+
+  it('should return null when content is "null" string and no tool calls', () => {
+    const { container } = render(
+      <SimpleMessageBlock label="USER" content="null" />
+    );
+    expect(container.innerHTML).toBe("");
+  });
+
+  it("should render tool calls when present", () => {
+    render(
+      <SimpleMessageBlock
+        label="ASSISTANT"
+        toolCalls={[
+          { id: "tc1", name: "get_weather", arguments: { city: "Paris" } },
+        ]}
+      />
+    );
+    expect(screen.getByText("ASSISTANT")).toBeInTheDocument();
+    expect(screen.getByText("get_weather")).toBeInTheDocument();
+  });
+
+  it("should render content and tool calls together", () => {
+    render(
+      <SimpleMessageBlock
+        label="ASSISTANT"
+        content="Let me check the weather."
+        toolCalls={[
+          { id: "tc1", name: "get_weather", arguments: { city: "Paris" } },
+        ]}
+      />
+    );
+    expect(
+      screen.getByText("Let me check the weather.")
+    ).toBeInTheDocument();
+    expect(screen.getByText("get_weather")).toBeInTheDocument();
+  });
+});

ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/SimpleToolCallBlock.test.tsx

@@ -0,0 +1,51 @@
+import React from "react";
+import { render, screen } from "@testing-library/react";
+import { describe, it, expect } from "vitest";
+import { SimpleToolCallBlock } from "./SimpleToolCallBlock";
+
+describe("SimpleToolCallBlock", () => {
+  it("should render the tool name", () => {
+    render(
+      <SimpleToolCallBlock
+        tool={{ id: "1", name: "get_weather", arguments: {} }}
+      />
+    );
+    expect(screen.getByText("get_weather")).toBeInTheDocument();
+  });
+
+  it('should display "function" badge', () => {
+    render(
+      <SimpleToolCallBlock
+        tool={{ id: "1", name: "get_weather", arguments: {} }}
+      />
+    );
+    expect(screen.getByText("function")).toBeInTheDocument();
+  });
+
+  it("should render arguments when present", () => {
+    render(
+      <SimpleToolCallBlock
+        tool={{
+          id: "1",
+          name: "get_weather",
+          arguments: { city: "London", units: "metric" },
+        }}
+      />
+    );
+    expect(screen.getByText("city:")).toBeInTheDocument();
+    expect(screen.getByText('"London"')).toBeInTheDocument();
+    expect(screen.getByText("units:")).toBeInTheDocument();
+    expect(screen.getByText('"metric"')).toBeInTheDocument();
+  });
+
+  it("should not render arguments section when arguments are empty", () => {
+    const { container } = render(
+      <SimpleToolCallBlock
+        tool={{ id: "1", name: "get_weather", arguments: {} }}
+      />
+    );
+    // The tool name and "function" badge should be there, but no key: value pairs
+    expect(screen.getByText("get_weather")).toBeInTheDocument();
+    expect(screen.queryByText(/:$/)).not.toBeInTheDocument();
+  });
+});

ui/litellm-dashboard/src/contexts/AntdGlobalProvider.tsx

@@ -1,23 +1,27 @@
 "use client";
 
 import React, { useEffect, useRef } from "react";
-import { notification } from "antd";
+import { notification, message } from "antd";
 import { setNotificationInstance } from "@/components/molecules/notifications_manager";
+import { setMessageInstance } from "@/components/molecules/message_manager";
 
 export default function AntdGlobalProvider({ children }: { children: React.ReactNode }) {
-  const [api, contextHolder] = notification.useNotification();
+  const [notificationApi, notificationContextHolder] = notification.useNotification();
+  const [messageApi, messageContextHolder] = message.useMessage();
   const initialized = useRef(false);
 
   useEffect(() => {
     if (!initialized.current) {
-      setNotificationInstance(api);
+      setNotificationInstance(notificationApi);
+      setMessageInstance(messageApi);
       initialized.current = true;
     }
-  }, [api]);
+  }, [notificationApi, messageApi]);
 
   return (
     <>
-      {contextHolder}
+      {notificationContextHolder}
+      {messageContextHolder}
       {children}
     </>
   );

ui/litellm-dashboard/src/hooks/useWorker.ts

@@ -0,0 +1,65 @@
+"use client";
+
+import { useState, useCallback, useEffect } from "react";
+import { switchToWorkerUrl, WorkerInfo } from "@/components/networking";
+import { useUIConfig } from "@/app/(dashboard)/hooks/uiConfig/useUIConfig";
+
+const SELECTED_WORKER_KEY = "litellm_selected_worker_id";
+
+interface UseWorkerReturn {
+  isControlPlane: boolean;
+  workers: WorkerInfo[];
+  selectedWorkerId: string | null;
+  selectedWorker: WorkerInfo | null;
+  selectWorker: (workerId: string) => void;
+  disconnectFromWorker: () => void;
+}
+
+export const useWorker = (): UseWorkerReturn => {
+  const { data: uiConfig } = useUIConfig();
+  const isControlPlane = uiConfig?.is_control_plane ?? false;
+  const workers: WorkerInfo[] = uiConfig?.workers ?? [];
+
+  const [selectedWorkerId, setSelectedWorkerId] = useState<string | null>(() => {
+    if (typeof window === "undefined") return null;
+    return localStorage.getItem(SELECTED_WORKER_KEY);
+  });
+
+  // Once workers are loaded, restore proxyBaseUrl from the persisted selection
+  useEffect(() => {
+    if (!selectedWorkerId || workers.length === 0) return;
+    const worker = workers.find((w) => w.worker_id === selectedWorkerId);
+    if (worker) {
+      switchToWorkerUrl(worker.url);
+    }
+  }, [selectedWorkerId, workers]);
+
+  const selectedWorker =
+    workers.find((w) => w.worker_id === selectedWorkerId) ?? null;
+
+  const selectWorker = useCallback(
+    (workerId: string) => {
+      const worker = workers.find((w) => w.worker_id === workerId);
+      if (!worker) return;
+      setSelectedWorkerId(workerId);
+      localStorage.setItem(SELECTED_WORKER_KEY, workerId);
+      switchToWorkerUrl(worker.url);
+    },
+    [workers],
+  );
+
+  const disconnectFromWorker = useCallback(() => {
+    setSelectedWorkerId(null);
+    localStorage.removeItem(SELECTED_WORKER_KEY);
+    switchToWorkerUrl(null);
+  }, []);
+
+  return {
+    isControlPlane,
+    workers,
+    selectedWorkerId,
+    selectedWorker,
+    selectWorker,
+    disconnectFromWorker,
+  };
+};

ui/litellm-dashboard/tests/setupTests.ts

@@ -100,6 +100,11 @@ if (!document.getAnimations) {
   document.getAnimations = () => [];
 }
 
+// Stub URL.revokeObjectURL so vi.spyOn can intercept it in tests
+if (!URL.revokeObjectURL) {
+  URL.revokeObjectURL = () => {};
+}
+
 // Mock ResizeObserver for components that use it (e.g., Tremor UI components)
 // This prevents "ResizeObserver is not defined" errors in JSDOM
 global.ResizeObserver = class ResizeObserver {