Merge pull request #20386 from naaa760/fix/extra-head-chat-comp-brid

fix(proxy): forward extra headers in chat

docs/my-website/docs/a2a.md

@@ -68,116 +68,9 @@ Follow [this guide, to add your pydantic ai agent to LiteLLM Agent Gateway](./pr
 
 ## Invoking your Agents
 
-Use the [A2A Python SDK](https://pypi.org/project/a2a-sdk) to invoke agents through LiteLLM.
-
-This example shows how to:
-1. **List available agents** - Query `/v1/agents` to see which agents your key can access
-2. **Select an agent** - Pick an agent from the list
-3. **Invoke via A2A** - Use the A2A protocol to send messages to the agent
-
-```python showLineNumbers title="invoke_a2a_agent.py"
-from uuid import uuid4
-import httpx
-import asyncio
-from a2a.client import A2ACardResolver, A2AClient
-from a2a.types import MessageSendParams, SendMessageRequest
-
-# === CONFIGURE THESE ===
-LITELLM_BASE_URL = "http://localhost:4000"  # Your LiteLLM proxy URL
-LITELLM_VIRTUAL_KEY = "sk-1234"             # Your LiteLLM Virtual Key
-# =======================
-
-async def main():
-    headers = {"Authorization": f"Bearer {LITELLM_VIRTUAL_KEY}"}
-    
-    async with httpx.AsyncClient(headers=headers) as client:
-        # Step 1: List available agents
-        response = await client.get(f"{LITELLM_BASE_URL}/v1/agents")
-        agents = response.json()
-        
-        print("Available agents:")
-        for agent in agents:
-            print(f"  - {agent['agent_name']} (ID: {agent['agent_id']})")
-        
-        if not agents:
-            print("No agents available for this key")
-            return
-        
-        # Step 2: Select an agent and invoke it
-        selected_agent = agents[0]
-        agent_id = selected_agent["agent_id"]
-        agent_name = selected_agent["agent_name"]
-        print(f"\nInvoking: {agent_name}")
-        
-        # Step 3: Use A2A protocol to invoke the agent
-        base_url = f"{LITELLM_BASE_URL}/a2a/{agent_id}"
-        resolver = A2ACardResolver(httpx_client=client, base_url=base_url)
-        agent_card = await resolver.get_agent_card()
-        a2a_client = A2AClient(httpx_client=client, agent_card=agent_card)
-        
-        request = SendMessageRequest(
-            id=str(uuid4()),
-            params=MessageSendParams(
-                message={
-                    "role": "user",
-                    "parts": [{"kind": "text", "text": "Hello, what can you do?"}],
-                    "messageId": uuid4().hex,
-                }
-            ),
-        )
-        response = await a2a_client.send_message(request)
-        print(f"Response: {response.model_dump(mode='json', exclude_none=True, indent=4)}")
-
-if __name__ == "__main__":
-    asyncio.run(main())
-```
-
-### Streaming Responses
-
-For streaming responses, use `send_message_streaming`:
-
-```python showLineNumbers title="invoke_a2a_agent_streaming.py"
-from uuid import uuid4
-import httpx
-import asyncio
-from a2a.client import A2ACardResolver, A2AClient
-from a2a.types import MessageSendParams, SendStreamingMessageRequest
-
-# === CONFIGURE THESE ===
-LITELLM_BASE_URL = "http://localhost:4000"  # Your LiteLLM proxy URL
-LITELLM_VIRTUAL_KEY = "sk-1234"             # Your LiteLLM Virtual Key
-LITELLM_AGENT_NAME = "ij-local"             # Agent name registered in LiteLLM
-# =======================
-
-async def main():
-    base_url = f"{LITELLM_BASE_URL}/a2a/{LITELLM_AGENT_NAME}"
-    headers = {"Authorization": f"Bearer {LITELLM_VIRTUAL_KEY}"}
-    
-    async with httpx.AsyncClient(headers=headers) as httpx_client:
-        # Resolve agent card and create client
-        resolver = A2ACardResolver(httpx_client=httpx_client, base_url=base_url)
-        agent_card = await resolver.get_agent_card()
-        client = A2AClient(httpx_client=httpx_client, agent_card=agent_card)
-
-        # Send a streaming message
-        request = SendStreamingMessageRequest(
-            id=str(uuid4()),
-            params=MessageSendParams(
-                message={
-                    "role": "user",
-                    "parts": [{"kind": "text", "text": "Hello, what can you do?"}],
-                    "messageId": uuid4().hex,
-                }
-            ),
-        )
-        
-        # Stream the response
-        async for chunk in client.send_message_streaming(request):
-            print(chunk.model_dump(mode="json", exclude_none=True))
-
-if __name__ == "__main__":
-    asyncio.run(main())
-```
+See the [Invoking A2A Agents](./a2a_invoking_agents) guide to learn how to call your agents using:
+- **A2A SDK** - Native A2A protocol with full support for tasks and artifacts
+- **OpenAI SDK** - Familiar `/chat/completions` interface with `a2a/` model prefix
 
 ## Tracking Agent Logs
 

docs/my-website/docs/a2a_invoking_agents.md

@@ -0,0 +1,280 @@
+import Tabs from '@theme/Tabs';
+import TabItem from '@theme/TabItem';
+
+# Invoking A2A Agents
+
+Learn how to invoke A2A agents through LiteLLM using different methods.
+
+:::tip Deploy Your Own A2A Agent
+
+Want to test with your own agent? Deploy this template A2A agent powered by Google Gemini:
+
+[**shin-bot-litellm/a2a-gemini-agent**](https://github.com/shin-bot-litellm/a2a-gemini-agent) - Simple deployable A2A agent with streaming support
+
+:::
+
+## A2A SDK
+
+Use the [A2A Python SDK](https://pypi.org/project/a2a-sdk) to invoke agents through LiteLLM using the A2A protocol.
+
+### Non-Streaming
+
+This example shows how to:
+1. **List available agents** - Query `/v1/agents` to see which agents your key can access
+2. **Select an agent** - Pick an agent from the list
+3. **Invoke via A2A** - Use the A2A protocol to send messages to the agent
+
+```python showLineNumbers title="invoke_a2a_agent.py"
+from uuid import uuid4
+import httpx
+import asyncio
+from a2a.client import A2ACardResolver, A2AClient
+from a2a.types import MessageSendParams, SendMessageRequest
+
+# === CONFIGURE THESE ===
+LITELLM_BASE_URL = "http://localhost:4000"  # Your LiteLLM proxy URL
+LITELLM_VIRTUAL_KEY = "sk-1234"             # Your LiteLLM Virtual Key
+# =======================
+
+async def main():
+    headers = {"Authorization": f"Bearer {LITELLM_VIRTUAL_KEY}"}
+    
+    async with httpx.AsyncClient(headers=headers) as client:
+        # Step 1: List available agents
+        response = await client.get(f"{LITELLM_BASE_URL}/v1/agents")
+        agents = response.json()
+        
+        print("Available agents:")
+        for agent in agents:
+            print(f"  - {agent['agent_name']} (ID: {agent['agent_id']})")
+        
+        if not agents:
+            print("No agents available for this key")
+            return
+        
+        # Step 2: Select an agent and invoke it
+        selected_agent = agents[0]
+        agent_id = selected_agent["agent_id"]
+        agent_name = selected_agent["agent_name"]
+        print(f"\nInvoking: {agent_name}")
+        
+        # Step 3: Use A2A protocol to invoke the agent
+        base_url = f"{LITELLM_BASE_URL}/a2a/{agent_id}"
+        resolver = A2ACardResolver(httpx_client=client, base_url=base_url)
+        agent_card = await resolver.get_agent_card()
+        a2a_client = A2AClient(httpx_client=client, agent_card=agent_card)
+        
+        request = SendMessageRequest(
+            id=str(uuid4()),
+            params=MessageSendParams(
+                message={
+                    "role": "user",
+                    "parts": [{"kind": "text", "text": "Hello, what can you do?"}],
+                    "messageId": uuid4().hex,
+                }
+            ),
+        )
+        response = await a2a_client.send_message(request)
+        print(f"Response: {response.model_dump(mode='json', exclude_none=True, indent=4)}")
+
+if __name__ == "__main__":
+    asyncio.run(main())
+```
+
+### Streaming
+
+For streaming responses, use `send_message_streaming`:
+
+```python showLineNumbers title="invoke_a2a_agent_streaming.py"
+from uuid import uuid4
+import httpx
+import asyncio
+from a2a.client import A2ACardResolver, A2AClient
+from a2a.types import MessageSendParams, SendStreamingMessageRequest
+
+# === CONFIGURE THESE ===
+LITELLM_BASE_URL = "http://localhost:4000"  # Your LiteLLM proxy URL
+LITELLM_VIRTUAL_KEY = "sk-1234"             # Your LiteLLM Virtual Key
+LITELLM_AGENT_NAME = "ij-local"             # Agent name registered in LiteLLM
+# =======================
+
+async def main():
+    base_url = f"{LITELLM_BASE_URL}/a2a/{LITELLM_AGENT_NAME}"
+    headers = {"Authorization": f"Bearer {LITELLM_VIRTUAL_KEY}"}
+    
+    async with httpx.AsyncClient(headers=headers) as httpx_client:
+        # Resolve agent card and create client
+        resolver = A2ACardResolver(httpx_client=httpx_client, base_url=base_url)
+        agent_card = await resolver.get_agent_card()
+        client = A2AClient(httpx_client=httpx_client, agent_card=agent_card)
+
+        # Send a streaming message
+        request = SendStreamingMessageRequest(
+            id=str(uuid4()),
+            params=MessageSendParams(
+                message={
+                    "role": "user",
+                    "parts": [{"kind": "text", "text": "Tell me a long story"}],
+                    "messageId": uuid4().hex,
+                }
+            ),
+        )
+        
+        # Stream the response
+        async for chunk in client.send_message_streaming(request):
+            print(chunk.model_dump(mode="json", exclude_none=True))
+
+if __name__ == "__main__":
+    asyncio.run(main())
+```
+
+## /chat/completions API (OpenAI SDK)
+
+You can also invoke A2A agents using the familiar OpenAI SDK by using the `a2a/` model prefix.
+
+### Non-Streaming
+
+<Tabs>
+<TabItem value="python" label="Python" default>
+
+```python showLineNumbers title="openai_non_streaming.py"
+import openai
+
+client = openai.OpenAI(
+    api_key="sk-1234",  # Your LiteLLM Virtual Key
+    base_url="http://localhost:4000"  # Your LiteLLM proxy URL
+)
+
+response = client.chat.completions.create(
+    model="a2a/my-agent",  # Use a2a/ prefix with your agent name
+    messages=[
+        {"role": "user", "content": "Hello, what can you do?"}
+    ]
+)
+
+print(response.choices[0].message.content)
+```
+
+</TabItem>
+<TabItem value="typescript" label="TypeScript">
+
+```typescript showLineNumbers title="openai_non_streaming.ts"
+import OpenAI from 'openai';
+
+const client = new OpenAI({
+  apiKey: 'sk-1234',  // Your LiteLLM Virtual Key
+  baseURL: 'http://localhost:4000'  // Your LiteLLM proxy URL
+});
+
+const response = await client.chat.completions.create({
+  model: 'a2a/my-agent',  // Use a2a/ prefix with your agent name
+  messages: [
+    { role: 'user', content: 'Hello, what can you do?' }
+  ]
+});
+
+console.log(response.choices[0].message.content);
+```
+
+</TabItem>
+<TabItem value="curl" label="cURL">
+
+```bash showLineNumbers title="curl_non_streaming.sh"
+curl -X POST http://localhost:4000/v1/chat/completions \
+  -H "Authorization: Bearer sk-1234" \
+  -H "Content-Type: application/json" \
+  -d '{
+    "model": "a2a/my-agent",
+    "messages": [
+      {"role": "user", "content": "Hello, what can you do?"}
+    ]
+  }'
+```
+
+</TabItem>
+</Tabs>
+
+### Streaming
+
+<Tabs>
+<TabItem value="python" label="Python" default>
+
+```python showLineNumbers title="openai_streaming.py"
+import openai
+
+client = openai.OpenAI(
+    api_key="sk-1234",  # Your LiteLLM Virtual Key
+    base_url="http://localhost:4000"  # Your LiteLLM proxy URL
+)
+
+stream = client.chat.completions.create(
+    model="a2a/my-agent",  # Use a2a/ prefix with your agent name
+    messages=[
+        {"role": "user", "content": "Tell me a long story"}
+    ],
+    stream=True
+)
+
+for chunk in stream:
+    if chunk.choices[0].delta.content:
+        print(chunk.choices[0].delta.content, end="", flush=True)
+```
+
+</TabItem>
+<TabItem value="typescript" label="TypeScript">
+
+```typescript showLineNumbers title="openai_streaming.ts"
+import OpenAI from 'openai';
+
+const client = new OpenAI({
+  apiKey: 'sk-1234',  // Your LiteLLM Virtual Key
+  baseURL: 'http://localhost:4000'  // Your LiteLLM proxy URL
+});
+
+const stream = await client.chat.completions.create({
+  model: 'a2a/my-agent',  // Use a2a/ prefix with your agent name
+  messages: [
+    { role: 'user', content: 'Tell me a long story' }
+  ],
+  stream: true
+});
+
+for await (const chunk of stream) {
+  const content = chunk.choices[0]?.delta?.content;
+  if (content) {
+    process.stdout.write(content);
+  }
+}
+```
+
+</TabItem>
+<TabItem value="curl" label="cURL">
+
+```bash showLineNumbers title="curl_streaming.sh"
+curl -X POST http://localhost:4000/v1/chat/completions \
+  -H "Authorization: Bearer sk-1234" \
+  -H "Content-Type: application/json" \
+  -d '{
+    "model": "a2a/my-agent",
+    "messages": [
+      {"role": "user", "content": "Tell me a long story"}
+    ],
+    "stream": true
+  }'
+```
+
+</TabItem>
+</Tabs>
+
+## Key Differences
+
+| Method | Use Case | Advantages |
+|--------|----------|------------|
+| **A2A SDK** | Native A2A protocol integration | • Full A2A protocol support<br/>• Access to task states and artifacts<br/>• Context management |
+| **OpenAI SDK** | Familiar OpenAI-style interface | • Drop-in replacement for OpenAI calls<br/>• Easier migration from LLM to agent workflows<br/>• Works with existing OpenAI tooling |
+
+:::tip Model Prefix
+
+When using the OpenAI SDK, always prefix your agent name with `a2a/` (e.g., `a2a/my-agent`) to route requests to the A2A agent instead of an LLM provider.
+
+:::

docs/my-website/sidebars.js

@@ -469,6 +469,7 @@ const sidebars = {
           label: "/a2a - A2A Agent Gateway",
           items: [
             "a2a",
+            "a2a_invoking_agents",
             "a2a_cost_tracking",
             "a2a_agent_permissions"
           ],

enterprise/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "litellm-enterprise"
-version = "0.1.28"
+version = "0.1.29"
 description = "Package for LiteLLM Enterprise features"
 authors = ["BerriAI"]
 readme = "README.md"
@@ -22,7 +22,7 @@ requires = ["poetry-core"]
 build-backend = "poetry.core.masonry.api"
 
 [tool.commitizen]
-version = "0.1.28"
+version = "0.1.29"
 version_files = [
     "pyproject.toml:version",
     "../requirements.txt:litellm-enterprise==",

litellm/__init__.py

@@ -1380,6 +1380,7 @@ if TYPE_CHECKING:
     from .llms.topaz.image_variations.transformation import TopazImageVariationConfig as TopazImageVariationConfig
     from litellm.llms.openai.completion.transformation import OpenAITextCompletionConfig as OpenAITextCompletionConfig
     from .llms.groq.chat.transformation import GroqChatConfig as GroqChatConfig
+    from .llms.a2a.chat.transformation import A2AConfig as A2AConfig
     from .llms.voyage.embedding.transformation import VoyageEmbeddingConfig as VoyageEmbeddingConfig
     from .llms.voyage.embedding.transformation_contextual import VoyageContextualEmbeddingConfig as VoyageContextualEmbeddingConfig
     from .llms.infinity.embedding.transformation import InfinityEmbeddingConfig as InfinityEmbeddingConfig

litellm/_lazy_imports_registry.py

@@ -213,6 +213,7 @@ LLM_CONFIG_NAMES = (
     "TopazImageVariationConfig",
     "OpenAITextCompletionConfig",
     "GroqChatConfig",
+    "A2AConfig",
     "GenAIHubOrchestrationConfig",
     "VoyageEmbeddingConfig",
     "VoyageContextualEmbeddingConfig",
@@ -850,6 +851,7 @@ _LLM_CONFIGS_IMPORT_MAP = {
         "OpenAITextCompletionConfig",
     ),
     "GroqChatConfig": (".llms.groq.chat.transformation", "GroqChatConfig"),
+    "A2AConfig": (".llms.a2a.chat.transformation", "A2AConfig"),
     "GenAIHubOrchestrationConfig": (
         ".llms.sap.chat.transformation",
         "GenAIHubOrchestrationConfig",

litellm/completion_extras/litellm_responses_transformation/transformation.py

@@ -329,6 +329,9 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge):
                 else:
                     request_data[key] = value
 
+        if headers:
+            request_data["extra_headers"] = headers
+
         return request_data
 
     @staticmethod

litellm/llms/a2a/__init__.py

@@ -0,0 +1,6 @@
+"""
+A2A (Agent-to-Agent) Protocol Provider for LiteLLM
+"""
+from .chat.transformation import A2AConfig
+
+__all__ = ["A2AConfig"]

litellm/llms/a2a/chat/__init__.py

@@ -0,0 +1,6 @@
+"""
+A2A Chat Completion Implementation
+"""
+from .transformation import A2AConfig
+
+__all__ = ["A2AConfig"]

litellm/llms/a2a/chat/streaming_iterator.py

@@ -0,0 +1,103 @@
+"""
+A2A Streaming Response Iterator
+"""
+from typing import Optional, Union
+
+from litellm.llms.base_llm.base_model_iterator import BaseModelResponseIterator
+from litellm.types.utils import GenericStreamingChunk, ModelResponseStream
+
+from ..common_utils import extract_text_from_a2a_response
+
+
+class A2AModelResponseIterator(BaseModelResponseIterator):
+    """
+    Iterator for parsing A2A streaming responses.
+    
+    Converts A2A JSON-RPC streaming chunks to OpenAI-compatible format.
+    """
+    
+    def __init__(
+        self,
+        streaming_response,
+        sync_stream: bool,
+        json_mode: Optional[bool] = False,
+        model: str = "a2a/agent",
+    ):
+        super().__init__(
+            streaming_response=streaming_response,
+            sync_stream=sync_stream,
+            json_mode=json_mode,
+        )
+        self.model = model
+    
+    def chunk_parser(self, chunk: dict) -> Union[GenericStreamingChunk, ModelResponseStream]:
+        """
+        Parse A2A streaming chunk to OpenAI format.
+        
+        A2A chunk format:
+        {
+            "jsonrpc": "2.0",
+            "id": "request-id",
+            "result": {
+                "message": {
+                    "parts": [{"kind": "text", "text": "content"}]
+                }
+            }
+        }
+        
+        Or for tasks:
+        {
+            "jsonrpc": "2.0",
+            "result": {
+                "kind": "task",
+                "status": {"state": "running"},
+                "artifacts": [{"parts": [{"kind": "text", "text": "content"}]}]
+            }
+        }
+        """
+        try:
+            # Extract text from A2A response
+            text = extract_text_from_a2a_response(chunk)
+            
+            # Determine finish reason
+            finish_reason = self._get_finish_reason(chunk)
+            
+            # Return generic streaming chunk
+            return GenericStreamingChunk(
+                text=text,
+                is_finished=bool(finish_reason),
+                finish_reason=finish_reason or "",
+                usage=None,
+                index=0,
+                tool_use=None,
+            )
+        except Exception:
+            # Return empty chunk on parse error
+            return GenericStreamingChunk(
+                text="",
+                is_finished=False,
+                finish_reason="",
+                usage=None,
+                index=0,
+                tool_use=None,
+            )
+    
+    def _get_finish_reason(self, chunk: dict) -> Optional[str]:
+        """Extract finish reason from A2A chunk"""
+        result = chunk.get("result", {})
+        
+        # Check for task completion
+        if isinstance(result, dict):
+            status = result.get("status", {})
+            if isinstance(status, dict):
+                state = status.get("state")
+                if state == "completed":
+                    return "stop"
+                elif state == "failed":
+                    return "stop"  # Map failed state to 'stop' (valid finish_reason)
+        
+        # Check for [DONE] marker
+        if chunk.get("done") is True:
+            return "stop"
+        
+        return None

litellm/llms/a2a/chat/transformation.py

@@ -0,0 +1,370 @@
+"""
+A2A Protocol Transformation for LiteLLM
+"""
+import uuid
+from typing import Any, Dict, Iterator, List, Optional, Union
+
+import httpx
+
+from litellm.llms.base_llm.base_model_iterator import BaseModelResponseIterator
+from litellm.llms.base_llm.chat.transformation import BaseConfig, BaseLLMException
+from litellm.types.llms.openai import AllMessageValues
+from litellm.types.utils import Choices, Message, ModelResponse
+
+from ..common_utils import (
+    A2AError,
+    convert_messages_to_prompt,
+    extract_text_from_a2a_response,
+)
+from .streaming_iterator import A2AModelResponseIterator
+
+
+class A2AConfig(BaseConfig):
+    """
+    Configuration for A2A (Agent-to-Agent) Protocol.
+    
+    Handles transformation between OpenAI and A2A JSON-RPC 2.0 formats.
+    """
+    
+    @staticmethod
+    def resolve_agent_config_from_registry(
+        model: str,
+        api_base: Optional[str],
+        api_key: Optional[str],
+        headers: Optional[Dict[str, Any]],
+        optional_params: Dict[str, Any],
+    ) -> tuple[Optional[str], Optional[str], Optional[Dict[str, Any]]]:
+        """
+        Resolve agent configuration from registry if model format is "a2a/<agent-name>".
+        
+        Extracts agent name from model string and looks up configuration in the
+        agent registry (if available in proxy context).
+        
+        Args:
+            model: Model string (e.g., "a2a/my-agent")
+            api_base: Explicit api_base (takes precedence over registry)
+            api_key: Explicit api_key (takes precedence over registry)
+            headers: Explicit headers (takes precedence over registry)
+            optional_params: Dict to merge additional litellm_params into
+        
+        Returns:
+            Tuple of (api_base, api_key, headers) with registry values filled in
+        """
+        # Extract agent name from model (e.g., "a2a/my-agent" -> "my-agent")
+        agent_name = model.split("/", 1)[1] if "/" in model else None
+        
+        # Only lookup if agent name exists and some config is missing
+        if not agent_name or (api_base is not None and api_key is not None and headers is not None):
+            return api_base, api_key, headers
+        
+        # Try registry lookup (only available in proxy context)
+        try:
+            from litellm.proxy.agent_endpoints.agent_registry import (
+                global_agent_registry,
+            )
+            
+            agent = global_agent_registry.get_agent_by_name(agent_name)
+            if agent:
+                # Get api_base from agent card URL
+                if api_base is None and agent.agent_card_params:
+                    api_base = agent.agent_card_params.get("url")
+                
+                # Get api_key, headers, and other params from litellm_params
+                if agent.litellm_params:
+                    if api_key is None:
+                        api_key = agent.litellm_params.get("api_key")
+                    
+                    if headers is None:
+                        agent_headers = agent.litellm_params.get("headers")
+                        if agent_headers:
+                            headers = agent_headers
+                    
+                    # Merge other litellm_params (timeout, max_retries, etc.)
+                    for key, value in agent.litellm_params.items():
+                        if key not in ["api_key", "api_base", "headers", "model"] and key not in optional_params:
+                            optional_params[key] = value
+        except ImportError:
+            pass  # Registry not available (not running in proxy context)
+        
+        return api_base, api_key, headers
+    
+    def get_supported_openai_params(self, model: str) -> List[str]:
+        """Return list of supported OpenAI parameters"""
+        return [
+            "stream",
+            "temperature",
+            "max_tokens",
+            "top_p",
+        ]
+    
+    def map_openai_params(
+        self,
+        non_default_params: dict,
+        optional_params: dict,
+        model: str,
+        drop_params: bool,
+    ) -> dict:
+        """
+        Map OpenAI parameters to A2A parameters.
+        
+        For A2A protocol, we need to map the stream parameter so
+        transform_request can determine which JSON-RPC method to use.
+        """
+        # Map stream parameter
+        for param, value in non_default_params.items():
+            if param == "stream" and value is True:
+                optional_params["stream"] = value
+        
+        return optional_params
+    
+    def validate_environment(
+        self,
+        headers: dict,
+        model: str,
+        messages: List[AllMessageValues],
+        optional_params: dict,
+        litellm_params: dict,
+        api_key: Optional[str] = None,
+        api_base: Optional[str] = None,
+    ) -> dict:
+        """
+        Validate environment and set headers for A2A requests.
+        
+        Args:
+            headers: Request headers dict
+            model: Model name
+            messages: Messages list
+            optional_params: Optional parameters
+            litellm_params: LiteLLM parameters
+            api_key: API key (optional for A2A)
+            api_base: API base URL
+        
+        Returns:
+            Updated headers dict
+        """
+        # Ensure Content-Type is set to application/json for JSON-RPC 2.0
+        if "content-type" not in headers and "Content-Type" not in headers:
+            headers["Content-Type"] = "application/json"
+        
+        # Add Authorization header if API key is provided
+        if api_key is not None:
+            headers["Authorization"] = f"Bearer {api_key}"
+        
+        return headers
+    
+    def get_complete_url(
+        self,
+        api_base: Optional[str],
+        api_key: Optional[str],
+        model: str,
+        optional_params: dict,
+        litellm_params: dict,
+        stream: Optional[bool] = None,
+    ) -> str:
+        """
+        Get the complete A2A agent endpoint URL.
+        
+        A2A agents use JSON-RPC 2.0 at the base URL, not specific paths.
+        The method (message/send or message/stream) is specified in the
+        JSON-RPC request body, not in the URL.
+        
+        Args:
+            api_base: Base URL of the A2A agent (e.g., "http://0.0.0.0:9999")
+            api_key: API key (not used for URL construction)
+            model: Model name (not used for A2A, agent determined by api_base)
+            optional_params: Optional parameters
+            litellm_params: LiteLLM parameters
+            stream: Whether this is a streaming request (affects JSON-RPC method)
+        
+        Returns:
+            Complete URL for the A2A endpoint (base URL)
+        """
+        if api_base is None:
+            raise ValueError("api_base is required for A2A provider")
+        
+        # A2A uses JSON-RPC 2.0 at the base URL
+        # Remove trailing slash for consistency
+        return api_base.rstrip("/")
+    
+    def transform_request(
+        self,
+        model: str,
+        messages: List[AllMessageValues],
+        optional_params: dict,
+        litellm_params: dict,
+        headers: dict,
+    ) -> dict:
+        """
+        Transform OpenAI request to A2A JSON-RPC 2.0 format.
+        
+        Args:
+            model: Model name
+            messages: List of OpenAI messages
+            optional_params: Optional parameters
+            litellm_params: LiteLLM parameters
+            headers: Request headers
+        
+        Returns:
+            A2A JSON-RPC 2.0 request dict
+        """
+        # Generate request ID
+        request_id = str(uuid.uuid4())
+        
+        if not messages:
+            raise ValueError("At least one message is required for A2A completion")
+        
+        # Convert all messages to maintain conversation history
+        # Use helper to format conversation with role prefixes
+        full_context = convert_messages_to_prompt(messages)
+        
+        # Create single A2A message with full conversation context
+        a2a_message = {
+            "role": "user",
+            "parts": [{"kind": "text", "text": full_context}],
+            "messageId": str(uuid.uuid4()),
+        }
+        
+        # Build JSON-RPC 2.0 request
+        # For A2A protocol, the method is "message/send" for non-streaming
+        # and "message/stream" for streaming
+        stream = optional_params.get("stream", False)
+        method = "message/stream" if stream else "message/send"
+        
+        request_data = {
+            "jsonrpc": "2.0",
+            "id": request_id,
+            "method": method,
+            "params": {
+                "message": a2a_message
+            }
+        }
+        
+        return request_data
+    
+    def transform_response(
+        self,
+        model: str,
+        raw_response: httpx.Response,
+        model_response: ModelResponse,
+        logging_obj: Any,
+        request_data: dict,
+        messages: List[AllMessageValues],
+        optional_params: dict,
+        litellm_params: dict,
+        encoding: Any,
+        api_key: Optional[str] = None,
+        json_mode: Optional[bool] = None,
+    ) -> ModelResponse:
+        """
+        Transform A2A JSON-RPC 2.0 response to OpenAI format.
+        
+        Args:
+            model: Model name
+            raw_response: HTTP response from A2A agent
+            model_response: Model response object to populate
+            logging_obj: Logging object
+            request_data: Original request data
+            messages: Original messages
+            optional_params: Optional parameters
+            litellm_params: LiteLLM parameters
+            encoding: Encoding object
+            api_key: API key
+            json_mode: JSON mode flag
+        
+        Returns:
+            Populated ModelResponse object
+        """
+        try:
+            response_json = raw_response.json()
+        except Exception as e:
+            raise A2AError(
+                status_code=raw_response.status_code,
+                message=f"Failed to parse A2A response: {str(e)}",
+                headers=dict(raw_response.headers),
+            )
+        
+        # Check for JSON-RPC error
+        if "error" in response_json:
+            error = response_json["error"]
+            raise A2AError(
+                status_code=raw_response.status_code,
+                message=f"A2A error: {error.get('message', 'Unknown error')}",
+                headers=dict(raw_response.headers),
+            )
+        
+        # Extract text from A2A response
+        text = extract_text_from_a2a_response(response_json)
+        
+        # Populate model response
+        model_response.choices = [
+            Choices(
+                finish_reason="stop",
+                index=0,
+                message=Message(
+                    content=text,
+                    role="assistant",
+                ),
+            )
+        ]
+        
+        # Set model
+        model_response.model = model
+        
+        # Set ID from response
+        model_response.id = response_json.get("id", str(uuid.uuid4()))
+        
+        return model_response
+    
+    def get_model_response_iterator(
+        self,
+        streaming_response: Union[Iterator, Any],
+        sync_stream: bool,
+        json_mode: Optional[bool] = False,
+    ) -> BaseModelResponseIterator:
+        """
+        Get streaming iterator for A2A responses.
+        
+        Args:
+            streaming_response: Streaming response iterator
+            sync_stream: Whether this is a sync stream
+            json_mode: JSON mode flag
+        
+        Returns:
+            A2A streaming iterator
+        """
+        return A2AModelResponseIterator(
+            streaming_response=streaming_response,
+            sync_stream=sync_stream,
+            json_mode=json_mode,
+        )
+    
+    def _openai_message_to_a2a_message(self, message: Dict[str, Any]) -> Dict[str, Any]:
+        """
+        Convert OpenAI message to A2A message format.
+        
+        Args:
+            message: OpenAI message dict
+        
+        Returns:
+            A2A message dict
+        """
+        content = message.get("content", "")
+        role = message.get("role", "user")
+        
+        return {
+            "role": role,
+            "parts": [{"kind": "text", "text": str(content)}],
+            "messageId": str(uuid.uuid4()),
+        }
+    
+    def get_error_class(
+        self, error_message: str, status_code: int, headers: Union[dict, httpx.Headers]
+    ) -> BaseLLMException:
+        """Return appropriate error class for A2A errors"""
+        # Convert headers to dict if needed
+        headers_dict = dict(headers) if isinstance(headers, httpx.Headers) else headers
+        return A2AError(
+            status_code=status_code,
+            message=error_message,
+            headers=headers_dict,
+        )

litellm/llms/a2a/common_utils.py

@@ -0,0 +1,152 @@
+"""
+Common utilities for A2A (Agent-to-Agent) Protocol
+"""
+from typing import Any, Dict, List
+
+from pydantic import BaseModel
+
+from litellm.litellm_core_utils.prompt_templates.common_utils import (
+    convert_content_list_to_str,
+)
+from litellm.llms.base_llm.chat.transformation import BaseLLMException
+from litellm.types.llms.openai import AllMessageValues
+
+
+class A2AError(BaseLLMException):
+    """Base exception for A2A protocol errors"""
+
+    def __init__(
+        self,
+        status_code: int,
+        message: str,
+        headers: Dict[str, Any] = {},
+    ):
+        super().__init__(
+            status_code=status_code,
+            message=message,
+            headers=headers,
+        )
+
+
+def convert_messages_to_prompt(messages: List[AllMessageValues]) -> str:
+    """
+    Convert OpenAI messages to a single prompt string for A2A agent.
+    
+    Formats each message as "{role}: {content}" and joins with newlines
+    to preserve conversation history. Handles both string and list content.
+    
+    Args:
+        messages: List of OpenAI-format messages
+    
+    Returns:
+        Formatted prompt string with full conversation context
+    """
+    conversation_parts = []
+    for msg in messages:
+        # Use LiteLLM's helper to extract text from content (handles both str and list)
+        content_text = convert_content_list_to_str(message=msg)
+        
+        # Get role
+        if isinstance(msg, BaseModel):
+            role = msg.model_dump().get("role", "user")
+        elif isinstance(msg, dict):
+            role = msg.get("role", "user")
+        else:
+            role = dict(msg).get("role", "user")  # type: ignore
+        
+        if content_text:
+            conversation_parts.append(f"{role}: {content_text}")
+    
+    return "\n".join(conversation_parts)
+
+
+def extract_text_from_a2a_message(
+    message: Dict[str, Any], depth: int = 0, max_depth: int = 10
+) -> str:
+    """
+    Extract text content from A2A message parts.
+    
+    Args:
+        message: A2A message dict with 'parts' containing text parts
+        depth: Current recursion depth (internal use)
+        max_depth: Maximum recursion depth to prevent infinite loops
+    
+    Returns:
+        Concatenated text from all text parts
+    """
+    if message is None or depth >= max_depth:
+        return ""
+    
+    parts = message.get("parts", [])
+    text_parts: List[str] = []
+    
+    for part in parts:
+        if part.get("kind") == "text":
+            text_parts.append(part.get("text", ""))
+        # Handle nested parts if they exist
+        elif "parts" in part:
+            nested_text = extract_text_from_a2a_message(part, depth + 1, max_depth)
+            if nested_text:
+                text_parts.append(nested_text)
+    
+    return " ".join(text_parts)
+
+
+def extract_text_from_a2a_response(
+    response_dict: Dict[str, Any], max_depth: int = 10
+) -> str:
+    """
+    Extract text content from A2A response result.
+    
+    Args:
+        response_dict: A2A response dict with 'result' containing message
+        max_depth: Maximum recursion depth to prevent infinite loops
+    
+    Returns:
+        Text from response message parts
+    """
+    result = response_dict.get("result", {})
+    if not isinstance(result, dict):
+        return ""
+    
+    # A2A response can have different formats:
+    # 1. Direct message: {"result": {"kind": "message", "parts": [...]}}
+    # 2. Nested message: {"result": {"message": {"parts": [...]}}}
+    # 3. Task with artifacts: {"result": {"kind": "task", "artifacts": [{"parts": [...]}]}}
+    # 4. Task with status message: {"result": {"kind": "task", "status": {"message": {"parts": [...]}}}}
+    # 5. Streaming artifact-update: {"result": {"kind": "artifact-update", "artifact": {"parts": [...]}}}
+    
+    # Check if result itself has parts (direct message)
+    if "parts" in result:
+        return extract_text_from_a2a_message(result, depth=0, max_depth=max_depth)
+    
+    # Check for nested message
+    message = result.get("message")
+    if message:
+        return extract_text_from_a2a_message(message, depth=0, max_depth=max_depth)
+    
+    # Check for streaming artifact-update (singular artifact)
+    artifact = result.get("artifact")
+    if artifact and isinstance(artifact, dict):
+        return extract_text_from_a2a_message(
+            artifact, depth=0, max_depth=max_depth
+        )
+    
+    # Check for task status message (common in Gemini A2A agents)
+    status = result.get("status", {})
+    if isinstance(status, dict):
+        status_message = status.get("message")
+        if status_message:
+            return extract_text_from_a2a_message(
+                status_message, depth=0, max_depth=max_depth
+            )
+    
+    # Handle task result with artifacts (plural, array)
+    artifacts = result.get("artifacts", [])
+    if artifacts and len(artifacts) > 0:
+        first_artifact = artifacts[0]
+        return extract_text_from_a2a_message(
+            first_artifact, depth=0, max_depth=max_depth
+        )
+    
+    return ""

litellm/llms/fireworks_ai/chat/transformation.py

@@ -236,6 +236,10 @@ class FireworksAIConfig(OpenAIGPTConfig):
                                 disable_add_transform_inline_image_block=disable_add_transform_inline_image_block,
                             )
             filter_value_from_dict(cast(dict, message), "cache_control")
+            # Remove fields not permitted by FireworksAI that may cause:
+            # "Not permitted, field: 'messages[n].provider_specific_fields'"
+            if isinstance(message, dict) and "provider_specific_fields" in message:
+                message.pop("provider_specific_fields", None)
 
         return messages
 

litellm/llms/gemini/files/transformation.py

@@ -210,7 +210,7 @@ class GoogleAIStudioFilesHandler(GeminiModelInfo, BaseFilesConfig):
         We expect file_id to be the URI (e.g. https://generativelanguage.googleapis.com/v1beta/files/...)
         as returned by the upload response.
         """
-        api_key = litellm_params.get("api_key")
+        api_key = litellm_params.get("api_key") or self.get_api_key()
         if not api_key:
             raise ValueError("api_key is required")
 
@@ -222,7 +222,8 @@ class GoogleAIStudioFilesHandler(GeminiModelInfo, BaseFilesConfig):
             api_base = api_base.rstrip("/")
             url = "{}/v1beta/{}?key={}".format(api_base, file_id, api_key)
 
-        return url, {"Content-Type": "application/json"}
+        # Return empty params dict - API key is already in URL, no query params needed
+        return url, {}
 
     def transform_retrieve_file_response(
         self,
@@ -299,7 +300,7 @@ class GoogleAIStudioFilesHandler(GeminiModelInfo, BaseFilesConfig):
             # Extract the file path from full URI
             file_name = file_id.split("/v1beta/")[-1]
         else:
-            file_name = file_id
+            file_name = file_id if file_id.startswith("files/") else f"files/{file_id}"
         
         # Construct the delete URL
         url = f"{api_base}/v1beta/{file_name}"

litellm/llms/openai/common_utils.py

@@ -22,7 +22,6 @@ from litellm.llms.custom_httpx.http_handler import (
     AsyncHTTPHandler,
     get_ssl_configuration,
 )
-from litellm.types.utils import LlmProviders
 
 
 class OpenAIError(BaseLLMException):
@@ -205,67 +204,30 @@ class BaseOpenAILLM:
         if litellm.aclient_session is not None:
             return litellm.aclient_session
 
-        # Use the global cached client system to prevent memory leaks (issue #14540)
-        # This routes through get_async_httpx_client() which provides TTL-based caching
-        from litellm.llms.custom_httpx.http_handler import get_async_httpx_client
+        # Get unified SSL configuration
+        ssl_config = get_ssl_configuration()
 
-        try:
-            # Get SSL config and include in params for proper cache key
-            ssl_config = get_ssl_configuration()
-            params = {"ssl_verify": ssl_config} if ssl_config is not None else {}
-            params["disable_aiohttp_transport"] = litellm.disable_aiohttp_transport
-
-            # Get a cached AsyncHTTPHandler which manages the httpx.AsyncClient
-            cached_handler = get_async_httpx_client(
-                llm_provider=LlmProviders.OPENAI,  # Cache key includes provider
-                params=params,  # Include SSL config in cache key
+        return httpx.AsyncClient(
+            verify=ssl_config,
+            transport=AsyncHTTPHandler._create_async_transport(
+                ssl_context=ssl_config
+                if isinstance(ssl_config, ssl.SSLContext)
+                else None,
+                ssl_verify=ssl_config if isinstance(ssl_config, bool) else None,
                 shared_session=shared_session,
-            )
-            # Return the underlying httpx client from the handler
-            return cached_handler.client
-        except (ImportError, AttributeError, KeyError) as e:
-            # Fallback to creating a client directly if caching system unavailable
-            # This preserves backwards compatibility
-            verbose_logger.debug(
-                f"Client caching unavailable ({type(e).__name__}), using direct client creation"
-            )
-            ssl_config = get_ssl_configuration()
-            return httpx.AsyncClient(
-                verify=ssl_config,
-                transport=AsyncHTTPHandler._create_async_transport(
-                    ssl_context=ssl_config
-                    if isinstance(ssl_config, ssl.SSLContext)
-                    else None,
-                    ssl_verify=ssl_config if isinstance(ssl_config, bool) else None,
-                    shared_session=shared_session,
-                ),
-                follow_redirects=True,
-            )
+            ),
+            follow_redirects=True,
+        )
 
     @staticmethod
     def _get_sync_http_client() -> Optional[httpx.Client]:
         if litellm.client_session is not None:
             return litellm.client_session
 
-        # Use the global cached client system to prevent memory leaks (issue #14540)
-        from litellm.llms.custom_httpx.http_handler import _get_httpx_client
+        # Get unified SSL configuration
+        ssl_config = get_ssl_configuration()
 
-        try:
-            # Get SSL config and include in params for proper cache key
-            ssl_config = get_ssl_configuration()
-            params = {"ssl_verify": ssl_config} if ssl_config is not None else None
-
-            # Get a cached HTTPHandler which manages the httpx.Client
-            cached_handler = _get_httpx_client(params=params)
-            # Return the underlying httpx client from the handler
-            return cached_handler.client
-        except (ImportError, AttributeError, KeyError) as e:
-            # Fallback to creating a client directly if caching system unavailable
-            verbose_logger.debug(
-                f"Client caching unavailable ({type(e).__name__}), using direct client creation"
-            )
-            ssl_config = get_ssl_configuration()
-            return httpx.Client(
-                verify=ssl_config,
-                follow_redirects=True,
-            )
+        return httpx.Client(
+            verify=ssl_config,
+            follow_redirects=True,
+        )

litellm/main.py

@@ -2206,6 +2206,48 @@ def completion(  # type: ignore # noqa: PLR0915
                 logging_obj=logging,  # model call logging done inside the class as we make need to modify I/O to fit aleph alpha's requirements
                 client=client,
             )
+        elif custom_llm_provider == "a2a":
+            # A2A (Agent-to-Agent) Protocol
+            # Resolve agent configuration from registry if model format is "a2a/<agent-name>"
+            api_base, api_key, headers = litellm.A2AConfig.resolve_agent_config_from_registry(
+                model=model,
+                api_base=api_base,
+                api_key=api_key,
+                headers=headers,
+                optional_params=optional_params,
+            )
+            
+            # Fall back to environment variables and defaults
+            api_base = api_base or litellm.api_base or get_secret_str("A2A_API_BASE")
+            
+            if api_base is None:
+                raise Exception(
+                    "api_base is required for A2A provider. "
+                    "Either provide api_base parameter, set A2A_API_BASE environment variable, "
+                    "or register the agent in the proxy with model='a2a/<agent-name>'."
+                )
+
+            headers = headers or litellm.headers
+
+            response = base_llm_http_handler.completion(
+                model=model,
+                stream=stream,
+                messages=messages,
+                acompletion=acompletion,
+                api_base=api_base,
+                model_response=model_response,
+                optional_params=optional_params,
+                litellm_params=litellm_params,
+                shared_session=shared_session,
+                custom_llm_provider=custom_llm_provider,
+                timeout=timeout,
+                headers=headers,
+                encoding=_get_encoding(),
+                api_key=api_key,
+                logging_obj=logging,
+                client=client,
+                provider_config=provider_config,
+            )
         elif custom_llm_provider == "gigachat":
             # GigaChat - Sber AI's LLM (Russia)
             api_key = (
@@ -3134,8 +3176,8 @@ def completion(  # type: ignore # noqa: PLR0915
                 api_key
                 or litellm.api_key
                 or litellm.openrouter_key
-                or get_secret("OPENROUTER_API_KEY")
-                or get_secret("OR_API_KEY")
+                or get_secret_str("OPENROUTER_API_KEY")
+                or get_secret_str("OR_API_KEY")
             )
 
             openrouter_site_url = get_secret("OR_SITE_URL") or "https://litellm.ai"
@@ -4912,8 +4954,8 @@ def embedding(  # noqa: PLR0915
                 api_key
                 or litellm.api_key
                 or litellm.openrouter_key
-                or get_secret("OPENROUTER_API_KEY")
-                or get_secret("OR_API_KEY")
+                or get_secret_str("OPENROUTER_API_KEY")
+                or get_secret_str("OR_API_KEY")
             )
 
             openrouter_site_url = get_secret("OR_SITE_URL") or "https://litellm.ai"

litellm/model_prices_and_context_window_backup.json

@@ -27113,6 +27113,34 @@
         "supports_reasoning": true,
         "supports_tool_choice": true
     },
+    "together_ai/zai-org/GLM-4.7": {
+        "input_cost_per_token": 4.5e-07,
+        "litellm_provider": "together_ai",
+        "max_input_tokens": 200000,
+        "max_output_tokens": 200000,
+        "max_tokens": 200000,
+        "mode": "chat",
+        "output_cost_per_token": 2e-06,
+        "source": "https://www.together.ai/models/glm-4-7",
+        "supports_function_calling": true,
+        "supports_parallel_function_calling": true,
+        "supports_reasoning": true,
+        "supports_tool_choice": true
+    },
+    "together_ai/moonshotai/Kimi-K2.5": {
+        "input_cost_per_token": 5e-07,
+        "litellm_provider": "together_ai",
+        "max_input_tokens": 256000,
+        "max_output_tokens": 256000,
+        "max_tokens": 256000,
+        "mode": "chat",
+        "output_cost_per_token": 2.8e-06,
+        "source": "https://www.together.ai/models/kimi-k2-5",
+        "supports_function_calling": true,
+        "supports_tool_choice": true,
+        "supports_vision": true,
+        "supports_reasoning": true
+    },
     "together_ai/moonshotai/Kimi-K2-Instruct-0905": {
         "input_cost_per_token": 1e-06,
         "litellm_provider": "together_ai",

litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py

@@ -387,6 +387,9 @@ class MCPRequestHandler:
             user_api_key_cache,
         )
 
+        verbose_logger.debug(
+            f"MCP team permission lookup: team_id={user_api_key_auth.team_id if user_api_key_auth else None}"
+        )
         if not user_api_key_auth or not user_api_key_auth.team_id or not prisma_client:
             return None
 

litellm/proxy/_types.py

@@ -3681,7 +3681,7 @@ class LiteLLM_JWTAuth(LiteLLMPydanticObjectBase):
     team_id_upsert: bool = False
     team_ids_jwt_field: Optional[str] = None
     upsert_sso_user_to_team: bool = False
-    team_allowed_routes: List[str] = ["openai_routes", "info_routes"]
+    team_allowed_routes: List[str] = ["openai_routes", "info_routes", "mcp_routes"]
     team_id_default: Optional[str] = Field(
         default=None,
         description="If no team_id given, default permissions/spend-tracking to this team.s",

litellm/proxy/agent_endpoints/a2a_routing.py

@@ -0,0 +1,53 @@
+"""
+A2A Agent Routing
+
+Handles routing for A2A agents (models with "a2a/<agent-name>" prefix).
+Looks up agents in the registry and injects their API base URL.
+"""
+
+from typing import Any, Optional
+
+import litellm
+from litellm._logging import verbose_proxy_logger
+
+
+async def route_a2a_agent_request(data: dict, route_type: str) -> Optional[Any]:
+    """
+    Route A2A agent requests directly to litellm with injected API base.
+    
+    Returns None if not an A2A request (allows normal routing to continue).
+    """
+    # Import here to avoid circular imports
+    from litellm.proxy.agent_endpoints.agent_registry import global_agent_registry
+    from litellm.proxy.route_llm_request import (
+        ROUTE_ENDPOINT_MAPPING,
+        ProxyModelNotFoundError,
+    )
+    
+    model_name = data.get("model", "")
+    
+    # Check if this is an A2A agent request
+    if not isinstance(model_name, str) or not model_name.startswith("a2a/"):
+        return None
+    
+    # Extract agent name (e.g., "a2a/my-agent" -> "my-agent")
+    agent_name = model_name[4:]
+    
+    # Look up agent in registry
+    agent = global_agent_registry.get_agent_by_name(agent_name)
+    if agent is None:
+        verbose_proxy_logger.error(f"[A2A] Agent '{agent_name}' not found in registry")
+        route_name = ROUTE_ENDPOINT_MAPPING.get(route_type, route_type)
+        raise ProxyModelNotFoundError(route=route_name, model_name=model_name)
+    
+    # Get API base URL from agent config
+    if not agent.agent_card_params or "url" not in agent.agent_card_params:
+        verbose_proxy_logger.error(f"[A2A] Agent '{agent_name}' has no URL configured")
+        route_name = ROUTE_ENDPOINT_MAPPING.get(route_type, route_type)
+        raise ProxyModelNotFoundError(route=route_name, model_name=model_name)
+    
+    # Inject API base and route to litellm
+    data["api_base"] = agent.agent_card_params["url"]
+    verbose_proxy_logger.debug(f"[A2A] Routing {model_name} to {data['api_base']}")
+    
+    return getattr(litellm, f"{route_type}")(**data)

litellm/proxy/agent_endpoints/model_list_helpers.py

@@ -0,0 +1,96 @@
+"""
+Helper functions for appending A2A agents to model lists.
+
+Used by proxy model endpoints to make agents appear in UI alongside models.
+"""
+from typing import List
+
+from litellm._logging import verbose_proxy_logger
+from litellm.proxy._types import UserAPIKeyAuth
+from litellm.types.proxy.management_endpoints.model_management_endpoints import (
+    ModelGroupInfoProxy,
+)
+
+
+async def append_agents_to_model_group(
+    model_groups: List[ModelGroupInfoProxy],
+    user_api_key_dict: UserAPIKeyAuth,
+) -> List[ModelGroupInfoProxy]:
+    """
+    Append A2A agents to model groups list for UI display.
+    
+    Converts agents to model format with "a2a/<agent-name>" naming
+    so they appear in playground and work with LiteLLM routing.
+    """
+    try:
+        from litellm.proxy.agent_endpoints.agent_registry import global_agent_registry
+        from litellm.proxy.agent_endpoints.auth.agent_permission_handler import (
+            AgentRequestHandler,
+        )
+
+        allowed_agent_ids = await AgentRequestHandler.get_allowed_agents(
+            user_api_key_auth=user_api_key_dict
+        )
+        
+        for agent_id in allowed_agent_ids:
+            agent = global_agent_registry.get_agent_by_id(agent_id)
+            if agent is not None:
+                model_groups.append(
+                    ModelGroupInfoProxy(
+                        model_group=f"a2a/{agent.agent_name}",
+                        mode="chat",
+                        providers=["a2a"],
+                    )
+                )
+    except Exception as e:
+        verbose_proxy_logger.debug(
+            f"Error appending agents to model_group/info: {e}"
+        )
+    
+    return model_groups
+
+
+async def append_agents_to_model_info(
+    models: List[dict],
+    user_api_key_dict: UserAPIKeyAuth,
+) -> List[dict]:
+    """
+    Append A2A agents to model info list for UI display.
+    
+    Converts agents to model format with "a2a/<agent-name>" naming
+    so they appear in models page and work with LiteLLM routing.
+    """
+    try:
+        from litellm.proxy.agent_endpoints.agent_registry import global_agent_registry
+        from litellm.proxy.agent_endpoints.auth.agent_permission_handler import (
+            AgentRequestHandler,
+        )
+
+        allowed_agent_ids = await AgentRequestHandler.get_allowed_agents(
+            user_api_key_auth=user_api_key_dict
+        )
+        
+        for agent_id in allowed_agent_ids:
+            agent = global_agent_registry.get_agent_by_id(agent_id)
+            if agent is not None:
+                models.append({
+                    "model_name": f"a2a/{agent.agent_name}",
+                    "litellm_params": {
+                        "model": f"a2a/{agent.agent_name}",
+                        "custom_llm_provider": "a2a",
+                    },
+                    "model_info": {
+                        "id": agent.agent_id,
+                        "mode": "chat",
+                        "db_model": True,
+                        "created_by": agent.created_by,
+                        "created_at": agent.created_at,
+                        "updated_at": agent.updated_at,
+                    },
+                })
+    except Exception as e:
+        verbose_proxy_logger.debug(
+            f"Error appending agents to v2/model/info: {e}"
+        )
+    
+    return models

litellm/proxy/auth/handle_jwt.py

@@ -976,6 +976,9 @@ class JWTAuthManager:
                             user_route=route,
                             litellm_proxy_roles=jwt_handler.litellm_jwtauth,
                         )
+                        verbose_proxy_logger.debug(
+                            f"JWT team route check: team_id={team_id}, route={route}, is_allowed={is_allowed}"
+                        )
                         if is_allowed:
                             return team_id, team_object
             except Exception:

litellm/proxy/management_endpoints/internal_user_endpoints.py

@@ -813,9 +813,12 @@ def _update_internal_user_params(
     data_json: dict, data: Union[UpdateUserRequest, UpdateUserRequestNoUserIDorEmail]
 ) -> dict:
     non_default_values = {}
+    fields_set = data.fields_set() if hasattr(data, 'fields_set') else set()
+    
     for k, v in data_json.items():
         if k == "max_budget":
-            non_default_values[k] = v
+            if "max_budget" in fields_set:
+                non_default_values[k] = v
         elif (
             v is not None
             and v

litellm/proxy/proxy_server.py

@@ -239,6 +239,10 @@ from litellm.proxy._types import *
 from litellm.proxy.agent_endpoints.a2a_endpoints import router as a2a_router
 from litellm.proxy.agent_endpoints.agent_registry import global_agent_registry
 from litellm.proxy.agent_endpoints.endpoints import router as agent_endpoints_router
+from litellm.proxy.agent_endpoints.model_list_helpers import (
+    append_agents_to_model_group,
+    append_agents_to_model_info,
+)
 from litellm.proxy.analytics_endpoints.analytics_endpoints import (
     router as analytics_router,
 )
@@ -8616,6 +8620,15 @@ async def model_info_v2(
         )
 
     verbose_proxy_logger.debug("all_models: %s", all_models)
+    
+    # Append A2A agents to models list
+    all_models = await append_agents_to_model_info(
+        models=all_models,
+        user_api_key_dict=user_api_key_dict,
+    )
+    
+    # Update total count to include agents
+    search_total_count = len(all_models)
 
     return _paginate_models_response(
         all_models=all_models,
@@ -9456,6 +9469,12 @@ async def model_group_info(
     model_groups: List[ModelGroupInfoProxy] = _get_model_group_info(
         llm_router=llm_router, all_models_str=all_models_str, model_group=model_group
     )
+    
+    # Append A2A agents to model groups
+    model_groups = await append_agents_to_model_group(
+        model_groups=model_groups,
+        user_api_key_dict=user_api_key_dict,
+    )
 
     return {"data": model_groups}
 

litellm/proxy/route_llm_request.py

@@ -12,6 +12,11 @@ else:
     LitellmRouter = Any
 
 
+def _is_a2a_agent_model(model_name: Any) -> bool:
+    """Check if the model name is for an A2A agent (a2a/ prefix)."""
+    return isinstance(model_name, str) and model_name.startswith("a2a/")
+
+
 ROUTE_ENDPOINT_MAPPING = {
     "acompletion": "/chat/completions",
     "atext_completion": "/completions",
@@ -322,6 +327,12 @@ async def route_request(
                 except Exception:
                     # If router fails (e.g., model not found in router), fall back to direct call
                     return getattr(litellm, f"{route_type}")(**data)
+            elif _is_a2a_agent_model(data.get("model", "")):
+                from litellm.proxy.agent_endpoints.a2a_routing import (
+                    route_a2a_agent_request,
+                )
+                
+                return await route_a2a_agent_request(data, route_type)
 
     elif user_model is not None:
         return getattr(litellm, f"{route_type}")(**data)

litellm/types/utils.py

@@ -3029,6 +3029,7 @@ class LlmProviders(str, Enum):
     MISTRAL = "mistral"
     MILVUS = "milvus"
     GROQ = "groq"
+    A2A = "a2a"
     GIGACHAT = "gigachat"
     NVIDIA_NIM = "nvidia_nim"
     CEREBRAS = "cerebras"

litellm/utils.py

@@ -1453,6 +1453,10 @@ def client(original_function):  # noqa: PLR0915
                 logging_obj, kwargs = function_setup(
                     original_function.__name__, rules_obj, start_time, *args, **kwargs
                 )
+            
+            # Type assertion: logging_obj is guaranteed to be non-None after function_setup
+            assert logging_obj is not None, "logging_obj should not be None after function_setup"
+            
             ## LOAD CREDENTIALS
             load_credentials_from_list(kwargs)
             kwargs["litellm_logging_obj"] = logging_obj
@@ -1771,6 +1775,9 @@ def client(original_function):  # noqa: PLR0915
                 logging_obj, kwargs = function_setup(
                     original_function.__name__, rules_obj, start_time, *args, **kwargs
                 )
+            
+            # Type assertion: logging_obj is guaranteed to be non-None after function_setup
+            assert logging_obj is not None, "logging_obj should not be None after function_setup"
 
             modified_kwargs = await async_pre_call_deployment_hook(kwargs, call_type)
             if modified_kwargs is not None:
@@ -7799,6 +7806,7 @@ class ProviderConfigManager:
             # Simple provider mappings (no model parameter needed)
             LlmProviders.DEEPSEEK: (lambda: litellm.DeepSeekChatConfig(), False),
             LlmProviders.GROQ: (lambda: litellm.GroqChatConfig(), False),
+            LlmProviders.A2A: (lambda: litellm.A2AConfig(), False),
             LlmProviders.BYTEZ: (lambda: litellm.BytezChatConfig(), False),
             LlmProviders.DATABRICKS: (lambda: litellm.DatabricksConfig(), False),
             LlmProviders.XAI: (lambda: litellm.XAIChatConfig(), False),

model_prices_and_context_window.json

@@ -27113,6 +27113,34 @@
         "supports_reasoning": true,
         "supports_tool_choice": true
     },
+    "together_ai/zai-org/GLM-4.7": {
+        "input_cost_per_token": 4.5e-07,
+        "litellm_provider": "together_ai",
+        "max_input_tokens": 200000,
+        "max_output_tokens": 200000,
+        "max_tokens": 200000,
+        "mode": "chat",
+        "output_cost_per_token": 2e-06,
+        "source": "https://www.together.ai/models/glm-4-7",
+        "supports_function_calling": true,
+        "supports_parallel_function_calling": true,
+        "supports_reasoning": true,
+        "supports_tool_choice": true
+    },
+    "together_ai/moonshotai/Kimi-K2.5": {
+        "input_cost_per_token": 5e-07,
+        "litellm_provider": "together_ai",
+        "max_input_tokens": 256000,
+        "max_output_tokens": 256000,
+        "max_tokens": 256000,
+        "mode": "chat",
+        "output_cost_per_token": 2.8e-06,
+        "source": "https://www.together.ai/models/kimi-k2-5",
+        "supports_function_calling": true,
+        "supports_tool_choice": true,
+        "supports_vision": true,
+        "supports_reasoning": true
+    },
     "together_ai/moonshotai/Kimi-K2-Instruct-0905": {
         "input_cost_per_token": 1e-06,
         "litellm_provider": "together_ai",

poetry.lock

@@ -1,4 +1,4 @@
-# This file is automatically @generated by Poetry 2.1.4 and should not be changed by hand.
+# This file is automatically @generated by Poetry 2.2.0 and should not be changed by hand.
 
 [[package]]
 name = "a2a-sdk"
@@ -5704,24 +5704,6 @@ pytest = ">=7.0.0"
 [package.extras]
 dev = ["black", "flake8", "isort", "mypy"]
 
-[[package]]
-name = "pytest-retry"
-version = "1.7.0"
-description = "Adds the ability to retry flaky tests in CI environments"
-optional = false
-python-versions = ">=3.9"
-groups = ["dev"]
-files = [
-    {file = "pytest_retry-1.7.0-py3-none-any.whl", hash = "sha256:a2dac85b79a4e2375943f1429479c65beb6c69553e7dae6b8332be47a60954f4"},
-    {file = "pytest_retry-1.7.0.tar.gz", hash = "sha256:f8d52339f01e949df47c11ba9ee8d5b362f5824dff580d3870ec9ae0057df80f"},
-]
-
-[package.dependencies]
-pytest = ">=7.0.0"
-
-[package.extras]
-dev = ["black", "flake8", "isort", "mypy"]
-
 [[package]]
 name = "python-dateutil"
 version = "2.9.0.post0"

provider_endpoints_support.json

@@ -32,6 +32,23 @@
     }
   },
   "providers": {
+    "a2a": {
+      "display_name": "A2A (Agent-to-Agent) (`a2a`)",
+      "url": "https://docs.litellm.ai/docs/providers/a2a",
+      "endpoints": {
+        "chat_completions": true,
+        "messages": false,
+        "responses": false,
+        "embeddings": false,
+        "image_generations": false,
+        "audio_transcriptions": false,
+        "audio_speech": false,
+        "moderations": false,
+        "batches": false,
+        "rerank": false,
+        "a2a": false
+      }
+    },
     "abliteration": {
       "display_name": "Abliteration (`abliteration`)",
       "url": "https://docs.litellm.ai/docs/providers/abliteration",

requirements.txt

@@ -73,4 +73,4 @@ pypdf>=6.6.2 # for PDF text extraction in RAG ingestion
 ########################
 # LITELLM ENTERPRISE DEPENDENCIES
 ########################
-litellm-enterprise==0.1.28
+litellm-enterprise==0.1.29

tests/code_coverage_tests/recursive_detector.py

@@ -40,6 +40,7 @@ IGNORE_FUNCTIONS = [
     "filter_exceptions_from_params",  # max depth set (default 20) to prevent infinite recursion.
     "__getattr__",  # lazy loading pattern in litellm/__init__.py with proper caching to prevent infinite recursion.
     "_validate_inheritance_chain",  # max depth set (default 100) to prevent infinite recursion in policy inheritance validation.
+    "extract_text_from_a2a_message",  # max depth set (default 10) to prevent infinite recursion in A2A message parsing.
 ]
 
 

tests/llm_translation/test_a2a.py

@@ -0,0 +1,132 @@
+"""
+Minimal E2E tests for A2A (Agent-to-Agent) Protocol provider.
+
+Tests validate that the endpoint is reachable and can handle both
+streaming and non-streaming requests.
+"""
+import os
+import sys
+
+import pytest
+
+sys.path.insert(0, os.path.abspath("../.."))
+
+import litellm
+
+
+@pytest.mark.asyncio
+async def test_a2a_completion_async_non_streaming():
+    """
+    Test A2A provider with async non-streaming request.
+    
+    Minimal test to validate endpoint reachability.
+    
+    Note: Requires an A2A agent running at http://0.0.0.0:9999
+    Set A2A_API_BASE environment variable to use a different endpoint.
+    """
+    api_base = os.environ.get("A2A_API_BASE", "http://0.0.0.0:9999")
+    
+    try:
+        response = await litellm.acompletion(
+            model="a2a/test-agent",
+            messages=[{"role": "user", "content": "Hello"}],
+            api_base=api_base,
+            stream=False,
+        )
+        
+        print(f"Response: {response}")
+        assert response is not None, "Expected non-None response"
+        print(f"✅ Async non-streaming test passed")
+        
+    except litellm.exceptions.APIConnectionError as e:
+        pytest.skip(f"A2A agent not reachable at {api_base}: {e}")
+    except Exception as e:
+        pytest.fail(f"Error occurred: {e}")
+
+
+@pytest.mark.asyncio
+async def test_a2a_completion_async_streaming():
+    """
+    Test A2A provider with async streaming request.
+    
+    Minimal test to validate streaming endpoint reachability.
+    """
+    api_base = os.environ.get("A2A_API_BASE", "http://0.0.0.0:9999")
+    
+    try:
+        response = await litellm.acompletion(
+            model="a2a/test-agent",
+            messages=[{"role": "user", "content": "Hello"}],
+            api_base=api_base,
+            stream=True,
+        )
+        
+        chunks = []
+        async for chunk in response:  # type: ignore
+            chunks.append(chunk)
+            print(f"Chunk: {chunk}")
+        
+        assert len(chunks) > 0, "Expected at least one chunk in streaming response"
+        print(f"✅ Async streaming test passed: received {len(chunks)} chunks")
+        
+    except litellm.exceptions.APIConnectionError as e:
+        pytest.skip(f"A2A agent not reachable at {api_base}: {e}")
+    except Exception as e:
+        pytest.fail(f"Error occurred: {e}")
+
+
+def test_a2a_completion_sync():
+    """
+    Test A2A provider with synchronous non-streaming request.
+    
+    Minimal test to validate sync endpoint reachability.
+    """
+    api_base = os.environ.get("A2A_API_BASE", "http://0.0.0.0:9999")
+    
+    try:
+        response = litellm.completion(
+            model="a2a/test-agent",
+            messages=[{"role": "user", "content": "Hello"}],
+            api_base=api_base,
+            stream=False,
+        )
+        
+        print(f"Response: {response}")
+        assert response is not None, "Expected non-None response"
+        print(f"✅ Sync non-streaming test passed")
+        
+    except litellm.exceptions.APIConnectionError as e:
+        pytest.skip(f"A2A agent not reachable at {api_base}: {e}")
+    except Exception as e:
+        pytest.fail(f"Error occurred: {e}")
+
+
+def test_a2a_completion_sync_streaming():
+    """
+    Test A2A provider with synchronous streaming request.
+    
+    Minimal test to validate sync streaming endpoint reachability.
+    """
+    api_base = os.environ.get("A2A_API_BASE", "http://0.0.0.0:9999")
+    
+    try:
+        response = litellm.completion(
+            model="a2a/test-agent",
+            messages=[{"role": "user", "content": "Hello"}],
+            api_base=api_base,
+            stream=True,
+        )
+        
+        chunks = []
+        for chunk in response:  # type: ignore
+            chunks.append(chunk)
+            print(f"Chunk: {chunk}")
+        
+        assert len(chunks) > 0, "Expected at least one chunk in streaming response"
+        print(f"✅ Sync streaming test passed: received {len(chunks)} chunks")
+        
+    except litellm.exceptions.APIConnectionError as e:
+        pytest.skip(f"A2A agent not reachable at {api_base}: {e}")
+    except Exception as e:
+        pytest.fail(f"Error occurred: {e}")
+

tests/test_litellm/completion_extras/test_litellm_responses_transformation_transformation.py

@@ -134,3 +134,25 @@ def test_transform_request_with_response_format():
     assert result["text"]["format"]["type"] == "json_schema"
     assert result["text"]["format"]["name"] == "person_schema"
     assert "schema" in result["text"]["format"]
+
+
+def test_transform_request_includes_extra_headers():
+    """Test that transform_request forwards headers as extra_headers for upstream call."""
+    handler = LiteLLMResponsesTransformationHandler()
+    messages = [{"role": "user", "content": "Hello"}]
+    optional_params = {}
+    litellm_params = {}
+
+    class MockLoggingObj:
+        pass
+
+    headers = {"cf-aig-authorization": "secret-token"}
+    result = handler.transform_request(
+        model="gpt-5-pro",
+        messages=messages,
+        optional_params=optional_params,
+        litellm_params=litellm_params,
+        headers=headers,
+        litellm_logging_obj=MockLoggingObj(),
+    )
+    assert result.get("extra_headers") == headers

tests/test_litellm/llms/fireworks_ai/chat/test_fireworks_ai_chat_transformation.py

@@ -108,3 +108,32 @@ def test_get_supported_openai_params_reasoning_effort():
         "fireworks_ai/accounts/fireworks/models/llama-v3-70b-instruct"
     )
     assert "reasoning_effort" not in unsupported_params
+
+
+def test_transform_messages_helper_removes_provider_specific_fields():
+    """
+    Test that _transform_messages_helper removes provider_specific_fields from messages.
+    """
+    config = FireworksAIConfig()
+    # Simulated messages, as dicts, including provider_specific_fields
+    messages = [
+        {
+            "role": "user",
+            "content": "Hello!",
+            "provider_specific_fields": {"extra": "should be removed"},
+        },
+        {
+            "role": "assistant",
+            "content": "Hi there!",
+            "provider_specific_fields": {"more": "remove this"},
+        },
+        {
+            "role": "user",
+            "content": "How are you?",
+            # no provider_specific_fields
+        }
+    ]
+    # Call helper
+    out = config._transform_messages_helper(messages, model="fireworks/test", litellm_params={})
+    for msg in out:
+        assert "provider_specific_fields" not in msg

tests/test_litellm/llms/gemini/files/__init__.py

@@ -0,0 +1 @@
+"""Tests for Gemini files functionality"""

tests/test_litellm/llms/gemini/files/test_gemini_files_transformation.py

@@ -0,0 +1,298 @@
+"""
+Test Google AI Studio (Gemini) files transformation functionality
+"""
+
+import os
+import pytest
+from unittest.mock import Mock, patch
+
+import httpx
+
+from litellm.llms.gemini.files.transformation import GoogleAIStudioFilesHandler
+from litellm.types.llms.openai import OpenAIFileObject
+
+
+class TestGoogleAIStudioFilesTransformation:
+    """Test Google AI Studio files transformation"""
+
+    def setup_method(self):
+        """Setup test method"""
+        self.handler = GoogleAIStudioFilesHandler()
+
+    def test_transform_retrieve_file_request_with_full_uri(self):
+        """
+        Test that transform_retrieve_file_request returns empty params dict
+        to avoid 'Content-Type' query parameter error
+        
+        Regression test for: https://github.com/BerriAI/litellm/issues/XXX
+        When retrieving a file, the API was incorrectly trying to pass Content-Type
+        as a query parameter, which Gemini API rejected.
+        """
+        file_id = "https://generativelanguage.googleapis.com/v1beta/files/test123"
+        litellm_params = {"api_key": "test-api-key"}
+
+        url, params = self.handler.transform_retrieve_file_request(
+            file_id=file_id,
+            optional_params={},
+            litellm_params=litellm_params,
+        )
+
+        # Verify URL is constructed correctly with API key
+        assert "key=test-api-key" in url
+        assert file_id in url
+
+        # CRITICAL: params should be empty dict, not contain Content-Type or any other params
+        # These would be incorrectly interpreted as query parameters
+        assert params == {}, f"Expected empty params dict, got: {params}"
+        assert "Content-Type" not in params, "Content-Type should not be in query params"
+
+    def test_transform_retrieve_file_request_with_file_name_only(self):
+        """
+        Test that transform_retrieve_file_request handles file_id without full URI
+        """
+        file_id = "files/test123"
+        litellm_params = {"api_key": "test-api-key"}
+
+        url, params = self.handler.transform_retrieve_file_request(
+            file_id=file_id,
+            optional_params={},
+            litellm_params=litellm_params,
+        )
+
+        # Verify URL is constructed correctly
+        assert "generativelanguage.googleapis.com" in url
+        assert file_id in url
+        assert "key=test-api-key" in url
+
+        # CRITICAL: params should be empty dict
+        assert params == {}, f"Expected empty params dict, got: {params}"
+        assert "Content-Type" not in params, "Content-Type should not be in query params"
+
+    @patch.dict('os.environ', {}, clear=True)
+    @patch('litellm.llms.gemini.common_utils.get_secret_str', return_value=None)
+    def test_transform_retrieve_file_request_missing_api_key(self, mock_get_secret):
+        """Test that transform_retrieve_file_request raises error when API key is missing"""
+        file_id = "files/test123"
+        litellm_params = {}
+
+        with pytest.raises(ValueError, match="api_key is required"):
+            self.handler.transform_retrieve_file_request(
+                file_id=file_id,
+                optional_params={},
+                litellm_params=litellm_params,
+            )
+
+    def test_transform_retrieve_file_response_success(self):
+        """Test successful transformation of Gemini file retrieval response"""
+        # Mock response data from Gemini API
+        mock_response_data = {
+            "name": "files/test123",
+            "displayName": "test_file.pdf",
+            "mimeType": "application/pdf",
+            "sizeBytes": "1024",
+            "createTime": "2024-01-15T10:30:00.123456Z",
+            "updateTime": "2024-01-15T10:30:00.123456Z",
+            "expirationTime": "2024-01-17T10:30:00.123456Z",
+            "sha256Hash": "abcd1234",
+            "uri": "https://generativelanguage.googleapis.com/v1beta/files/test123",
+            "state": "ACTIVE",
+        }
+
+        # Create mock httpx response
+        mock_response = Mock(spec=httpx.Response)
+        mock_response.json.return_value = mock_response_data
+
+        # Create mock logging object
+        mock_logging_obj = Mock()
+
+        # Transform response
+        result = self.handler.transform_retrieve_file_response(
+            raw_response=mock_response,
+            logging_obj=mock_logging_obj,
+            litellm_params={},
+        )
+
+        # Verify transformation
+        assert isinstance(result, OpenAIFileObject)
+        assert result.id == mock_response_data["uri"]
+        assert result.filename == mock_response_data["displayName"]
+        assert result.bytes == int(mock_response_data["sizeBytes"])
+        assert result.object == "file"
+        assert result.purpose == "user_data"
+        assert result.status == "processed"  # ACTIVE state maps to processed
+        assert result.status_details is None
+
+    def test_transform_retrieve_file_response_failed_state(self):
+        """Test transformation of Gemini file retrieval response with FAILED state"""
+        mock_response_data = {
+            "name": "files/test123",
+            "displayName": "test_file.pdf",
+            "mimeType": "application/pdf",
+            "sizeBytes": "1024",
+            "createTime": "2024-01-15T10:30:00.123456Z",
+            "uri": "https://generativelanguage.googleapis.com/v1beta/files/test123",
+            "state": "FAILED",
+            "error": {"message": "Upload failed", "code": "INTERNAL"},
+        }
+
+        mock_response = Mock(spec=httpx.Response)
+        mock_response.json.return_value = mock_response_data
+        mock_logging_obj = Mock()
+
+        result = self.handler.transform_retrieve_file_response(
+            raw_response=mock_response,
+            logging_obj=mock_logging_obj,
+            litellm_params={},
+        )
+
+        # Verify error state handling
+        assert result.status == "error"
+        assert result.status_details is not None
+        assert "message" in result.status_details
+
+    def test_transform_retrieve_file_response_processing_state(self):
+        """Test transformation of Gemini file retrieval response with PROCESSING state"""
+        mock_response_data = {
+            "name": "files/test123",
+            "displayName": "test_file.pdf",
+            "mimeType": "application/pdf",
+            "sizeBytes": "1024",
+            "createTime": "2024-01-15T10:30:00.123456Z",
+            "uri": "https://generativelanguage.googleapis.com/v1beta/files/test123",
+            "state": "PROCESSING",
+        }
+
+        mock_response = Mock(spec=httpx.Response)
+        mock_response.json.return_value = mock_response_data
+        mock_logging_obj = Mock()
+
+        result = self.handler.transform_retrieve_file_response(
+            raw_response=mock_response,
+            logging_obj=mock_logging_obj,
+            litellm_params={},
+        )
+
+        # PROCESSING state should map to "uploaded" status
+        assert result.status == "uploaded"
+
+    def test_transform_retrieve_file_response_missing_createTime(self):
+        """
+        Test that transform_retrieve_file_response raises proper error when createTime is missing
+        
+        This tests the error scenario that occurs when API returns an error response
+        without the expected file metadata fields.
+        """
+        # Mock error response from Gemini API (missing createTime)
+        mock_response_data = {
+            "error": {
+                "code": 400,
+                "message": "Invalid request",
+                "status": "INVALID_ARGUMENT",
+            }
+        }
+
+        mock_response = Mock(spec=httpx.Response)
+        mock_response.json.return_value = mock_response_data
+        mock_logging_obj = Mock()
+
+        # Should raise ValueError with helpful message
+        with pytest.raises(ValueError, match="Error parsing file retrieve response"):
+            self.handler.transform_retrieve_file_response(
+                raw_response=mock_response,
+                logging_obj=mock_logging_obj,
+                litellm_params={},
+            )
+
+    def test_validate_environment(self):
+        """Test that validate_environment properly adds API key to headers"""
+        headers = {}
+        api_key = "test-gemini-api-key"
+
+        result_headers = self.handler.validate_environment(
+            headers=headers,
+            model="gemini-pro",
+            messages=[],
+            optional_params={},
+            litellm_params={},
+            api_key=api_key,
+        )
+
+        # Verify API key is added to headers
+        assert "x-goog-api-key" in result_headers
+        assert result_headers["x-goog-api-key"] == api_key
+
+    @patch.dict('os.environ', {}, clear=True)
+    @patch('litellm.llms.gemini.common_utils.get_secret_str', return_value=None)
+    def test_validate_environment_missing_api_key(self, mock_get_secret):
+        """Test that validate_environment raises error when API key is missing"""
+        headers = {}
+
+        with pytest.raises(
+            ValueError, match="GEMINI_API_KEY is required for Google AI Studio file operations"
+        ):
+            self.handler.validate_environment(
+                headers=headers,
+                model="gemini-pro",
+                messages=[],
+                optional_params={},
+                litellm_params={},
+                api_key=None,
+            )
+
+    def test_get_complete_url(self):
+        """Test that get_complete_url constructs proper upload URL"""
+        api_base = "https://generativelanguage.googleapis.com"
+        api_key = "test-api-key"
+        
+        url = self.handler.get_complete_url(
+            api_base=api_base,
+            api_key=api_key,
+            model="gemini-pro",
+            optional_params={},
+            litellm_params={},
+        )
+
+        # Verify URL structure
+        assert api_base in url
+        assert "upload/v1beta/files" in url
+        assert f"key={api_key}" in url
+
+    def test_transform_delete_file_request_with_full_uri(self):
+        """Test delete file request transformation with full URI"""
+        file_id = "https://generativelanguage.googleapis.com/v1beta/files/test123"
+        litellm_params = {
+            "api_key": "test-api-key",
+            "api_base": "https://generativelanguage.googleapis.com",
+        }
+
+        url, params = self.handler.transform_delete_file_request(
+            file_id=file_id,
+            optional_params={},
+            litellm_params=litellm_params,
+        )
+
+        # Verify URL extraction
+        assert "files/test123" in url
+        assert "generativelanguage.googleapis.com" in url
+        
+        # Params should be empty (API key goes in header via validate_environment)
+        assert params == {}
+
+    def test_transform_delete_file_request_with_file_name_only(self):
+        """Test delete file request transformation with file name only"""
+        file_id = "files/test123"
+        litellm_params = {
+            "api_key": "test-api-key",
+            "api_base": "https://generativelanguage.googleapis.com",
+        }
+
+        url, params = self.handler.transform_delete_file_request(
+            file_id=file_id,
+            optional_params={},
+            litellm_params=litellm_params,
+        )
+
+        # Verify URL construction
+        assert file_id in url
+        assert "generativelanguage.googleapis.com" in url
+        assert params == {}

tests/test_litellm/llms/test_lifecycle_fix.py

@@ -0,0 +1,46 @@
+"""
+Verifies that the httpx client used by AsyncOpenAI is NOT closed
+when AsyncHTTPHandler instances are garbage collected.
+"""
+import asyncio
+import gc
+import httpx
+from litellm.llms.openai.common_utils import BaseOpenAILLM
+from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler
+
+
+async def test_httpx_client_not_closed_by_handler_gc():
+    """
+    Before the fix: _get_async_http_client() returned handler.client,
+    so when handler was GC'd its __del__ closed the client.
+    After the fix: returns a standalone httpx.AsyncClient, no handler involved.
+    """
+    # Get the client the same way AsyncOpenAI would
+    client = BaseOpenAILLM._get_async_http_client()
+    assert isinstance(client, httpx.AsyncClient)
+
+    # Simulate what the old code did: create an AsyncHTTPHandler and GC it
+    handler = AsyncHTTPHandler()
+    handler_client = handler.client
+    del handler
+    gc.collect()
+
+    # The client from _get_async_http_client should still be open
+    # because it's NOT tied to any AsyncHTTPHandler
+    assert not client.is_closed, "Client was closed prematurely!"
+
+    # Verify it can actually send (build a request without sending)
+    try:
+        req = client.build_request("GET", "https://example.com")
+        print("PASS: Client is still usable after handler GC")
+    except RuntimeError as e:
+        if "closed" in str(e):
+            print(f"FAIL: {e}")
+            raise
+        raise
+
+    await client.aclose()
+    print("All checks passed!")
+
+
+asyncio.run(test_httpx_client_not_closed_by_handler_gc())

tests/test_litellm/proxy/_experimental/mcp_server/test_jwt_mcp_enforcement.py

@@ -0,0 +1,480 @@
+"""
+Test to verify Team MCP permissions are enforced when using JWT authentication.
+
+Scenario:
+1. Team "ABC" exists with models configured and MCPs assigned
+2. User JWT has team "ABC" in groups (via team_ids_jwt_field)
+3. Call MCP list endpoint
+4. EXPECTED: Team MCP permissions should be enforced
+"""
+
+import pytest
+from unittest.mock import AsyncMock, MagicMock, patch
+
+from litellm.proxy._types import (
+    LiteLLM_JWTAuth,
+    LiteLLM_TeamTable,
+    LiteLLM_ObjectPermissionTable,
+    UserAPIKeyAuth,
+)
+from litellm.proxy.auth.handle_jwt import JWTAuthManager, JWTHandler
+from litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp import (
+    MCPRequestHandler,
+)
+from litellm.proxy._experimental.mcp_server.mcp_server_manager import MCPServerManager
+
+
+@pytest.mark.asyncio
+async def test_reproduce_jwt_mcp_enforcement_issue(monkeypatch):
+    """
+    Reproduce the bug where Team MCP permissions are NOT enforced when using JWT.
+    
+    Setup:
+    - Team "ABC" has models ["gpt-4"] and MCPs ["mcp-server-1"] assigned
+    - JWT has team "ABC" in groups field
+    - User calls MCP list endpoint (no model requested)
+    
+    Expected: team_id should be set to "ABC" so MCP permissions are enforced
+    Actual (BUG): team_id is None because route check fails for MCP routes
+    """
+    from litellm.caching import DualCache
+    from litellm.proxy.utils import ProxyLogging
+    from litellm.router import Router
+
+    # Setup mock router
+    router = Router(model_list=[{"model_name": "gpt-4", "litellm_params": {"model": "gpt-4"}}])
+    import sys
+    import types
+    proxy_server_module = types.ModuleType("proxy_server")
+    proxy_server_module.llm_router = router
+    monkeypatch.setitem(sys.modules, "litellm.proxy.proxy_server", proxy_server_module)
+
+    # Team "ABC" has models configured AND MCPs assigned
+    team_with_mcp = LiteLLM_TeamTable(
+        team_id="ABC",
+        models=["gpt-4"],  # Team HAS models
+        object_permission=LiteLLM_ObjectPermissionTable(
+            object_permission_id="perm-123",
+            mcp_servers=["mcp-server-1"],  # Team has MCPs assigned
+        ),
+    )
+
+    async def mock_get_team_object(*args, **kwargs):
+        team_id = kwargs.get("team_id") or args[0]
+        if team_id == "ABC":
+            return team_with_mcp
+        return None
+
+    monkeypatch.setattr(
+        "litellm.proxy.auth.handle_jwt.get_team_object", mock_get_team_object
+    )
+
+    # Setup JWT handler with team_ids_jwt_field (groups)
+    jwt_handler = JWTHandler()
+    jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth(
+        team_ids_jwt_field="groups",  # Use groups field for teams
+        # NOTE: team_allowed_routes defaults to ["openai_routes", "info_routes"]
+        # which does NOT include "mcp_routes"
+    )
+
+    user_api_key_cache = DualCache()
+    proxy_logging_obj = ProxyLogging(user_api_key_cache=user_api_key_cache)
+
+    # Simulate JWT payload with team in groups
+    jwt_token = {
+        "sub": "user-123",
+        "groups": ["ABC"],  # Team "ABC" is in groups
+        "scope": "",
+    }
+
+    # Mock auth_jwt to return our token
+    with patch.object(jwt_handler, "auth_jwt", new_callable=AsyncMock) as mock_auth_jwt:
+        mock_auth_jwt.return_value = jwt_token
+
+        # Call auth_builder for MCP route (like /mcp/tools/list)
+        result = await JWTAuthManager.auth_builder(
+            api_key="test-jwt-token",
+            jwt_handler=jwt_handler,
+            request_data={},  # No model in request (MCP endpoint)
+            general_settings={},
+            route="/mcp/tools/list",  # MCP route
+            prisma_client=None,
+            user_api_key_cache=user_api_key_cache,
+            parent_otel_span=None,
+            proxy_logging_obj=proxy_logging_obj,
+        )
+
+        # THIS IS THE BUG: team_id should be "ABC" but it's None!
+        print(f"Result team_id: {result['team_id']}")
+        print(f"Result team_object: {result['team_object']}")
+        
+        # The test should FAIL if the bug exists (team_id is None)
+        # If the fix is applied, team_id should be "ABC"
+        assert result["team_id"] == "ABC", (
+            f"BUG: team_id should be 'ABC' but got '{result['team_id']}'. "
+            f"This happens because default team_allowed_routes does not include 'mcp_routes', "
+            f"so allowed_routes_check() fails and the team is skipped in find_team_with_model_access()."
+        )
+
+
+@pytest.mark.asyncio  
+async def test_verify_mcp_routes_in_default_team_allowed_routes():
+    """
+    Verify that mcp_routes IS in the default team_allowed_routes.
+    This is required for team MCP permissions to work with JWT auth.
+    """
+    default_jwt_auth = LiteLLM_JWTAuth()
+    
+    print(f"Default team_allowed_routes: {default_jwt_auth.team_allowed_routes}")
+    
+    # mcp_routes must be in defaults for team MCP permissions to work
+    assert "mcp_routes" in default_jwt_auth.team_allowed_routes, (
+        "mcp_routes must be in default team_allowed_routes for JWT MCP enforcement to work"
+    )
+
+
+@pytest.mark.asyncio
+async def test_mcp_route_check_passes_for_team():
+    """
+    Verify that allowed_routes_check returns True for MCP routes with default settings.
+    This is required for teams to access MCP endpoints with JWT auth.
+    """
+    from litellm.proxy._types import LitellmUserRoles
+    from litellm.proxy.auth.auth_checks import allowed_routes_check
+    
+    jwt_auth = LiteLLM_JWTAuth()  # Use defaults
+    
+    # Check if MCP route is allowed for TEAM role
+    is_allowed = allowed_routes_check(
+        user_role=LitellmUserRoles.TEAM,
+        user_route="/mcp/tools/list",
+        litellm_proxy_roles=jwt_auth,
+    )
+    
+    print(f"Is /mcp/tools/list allowed for TEAM with defaults? {is_allowed}")
+    
+    # MCP routes should be allowed by default for teams
+    assert is_allowed is True, (
+        "MCP routes must be allowed by default for teams for JWT MCP enforcement to work"
+    )
+
+
+@pytest.mark.asyncio
+async def test_e2e_jwt_team_mcp_permissions_enforced(monkeypatch):
+    """
+    End-to-end test verifying that team MCP permissions are properly enforced
+    when using JWT authentication with teams in groups.
+    
+    This test verifies the complete flow:
+    1. JWT token contains team "ABC" in groups field
+    2. Team "ABC" exists with MCP servers ["mcp-server-1", "mcp-server-2"] assigned
+    3. JWT auth properly sets team_id on UserAPIKeyAuth
+    4. MCPRequestHandler.get_allowed_mcp_servers() returns team's MCP servers
+    """
+    from litellm.caching import DualCache
+    from litellm.proxy.utils import ProxyLogging
+    from litellm.router import Router
+
+    # Setup mock router
+    router = Router(model_list=[{"model_name": "gpt-4", "litellm_params": {"model": "gpt-4"}}])
+    import sys
+    import types
+    proxy_server_module = types.ModuleType("proxy_server")
+    proxy_server_module.llm_router = router
+    proxy_server_module.prisma_client = MagicMock()  # Mock prisma client
+    proxy_server_module.user_api_key_cache = DualCache()
+    proxy_server_module.proxy_logging_obj = MagicMock()
+    monkeypatch.setitem(sys.modules, "litellm.proxy.proxy_server", proxy_server_module)
+
+    # Team "ABC" has MCP servers assigned via object_permission
+    team_mcp_servers = ["mcp-server-1", "mcp-server-2"]
+    team_object_permission = LiteLLM_ObjectPermissionTable(
+        object_permission_id="perm-abc-123",
+        mcp_servers=team_mcp_servers,
+        mcp_access_groups=[],
+        vector_stores=[],
+    )
+    
+    team_with_mcp = LiteLLM_TeamTable(
+        team_id="ABC",
+        models=["gpt-4"],
+        object_permission=team_object_permission,
+        object_permission_id="perm-abc-123",
+    )
+
+    async def mock_get_team_object(*args, **kwargs):
+        team_id = kwargs.get("team_id") or (args[0] if args else None)
+        if team_id == "ABC":
+            return team_with_mcp
+        return None
+
+    monkeypatch.setattr(
+        "litellm.proxy.auth.handle_jwt.get_team_object", mock_get_team_object
+    )
+    monkeypatch.setattr(
+        "litellm.proxy.auth.auth_checks.get_team_object", mock_get_team_object
+    )
+
+    # Setup JWT handler with team_ids_jwt_field (groups)
+    jwt_handler = JWTHandler()
+    jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth(
+        team_ids_jwt_field="groups",
+    )
+
+    user_api_key_cache = DualCache()
+    proxy_logging_obj = ProxyLogging(user_api_key_cache=user_api_key_cache)
+
+    # Simulate JWT payload with team in groups
+    jwt_token = {
+        "sub": "user-123",
+        "groups": ["ABC"],
+        "scope": "",
+    }
+
+    # Step 1: Verify JWT auth returns correct team_id
+    with patch.object(jwt_handler, "auth_jwt", new_callable=AsyncMock) as mock_auth_jwt:
+        mock_auth_jwt.return_value = jwt_token
+
+        result = await JWTAuthManager.auth_builder(
+            api_key="test-jwt-token",
+            jwt_handler=jwt_handler,
+            request_data={},
+            general_settings={},
+            route="/mcp/tools/list",
+            prisma_client=None,
+            user_api_key_cache=user_api_key_cache,
+            parent_otel_span=None,
+            proxy_logging_obj=proxy_logging_obj,
+        )
+
+        # Verify team_id is set correctly
+        assert result["team_id"] == "ABC", f"Expected team_id='ABC', got '{result['team_id']}'"
+        assert result["team_object"] is not None, "team_object should not be None"
+        
+        # Step 2: Create UserAPIKeyAuth with the team_id from JWT auth
+        user_api_key_auth = UserAPIKeyAuth(
+            api_key=None,
+            team_id=result["team_id"],
+            user_id=result["user_id"],
+        )
+        
+        # Step 3: Verify MCPRequestHandler returns team's MCP servers
+        # Mock _get_team_object_permission to return our team's object_permission
+        with patch.object(
+            MCPRequestHandler, "_get_team_object_permission"
+        ) as mock_get_team_perm:
+            mock_get_team_perm.return_value = team_object_permission
+            
+            # Mock _get_allowed_mcp_servers_for_key to return empty (no key-level permissions)
+            with patch.object(
+                MCPRequestHandler, "_get_allowed_mcp_servers_for_key"
+            ) as mock_key_servers:
+                mock_key_servers.return_value = []
+                
+                # Mock _get_mcp_servers_from_access_groups to return empty
+                with patch.object(
+                    MCPRequestHandler, "_get_mcp_servers_from_access_groups"
+                ) as mock_access_groups:
+                    mock_access_groups.return_value = []
+                    
+                    allowed_servers = await MCPRequestHandler.get_allowed_mcp_servers(
+                        user_api_key_auth
+                    )
+                    
+                    print(f"Allowed MCP servers: {allowed_servers}")
+                    
+                    # Verify team's MCP servers are returned
+                    assert set(allowed_servers) == set(team_mcp_servers), (
+                        f"Expected team MCP servers {team_mcp_servers}, got {allowed_servers}"
+                    )
+
+
+@pytest.mark.asyncio
+async def test_e2e_jwt_without_team_no_mcp_servers(monkeypatch):
+    """
+    End-to-end test verifying that when JWT has no teams, no MCP servers are returned.
+    
+    This ensures:
+    1. JWT token with no groups returns no team_id
+    2. MCPRequestHandler.get_allowed_mcp_servers() returns empty list
+    """
+    from litellm.caching import DualCache
+    from litellm.proxy.utils import ProxyLogging
+    from litellm.router import Router
+
+    # Setup mock router
+    router = Router(model_list=[])
+    import sys
+    import types
+    proxy_server_module = types.ModuleType("proxy_server")
+    proxy_server_module.llm_router = router
+    monkeypatch.setitem(sys.modules, "litellm.proxy.proxy_server", proxy_server_module)
+
+    async def mock_get_team_object(*args, **kwargs):
+        return None
+
+    monkeypatch.setattr(
+        "litellm.proxy.auth.handle_jwt.get_team_object", mock_get_team_object
+    )
+
+    # Setup JWT handler
+    jwt_handler = JWTHandler()
+    jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth(
+        team_ids_jwt_field="groups",
+    )
+
+    user_api_key_cache = DualCache()
+    proxy_logging_obj = ProxyLogging(user_api_key_cache=user_api_key_cache)
+
+    # JWT payload with empty groups
+    jwt_token = {
+        "sub": "user-123",
+        "groups": [],  # No teams
+        "scope": "",
+    }
+
+    with patch.object(jwt_handler, "auth_jwt", new_callable=AsyncMock) as mock_auth_jwt:
+        mock_auth_jwt.return_value = jwt_token
+
+        result = await JWTAuthManager.auth_builder(
+            api_key="test-jwt-token",
+            jwt_handler=jwt_handler,
+            request_data={},
+            general_settings={},
+            route="/mcp/tools/list",
+            prisma_client=None,
+            user_api_key_cache=user_api_key_cache,
+            parent_otel_span=None,
+            proxy_logging_obj=proxy_logging_obj,
+        )
+
+        # Verify no team_id is set
+        assert result["team_id"] is None, f"Expected team_id=None, got '{result['team_id']}'"
+        
+        # Create UserAPIKeyAuth without team_id
+        user_api_key_auth = UserAPIKeyAuth(
+            api_key=None,
+            team_id=None,
+            user_id=result["user_id"],
+        )
+        
+        # Verify no MCP servers are returned when there's no team
+        allowed_servers = await MCPRequestHandler._get_allowed_mcp_servers_for_team(
+            user_api_key_auth
+        )
+        
+        assert allowed_servers == [], f"Expected empty list, got {allowed_servers}"
+
+
+@pytest.mark.asyncio
+async def test_e2e_jwt_team_mcp_key_intersection(monkeypatch):
+    """
+    End-to-end test verifying MCP permission intersection between key and team.
+    
+    Scenario:
+    - Team has MCP servers: ["server-1", "server-2", "server-3"]
+    - Key has MCP servers: ["server-2", "server-4"]
+    - Result should be intersection: ["server-2"]
+    """
+    from litellm.caching import DualCache
+    from litellm.proxy.utils import ProxyLogging
+    from litellm.router import Router
+
+    # Setup mock router
+    router = Router(model_list=[{"model_name": "gpt-4", "litellm_params": {"model": "gpt-4"}}])
+    import sys
+    import types
+    proxy_server_module = types.ModuleType("proxy_server")
+    proxy_server_module.llm_router = router
+    proxy_server_module.prisma_client = MagicMock()
+    proxy_server_module.user_api_key_cache = DualCache()
+    proxy_server_module.proxy_logging_obj = MagicMock()
+    monkeypatch.setitem(sys.modules, "litellm.proxy.proxy_server", proxy_server_module)
+
+    # Team MCP servers
+    team_mcp_servers = ["server-1", "server-2", "server-3"]
+    team_object_permission = LiteLLM_ObjectPermissionTable(
+        object_permission_id="team-perm",
+        mcp_servers=team_mcp_servers,
+    )
+    
+    team_with_mcp = LiteLLM_TeamTable(
+        team_id="TEAM-X",
+        models=["gpt-4"],
+        object_permission=team_object_permission,
+    )
+
+    # Key MCP servers
+    key_mcp_servers = ["server-2", "server-4"]
+    key_object_permission = LiteLLM_ObjectPermissionTable(
+        object_permission_id="key-perm",
+        mcp_servers=key_mcp_servers,
+    )
+
+    async def mock_get_team_object(*args, **kwargs):
+        team_id = kwargs.get("team_id") or (args[0] if args else None)
+        if team_id == "TEAM-X":
+            return team_with_mcp
+        return None
+
+    monkeypatch.setattr(
+        "litellm.proxy.auth.handle_jwt.get_team_object", mock_get_team_object
+    )
+
+    jwt_handler = JWTHandler()
+    jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth(team_ids_jwt_field="groups")
+
+    user_api_key_cache = DualCache()
+    proxy_logging_obj = ProxyLogging(user_api_key_cache=user_api_key_cache)
+
+    jwt_token = {"sub": "user-123", "groups": ["TEAM-X"], "scope": ""}
+
+    with patch.object(jwt_handler, "auth_jwt", new_callable=AsyncMock) as mock_auth_jwt:
+        mock_auth_jwt.return_value = jwt_token
+
+        result = await JWTAuthManager.auth_builder(
+            api_key="test-jwt-token",
+            jwt_handler=jwt_handler,
+            request_data={},
+            general_settings={},
+            route="/mcp/tools/list",
+            prisma_client=None,
+            user_api_key_cache=user_api_key_cache,
+            parent_otel_span=None,
+            proxy_logging_obj=proxy_logging_obj,
+        )
+
+        assert result["team_id"] == "TEAM-X"
+        
+        user_api_key_auth = UserAPIKeyAuth(
+            api_key=None,
+            team_id=result["team_id"],
+            user_id=result["user_id"],
+            object_permission=key_object_permission,  # Key has its own permissions
+        )
+        
+        # Mock the helper methods to return our test data
+        with patch.object(
+            MCPRequestHandler, "_get_team_object_permission"
+        ) as mock_team_perm:
+            mock_team_perm.return_value = team_object_permission
+            
+            with patch.object(
+                MCPRequestHandler, "_get_key_object_permission"
+            ) as mock_key_perm:
+                mock_key_perm.return_value = key_object_permission
+                
+                with patch.object(
+                    MCPRequestHandler, "_get_mcp_servers_from_access_groups"
+                ) as mock_access_groups:
+                    mock_access_groups.return_value = []
+                    
+                    allowed_servers = await MCPRequestHandler.get_allowed_mcp_servers(
+                        user_api_key_auth
+                    )
+                    
+                    # Should be intersection: only server-2 is in both
+                    expected = ["server-2"]
+                    assert sorted(allowed_servers) == sorted(expected), (
+                        f"Expected intersection {expected}, got {allowed_servers}"
+                    )

tests/test_litellm/proxy/_experimental/mcp_server/test_jwt_mcp_simple.py

@@ -0,0 +1,277 @@
+"""
+Simple test to validate MCP permissions are enforced when calling MCP routes with JWT.
+"""
+
+import pytest
+from unittest.mock import AsyncMock, MagicMock, patch
+
+from litellm.proxy._types import (
+    LiteLLM_JWTAuth,
+    LiteLLM_TeamTable,
+    LiteLLM_ObjectPermissionTable,
+    UserAPIKeyAuth,
+)
+
+
+@pytest.mark.asyncio
+async def test_simple_jwt_mcp_permissions_enforced():
+    """
+    Simple test: Call MCP route with JWT, verify team's MCP servers are returned.
+    
+    Setup:
+    - Team "my-team" has MCP servers: ["github-mcp", "slack-mcp"]
+    - JWT user belongs to "my-team"
+    
+    Expected: Only ["github-mcp", "slack-mcp"] should be allowed
+    """
+    from litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp import (
+        MCPRequestHandler,
+    )
+    
+    # 1. Create a user authenticated via JWT with team_id set
+    user_auth = UserAPIKeyAuth(
+        api_key=None,  # JWT auth doesn't have api_key
+        user_id="jwt-user-123",
+        team_id="my-team",  # This is set by JWT auth when team is in groups
+    )
+    
+    # 2. Team's MCP permissions
+    team_mcp_servers = ["github-mcp", "slack-mcp"]
+    team_object_permission = LiteLLM_ObjectPermissionTable(
+        object_permission_id="perm-123",
+        mcp_servers=team_mcp_servers,
+    )
+    
+    # 3. Mock the team permission lookup
+    with patch.object(
+        MCPRequestHandler, "_get_team_object_permission", new_callable=AsyncMock
+    ) as mock_team_perm:
+        mock_team_perm.return_value = team_object_permission
+        
+        # Mock key permissions (empty - user has no key-level MCP permissions)
+        with patch.object(
+            MCPRequestHandler, "_get_key_object_permission", new_callable=AsyncMock
+        ) as mock_key_perm:
+            mock_key_perm.return_value = None
+            
+            # Mock access groups (empty)
+            with patch.object(
+                MCPRequestHandler, "_get_mcp_servers_from_access_groups", new_callable=AsyncMock
+            ) as mock_access_groups:
+                mock_access_groups.return_value = []
+                
+                # 4. Call get_allowed_mcp_servers - this is what MCP routes use
+                allowed = await MCPRequestHandler.get_allowed_mcp_servers(user_auth)
+                
+                # 5. Verify only team's MCP servers are returned
+                assert sorted(allowed) == sorted(team_mcp_servers), (
+                    f"Expected {team_mcp_servers}, got {allowed}"
+                )
+                
+                # Verify team permission was looked up
+                mock_team_perm.assert_called_once_with(user_auth)
+
+
+@pytest.mark.asyncio
+async def test_simple_jwt_no_team_no_mcp_servers():
+    """
+    Simple test: JWT user with no team should get no MCP servers.
+    """
+    from litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp import (
+        MCPRequestHandler,
+    )
+    
+    # User with no team_id (JWT didn't have teams in groups)
+    user_auth = UserAPIKeyAuth(
+        api_key=None,
+        user_id="jwt-user-no-team",
+        team_id=None,  # No team
+    )
+    
+    # _get_allowed_mcp_servers_for_team returns [] when team_id is None
+    allowed = await MCPRequestHandler._get_allowed_mcp_servers_for_team(user_auth)
+    
+    assert allowed == [], f"Expected [], got {allowed}"
+
+
+@pytest.mark.asyncio
+async def test_simple_jwt_team_id_required_for_mcp_permissions():
+    """
+    Simple test: Verify that team_id must be set for team MCP permissions to work.
+    
+    This is the key insight - if JWT auth doesn't set team_id, 
+    team MCP permissions won't be enforced.
+    """
+    from litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp import (
+        MCPRequestHandler,
+    )
+    
+    # Case 1: team_id is set -> team permissions should be checked
+    user_with_team = UserAPIKeyAuth(
+        api_key=None,
+        user_id="user-1",
+        team_id="team-abc",
+    )
+    
+    team_mcp_servers = ["server-1", "server-2"]
+    team_perm = LiteLLM_ObjectPermissionTable(
+        object_permission_id="perm-1",
+        mcp_servers=team_mcp_servers,
+    )
+    
+    with patch.object(
+        MCPRequestHandler, "_get_team_object_permission", new_callable=AsyncMock
+    ) as mock_perm:
+        mock_perm.return_value = team_perm
+        
+        with patch.object(
+            MCPRequestHandler, "_get_mcp_servers_from_access_groups", new_callable=AsyncMock
+        ) as mock_groups:
+            mock_groups.return_value = []
+            
+            result = await MCPRequestHandler._get_allowed_mcp_servers_for_team(user_with_team)
+            
+            assert sorted(result) == sorted(team_mcp_servers)
+            mock_perm.assert_called_once()  # Permission WAS checked
+    
+    # Case 2: team_id is None -> team permissions NOT checked
+    user_without_team = UserAPIKeyAuth(
+        api_key=None,
+        user_id="user-2",
+        team_id=None,
+    )
+    
+    result = await MCPRequestHandler._get_allowed_mcp_servers_for_team(user_without_team)
+    assert result == []  # No permissions returned
+
+
+@pytest.mark.asyncio 
+async def test_jwt_auth_sets_team_id_for_mcp_route():
+    """
+    Test that JWT auth properly sets team_id when accessing MCP routes.
+    
+    This is the critical test - when user calls /mcp/tools/list with JWT,
+    the team_id from JWT groups must be set on UserAPIKeyAuth.
+    """
+    from litellm.proxy.auth.handle_jwt import JWTAuthManager, JWTHandler
+    from litellm.caching import DualCache
+    from litellm.proxy.utils import ProxyLogging
+    
+    # Setup
+    jwt_handler = JWTHandler()
+    jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth(
+        team_ids_jwt_field="groups",  # Teams come from "groups" field in JWT
+    )
+    
+    # Team exists with models
+    team = LiteLLM_TeamTable(
+        team_id="team-from-jwt",
+        models=["gpt-4"],
+    )
+    
+    user_api_key_cache = DualCache()
+    proxy_logging_obj = ProxyLogging(user_api_key_cache=user_api_key_cache)
+    
+    # Mock JWT token with team in groups
+    jwt_payload = {
+        "sub": "user-123",
+        "groups": ["team-from-jwt"],
+        "scope": "",
+    }
+    
+    with patch.object(jwt_handler, "auth_jwt", new_callable=AsyncMock) as mock_auth:
+        mock_auth.return_value = jwt_payload
+        
+        with patch(
+            "litellm.proxy.auth.handle_jwt.get_team_object", new_callable=AsyncMock
+        ) as mock_get_team:
+            mock_get_team.return_value = team
+            
+            # Simulate calling MCP route
+            result = await JWTAuthManager.auth_builder(
+                api_key="jwt-token",
+                jwt_handler=jwt_handler,
+                request_data={},
+                general_settings={},
+                route="/mcp/tools/list",  # MCP route
+                prisma_client=None,
+                user_api_key_cache=user_api_key_cache,
+                parent_otel_span=None,
+                proxy_logging_obj=proxy_logging_obj,
+            )
+            
+            # THE KEY ASSERTION: team_id must be set
+            assert result["team_id"] == "team-from-jwt", (
+                f"team_id should be 'team-from-jwt' but got '{result['team_id']}'. "
+                "This means JWT auth is not properly setting team_id for MCP routes!"
+            )
+
+
+@pytest.mark.asyncio 
+async def test_mcp_route_without_model_still_returns_team_id():
+    """
+    Test that MCP routes (which don't specify a model) still get team_id assigned.
+    
+    Key insight: MCP routes don't require a model in the request, but the JWT auth
+    flow must still assign a team_id so that team MCP permissions are enforced.
+    
+    The flow is:
+    1. JWT token contains team in "groups" field
+    2. find_team_with_model_access() is called with requested_model=None
+    3. Since `not requested_model` is True, model check passes
+    4. Route check passes because "mcp_routes" is in team_allowed_routes
+    5. team_id is returned and set on UserAPIKeyAuth
+    """
+    from litellm.proxy.auth.handle_jwt import JWTAuthManager, JWTHandler
+    from litellm.caching import DualCache
+    from litellm.proxy.utils import ProxyLogging
+    
+    # Setup
+    jwt_handler = JWTHandler()
+    jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth(
+        team_ids_jwt_field="groups",
+    )
+    
+    # Team exists - note: models is a list (can be empty or have values)
+    # The key is that when no model is requested, model check is skipped
+    team = LiteLLM_TeamTable(
+        team_id="my-team",
+        models=["gpt-4", "gpt-3.5-turbo"],  # Team has models, but MCP request won't specify one
+    )
+    
+    user_api_key_cache = DualCache()
+    proxy_logging_obj = ProxyLogging(user_api_key_cache=user_api_key_cache)
+    
+    # JWT with team in groups
+    jwt_payload = {
+        "sub": "user-abc",
+        "groups": ["my-team"],
+        "scope": "",
+    }
+    
+    with patch.object(jwt_handler, "auth_jwt", new_callable=AsyncMock) as mock_auth:
+        mock_auth.return_value = jwt_payload
+        
+        with patch(
+            "litellm.proxy.auth.handle_jwt.get_team_object", new_callable=AsyncMock
+        ) as mock_get_team:
+            mock_get_team.return_value = team
+            
+            # Call MCP route with NO MODEL in request_data
+            result = await JWTAuthManager.auth_builder(
+                api_key="jwt-token",
+                jwt_handler=jwt_handler,
+                request_data={},  # <-- NO MODEL SPECIFIED
+                general_settings={},
+                route="/mcp/tools/list",  # MCP route
+                prisma_client=None,
+                user_api_key_cache=user_api_key_cache,
+                parent_otel_span=None,
+                proxy_logging_obj=proxy_logging_obj,
+            )
+            
+            # Team ID must still be set even though no model was requested
+            assert result["team_id"] == "my-team", (
+                f"Expected team_id='my-team' but got '{result['team_id']}'. "
+                "MCP routes without model should still get team_id from JWT!"
+            )

tests/test_litellm/proxy/agent_endpoints/test_model_list_helpers.py

@@ -0,0 +1,110 @@
+"""
+Test appending A2A agents to model lists.
+
+Maps to: litellm/proxy/agent_endpoints/model_list_helpers.py
+"""
+import os
+import sys
+
+sys.path.insert(0, os.path.abspath("../../../.."))
+
+from unittest.mock import AsyncMock, Mock, patch
+
+import pytest
+
+from litellm.proxy.agent_endpoints.model_list_helpers import (
+    append_agents_to_model_group,
+    append_agents_to_model_info,
+)
+from litellm.proxy.auth.user_api_key_auth import UserAPIKeyAuth
+from litellm.types.agents import AgentResponse
+from litellm.types.proxy.management_endpoints.model_management_endpoints import (
+    ModelGroupInfoProxy,
+)
+
+
+@pytest.mark.asyncio
+async def test_append_agents_to_model_group():
+    """Test agents are converted to model group format with a2a/ prefix"""
+
+    # Mock agent data
+    mock_agent = AgentResponse(
+        agent_id="test-agent-id",
+        agent_name="my-agent",
+        agent_card_params={"url": "http://example.com"},
+        litellm_params=None,
+    )
+
+    # Mock AgentRequestHandler at its source location
+    mock_get_allowed_agents = AsyncMock(return_value=["test-agent-id"])
+
+    # Mock global_agent_registry
+    mock_registry = Mock()
+    mock_registry.get_agent_by_id = Mock(return_value=mock_agent)
+
+    with patch(
+        "litellm.proxy.agent_endpoints.auth.agent_permission_handler.AgentRequestHandler.get_allowed_agents",
+        mock_get_allowed_agents,
+    ):
+        with patch(
+            "litellm.proxy.agent_endpoints.agent_registry.global_agent_registry",
+            mock_registry,
+        ):
+            model_groups = []
+            user_api_key_dict = Mock(spec=UserAPIKeyAuth)
+
+            result = await append_agents_to_model_group(
+                model_groups=model_groups,
+                user_api_key_dict=user_api_key_dict,
+            )
+
+            # Verify agent was converted with a2a/ prefix
+            assert len(result) == 1
+            assert result[0].model_group == "a2a/my-agent"
+            assert result[0].mode == "chat"
+            assert result[0].providers == ["a2a"]
+
+
+@pytest.mark.asyncio
+async def test_append_agents_to_model_info():
+    """Test agents are converted to model info format with a2a/ prefix"""
+
+    # Mock agent data
+    mock_agent = AgentResponse(
+        agent_id="agent-123",
+        agent_name="test-agent",
+        agent_card_params={"url": "http://example.com"},
+        litellm_params=None,
+        created_by="user-123",
+    )
+
+    # Mock AgentRequestHandler at its source location
+    mock_get_allowed_agents = AsyncMock(return_value=["agent-123"])
+
+    # Mock global_agent_registry
+    mock_registry = Mock()
+    mock_registry.get_agent_by_id = Mock(return_value=mock_agent)
+
+    with patch(
+        "litellm.proxy.agent_endpoints.auth.agent_permission_handler.AgentRequestHandler.get_allowed_agents",
+        mock_get_allowed_agents,
+    ):
+        with patch(
+            "litellm.proxy.agent_endpoints.agent_registry.global_agent_registry",
+            mock_registry,
+        ):
+            models = []
+            user_api_key_dict = Mock(spec=UserAPIKeyAuth)
+
+            result = await append_agents_to_model_info(
+                models=models,
+                user_api_key_dict=user_api_key_dict,
+            )
+
+            # Verify agent was converted with a2a/ prefix
+            assert len(result) == 1
+            assert result[0]["model_name"] == "a2a/test-agent"
+            assert result[0]["litellm_params"]["model"] == "a2a/test-agent"
+            assert result[0]["litellm_params"]["custom_llm_provider"] == "a2a"
+            assert result[0]["model_info"]["id"] == "agent-123"
+            assert result[0]["model_info"]["mode"] == "chat"

tests/test_litellm/proxy/management_endpoints/test_internal_user_endpoints.py

@@ -1133,6 +1133,24 @@ def test_update_internal_user_params_ignores_other_nones():
     assert non_default_values["max_budget"] == 100.0
 
 
+def test_update_internal_user_params_keeps_original_max_budget_when_not_provided():
+    """
+    Test that _update_internal_user_params does not include max_budget 
+    when it's not provided in the request (should keep original value).
+    """
+    # Create test data without max_budget
+    data_json = {"user_id": "test_user", "user_alias": "test_alias"}
+    data = UpdateUserRequest(user_id="test_user", user_alias="test_alias")
+
+    # Call the function
+    non_default_values = _update_internal_user_params(data_json=data_json, data=data)
+
+    # Assertions: max_budget should NOT be in non_default_values
+    assert "max_budget" not in non_default_values
+    assert "user_id" in non_default_values
+    assert "user_alias" in non_default_values
+
+
 def test_generate_request_base_validator():
     """
     Test that GenerateRequestBase validator converts empty string to None for max_budget

tests/test_litellm/proxy/test_route_a2a_models.py

@@ -0,0 +1,105 @@
+"""
+Test A2A model routing in proxy.
+
+Maps to: litellm/proxy/agent_endpoints/a2a_routing.py
+"""
+import os
+import sys
+
+sys.path.insert(0, os.path.abspath("../../.."))
+
+from unittest.mock import AsyncMock, Mock, patch
+
+import pytest
+
+from litellm.proxy.agent_endpoints.a2a_routing import route_a2a_agent_request
+from litellm.proxy.route_llm_request import route_request
+
+
+@pytest.mark.asyncio
+async def test_route_a2a_model_bypasses_router():
+    """Test that a2a/ prefixed models bypass router and go directly to litellm with api_base"""
+
+    # Mock data for chat completion with a2a model
+    data = {
+        "model": "a2a/test-agent",
+        "messages": [{"role": "user", "content": "Hello"}],
+    }
+
+    # Mock router that doesn't have the a2a model
+    mock_router = Mock()
+    mock_router.model_names = ["gpt-4", "gpt-3.5-turbo"]
+    mock_router.deployment_names = []
+    mock_router.has_model_id = Mock(return_value=False)
+    mock_router.model_group_alias = None
+    mock_router.router_general_settings = Mock(pass_through_all_models=False)
+    mock_router.default_deployment = None
+    mock_router.pattern_router = Mock(patterns=[])
+    mock_router.map_team_model = Mock(return_value=None)
+
+    # Mock agent in registry
+    from litellm.types.agents import AgentResponse
+    
+    mock_agent = AgentResponse(
+        agent_id="test-agent-id",
+        agent_name="test-agent",
+        agent_card_params={"url": "http://agent.example.com"},
+        litellm_params=None,
+    )
+    
+    mock_registry = Mock()
+    mock_registry.get_agent_by_name = Mock(return_value=mock_agent)
+
+    # Mock litellm.acompletion to verify it's called
+    mock_acompletion = AsyncMock(return_value={"id": "test-response"})
+
+    with patch("litellm.acompletion", mock_acompletion):
+        with patch(
+            "litellm.proxy.agent_endpoints.a2a_routing.global_agent_registry",
+            mock_registry,
+        ):
+            result = await route_request(
+                data=data,
+                llm_router=mock_router,
+                user_model=None,
+                route_type="acompletion",
+            )
+
+            # Verify litellm.acompletion was called with api_base injected
+            mock_acompletion.assert_called_once()
+            call_kwargs = mock_acompletion.call_args.kwargs
+            assert call_kwargs["model"] == "a2a/test-agent"
+            assert call_kwargs["api_base"] == "http://agent.example.com"
+
+
+@pytest.mark.asyncio  
+async def test_route_non_a2a_model_raises_error_if_not_in_router():
+    """Test that non-a2a models that aren't in router raise an error"""
+
+    # Mock data for chat completion with model not in router
+    data = {
+        "model": "unknown-model",
+        "messages": [{"role": "user", "content": "Hello"}],
+    }
+
+    # Mock router without the model
+    mock_router = Mock()
+    mock_router.model_names = ["gpt-4", "gpt-3.5-turbo"]
+    mock_router.deployment_names = []
+    mock_router.has_model_id = Mock(return_value=False)
+    mock_router.model_group_alias = None
+    mock_router.router_general_settings = Mock(pass_through_all_models=False)
+    mock_router.default_deployment = None
+    mock_router.pattern_router = Mock(patterns=[])
+    mock_router.map_team_model = Mock(return_value=None)
+
+    # Should raise ProxyModelNotFoundError
+    from litellm.proxy.route_llm_request import ProxyModelNotFoundError
+    
+    with pytest.raises(ProxyModelNotFoundError):
+        await route_request(
+            data=data,
+            llm_router=mock_router,
+            user_model=None,
+            route_type="acompletion",
+        )

tests/test_litellm/test_a2a_registry_lookup.py

@@ -0,0 +1,73 @@
+"""
+Test A2A provider registry lookup functionality.
+
+Maps to: litellm/llms/a2a/chat/transformation.py
+"""
+import os
+import sys
+
+sys.path.insert(0, os.path.abspath("../.."))
+
+import pytest
+
+import litellm
+from litellm.llms.a2a.chat.transformation import A2AConfig
+
+
+def test_resolve_agent_config_from_registry_static_method():
+    """Test the static helper method for registry resolution"""
+    
+    # Test 1: No agent name in model
+    api_base, api_key, headers = A2AConfig.resolve_agent_config_from_registry(
+        model="a2a",
+        api_base="http://test.com",
+        api_key=None,
+        headers=None,
+        optional_params={}
+    )
+    assert api_base == "http://test.com"
+    
+    # Test 2: All params provided - should not lookup registry
+    api_base, api_key, headers = A2AConfig.resolve_agent_config_from_registry(
+        model="a2a/test-agent",
+        api_base="http://explicit.com",
+        api_key="explicit-key",
+        headers={"X-Test": "value"},
+        optional_params={}
+    )
+    assert api_base == "http://explicit.com"
+    assert api_key == "explicit-key"
+
+
+def test_a2a_registry_integration():
+    """Test registry lookup in proxy context"""
+    
+    try:
+        from litellm.proxy.agent_endpoints.agent_registry import global_agent_registry
+        from litellm.types.agents import AgentResponse
+
+        # Create test agent
+        test_agent = AgentResponse(
+            agent_id="test-id",
+            agent_name="test-agent",
+            agent_card_params={"url": "http://registry-url.example.com:9999"},
+            litellm_params={"api_key": "registry-key"},
+        )
+        
+        # Register and test
+        original_agents = global_agent_registry.agent_list.copy()
+        global_agent_registry.register_agent(test_agent)
+        
+        try:
+            litellm.completion(
+                model="a2a/test-agent",
+                messages=[{"role": "user", "content": "Hello"}]
+            )
+        except Exception as e:
+            # Should use registry URL (connection error expected)
+            assert "registry-url.example.com" in str(e) or "APIConnectionError" in str(type(e).__name__)
+        finally:
+            global_agent_registry.agent_list = original_agents
+            
+    except ImportError:
+        pytest.skip("Registry not available (not in proxy context)")

ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/ModelsAndEndpointsView.tsx

@@ -400,6 +400,7 @@ const ModelsAndEndpointsView: React.FC<ModelDashboardProps> = ({ premiumUser, te
                     all_models_on_proxy={allModelsOnProxy}
                     getDisplayModelName={getDisplayModelName}
                     setSelectedModelId={setSelectedModelId}
+                    teams={teams}
                   />
                 </TabPanel>
                 <ModelRetrySettingsTab

ui/litellm-dashboard/src/components/AIHub/ModelHubTable.tsx

@@ -483,7 +483,7 @@ const ModelHubTable: React.FC<ModelHubTableProps> = ({ accessToken, publicPage,
       <Modal
         title="Public Model Hub"
         width={600}
-        visible={isPublicPageModalVisible}
+        open={isPublicPageModalVisible}
         footer={null}
         onOk={handleOk}
         onCancel={handleCancel}
@@ -505,7 +505,7 @@ const ModelHubTable: React.FC<ModelHubTableProps> = ({ accessToken, publicPage,
       <Modal
         title={selectedModel?.model_group || "Model Details"}
         width={1000}
-        visible={isModalVisible}
+        open={isModalVisible}
         footer={null}
         onOk={handleOk}
         onCancel={handleCancel}
@@ -656,7 +656,7 @@ print(response.choices[0].message.content)`}
       <Modal
         title={selectedAgent?.name || "Agent Details"}
         width={1000}
-        visible={isAgentModalVisible}
+        open={isAgentModalVisible}
         footer={null}
         onOk={handleOk}
         onCancel={handleCancel}
@@ -795,7 +795,7 @@ print(response.choices[0].message.content)`}
       <Modal
         title={selectedMcpServer?.server_name || "MCP Server Details"}
         width={1000}
-        visible={isMcpModalVisible}
+        open={isMcpModalVisible}
         footer={null}
         onOk={handleOk}
         onCancel={handleCancel}

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

@@ -1018,7 +1018,7 @@ const Teams: React.FC<TeamProps> = ({
           {canCreateOrManageTeams(userRole, userID, organizations) && (
             <Modal
               title="Create Team"
-              visible={isTeamModalVisible}
+              open={isTeamModalVisible}
               width={1000}
               footer={null}
               onOk={handleOk}

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

@@ -323,7 +323,7 @@ const SSOModals: React.FC<SSOModalsProps> = ({
     <>
       <Modal
         title={ssoConfigured ? "Edit SSO Settings" : "Add SSO"}
-        visible={isAddSSOModalVisible}
+        open={isAddSSOModalVisible}
         width={800}
         footer={null}
         onOk={handleAddSSOOk}
@@ -517,7 +517,7 @@ const SSOModals: React.FC<SSOModalsProps> = ({
       {/* Clear Confirmation Modal */}
       <Modal
         title="Confirm Clear SSO Settings"
-        visible={isClearConfirmModalVisible}
+        open={isClearConfirmModalVisible}
         onOk={handleClearSSO}
         onCancel={() => setIsClearConfirmModalVisible(false)}
         okText="Yes, Clear"
@@ -536,7 +536,7 @@ const SSOModals: React.FC<SSOModalsProps> = ({
 
       <Modal
         title="SSO Setup Instructions"
-        visible={isInstructionsModalVisible}
+        open={isInstructionsModalVisible}
         width={800}
         footer={null}
         onOk={handleInstructionsOk}

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

@@ -91,7 +91,7 @@ const AdminPanel: React.FC<AdminPanelProps> = ({
 
   const isLocal = process.env.NODE_ENV === "development";
   if (isLocal != true) {
-    console.log = function () {};
+    console.log = function () { };
   }
 
   const baseUrl = useBaseUrl();
@@ -565,7 +565,7 @@ const AdminPanel: React.FC<AdminPanelProps> = ({
               <Modal
                 title="Manage Allowed IP Addresses"
                 width={800}
-                visible={isAllowedIPModalVisible}
+                open={isAllowedIPModalVisible}
                 onCancel={() => setIsAllowedIPModalVisible(false)}
                 footer={[
                   <Button className="mx-1" key="add" onClick={() => setIsAddIPModalVisible(true)}>
@@ -602,7 +602,7 @@ const AdminPanel: React.FC<AdminPanelProps> = ({
 
               <Modal
                 title="Add Allowed IP Address"
-                visible={isAddIPModalVisible}
+                open={isAddIPModalVisible}
                 onCancel={() => setIsAddIPModalVisible(false)}
                 footer={null}
               >
@@ -618,7 +618,7 @@ const AdminPanel: React.FC<AdminPanelProps> = ({
 
               <Modal
                 title="Confirm Delete"
-                visible={isDeleteIPModalVisible}
+                open={isDeleteIPModalVisible}
                 onCancel={() => setIsDeleteIPModalVisible(false)}
                 onOk={confirmDeleteIP}
                 footer={[
@@ -636,7 +636,7 @@ const AdminPanel: React.FC<AdminPanelProps> = ({
               {/* UI Access Control Modal */}
               <Modal
                 title="UI Access Control Settings"
-                visible={isUIAccessControlModalVisible}
+                open={isUIAccessControlModalVisible}
                 width={600}
                 footer={null}
                 onOk={handleUIAccessControlOk}

ui/litellm-dashboard/src/components/budgets/budget_modal.tsx

@@ -43,7 +43,7 @@ const BudgetModal: React.FC<BudgetModalProps> = ({ isModalVisible, accessToken,
   return (
     <Modal
       title="Create Budget"
-      visible={isModalVisible}
+      open={isModalVisible}
       width={800}
       footer={null}
       onOk={handleOk}

ui/litellm-dashboard/src/components/budgets/edit_budget_modal.tsx

@@ -59,7 +59,7 @@ const EditBudgetModal: React.FC<BudgetModalProps> = ({
   return (
     <Modal
       title="Edit Budget"
-      visible={isModalVisible}
+      open={isModalVisible}
       width={800}
       footer={null}
       onOk={handleOk}

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

@@ -370,11 +370,11 @@ const BulkCreateUsersButton: React.FC<BulkCreateUsersProps> = ({
                 current.map((u, i) =>
                   i === index
                     ? {
-                        ...u,
-                        status: "success",
-                        key: response.key || response.user_id,
-                        invitation_link: invitationUrl,
-                      }
+                      ...u,
+                      status: "success",
+                      key: response.key || response.user_id,
+                      invitation_link: invitationUrl,
+                    }
                     : u,
                 ),
               );
@@ -386,11 +386,11 @@ const BulkCreateUsersButton: React.FC<BulkCreateUsersProps> = ({
                 current.map((u, i) =>
                   i === index
                     ? {
-                        ...u,
-                        status: "success",
-                        key: response.key || response.user_id,
-                        invitation_link: invitationUrl,
-                      }
+                      ...u,
+                      status: "success",
+                      key: response.key || response.user_id,
+                      invitation_link: invitationUrl,
+                    }
                     : u,
                 ),
               );
@@ -401,11 +401,11 @@ const BulkCreateUsersButton: React.FC<BulkCreateUsersProps> = ({
               current.map((u, i) =>
                 i === index
                   ? {
-                      ...u,
-                      status: "success",
-                      key: response.key || response.user_id,
-                      error: "User created but failed to generate invitation link",
-                    }
+                    ...u,
+                    status: "success",
+                    key: response.key || response.user_id,
+                    error: "User created but failed to generate invitation link",
+                  }
                   : u,
               ),
             );
@@ -546,7 +546,7 @@ const BulkCreateUsersButton: React.FC<BulkCreateUsersProps> = ({
 
       <Modal
         title="Bulk Invite Users"
-        visible={isModalVisible}
+        open={isModalVisible}
         width={800}
         onCancel={() => setIsModalVisible(false)}
         bodyStyle={{ maxHeight: "70vh", overflow: "auto" }}

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

@@ -226,7 +226,7 @@ const CloudZeroExportModal: React.FC<CloudZeroExportModalProps> = ({ isOpen, onC
   ];
 
   return (
-    <Modal title="Export Data" open={isOpen} onCancel={handleModalClose} footer={null} width={600} destroyOnClose>
+    <Modal title="Export Data" open={isOpen} onCancel={handleModalClose} footer={null} width={600} destroyOnHidden>
       <div className="space-y-4">
         {/* Export Type Selection */}
         <div>

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

@@ -228,7 +228,7 @@ const Createuser: React.FC<CreateuserProps> = ({
       <BulkCreateUsers accessToken={accessToken} teams={teams} possibleUIRoles={possibleUIRoles} />
       <Modal
         title="Invite User"
-        visible={isModalVisible}
+        open={isModalVisible}
         width={800}
         footer={null}
         onOk={handleOk}

ui/litellm-dashboard/src/components/edit_auto_router/edit_auto_router_modal.tsx

@@ -160,7 +160,7 @@ const EditAutoRouterModal: React.FC<EditAutoRouterModalProps> = ({
         </Button>,
       ]}
       width={1000}
-      destroyOnClose
+      destroyOnHidden
     >
       <div className="space-y-6">
         <Text className="text-gray-600">

ui/litellm-dashboard/src/components/edit_model/edit_model_modal.tsx

@@ -53,8 +53,8 @@ export const handleEditModelSubmit = async (
     model_info:
       model_info_model_id !== undefined
         ? {
-            id: model_info_model_id,
-          }
+          id: model_info_model_id,
+        }
         : undefined,
   };
 
@@ -119,7 +119,7 @@ const EditModelModal: React.FC<EditModelModalProps> = ({ visible, onCancel, mode
   return (
     <Modal
       title={"Edit '" + model_name + "' LiteLLM Params"}
-      visible={visible}
+      open={visible}
       width={800}
       footer={null}
       onOk={handleOk}

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

@@ -39,7 +39,7 @@ const EditUserModal: React.FC<EditUserModalProps> = ({ visible, possibleUIRoles,
   }
 
   return (
-    <Modal visible={visible} onCancel={handleCancel} footer={null} title={"Edit User " + user.user_id} width={1000}>
+    <Modal open={visible} onCancel={handleCancel} footer={null} title={"Edit User " + user.user_id} width={1000}>
       <Form
         form={form}
         onFinish={handleEditSubmit}

ui/litellm-dashboard/src/components/model_add/CredentialDeleteModal.tsx

@@ -43,7 +43,7 @@ const CredentialDeleteModal: React.FC<CredentialDeleteModalProps> = ({
       footer={null}
       onCancel={handleCancel}
       closable={true}
-      destroyOnClose={true}
+      destroyOnHidden={true}
       maskClosable={false}
     >
       <div className="mt-4">

ui/litellm-dashboard/src/components/model_add/reuse_credentials.tsx

@@ -32,7 +32,7 @@ const ReuseCredentialsModal: React.FC<ReuseCredentialsModalProps> = ({
   return (
     <Modal
       title="Reuse Credentials"
-      visible={isVisible}
+      open={isVisible}
       onCancel={() => {
         onCancel();
         form.resetFields();

ui/litellm-dashboard/src/components/model_dashboard/HealthCheckComponent.tsx

@@ -7,6 +7,7 @@ import { healthCheckColumns } from "./health_check_columns";
 import { errorPatterns } from "@/utils/errorPatterns";
 import { individualModelHealthCheckCall, latestHealthChecksCall } from "../networking";
 import { Table as TableInstance } from "@tanstack/react-table";
+import { Team } from "../key_team_helpers/key_list";
 
 interface HealthStatus {
   status: string;
@@ -24,6 +25,7 @@ interface HealthCheckComponentProps {
   all_models_on_proxy: string[];
   getDisplayModelName: (model: any) => string;
   setSelectedModelId?: (modelId: string) => void;
+  teams?: Team[] | null;
 }
 
 const HealthCheckComponent: React.FC<HealthCheckComponentProps> = ({
@@ -32,6 +34,7 @@ const HealthCheckComponent: React.FC<HealthCheckComponentProps> = ({
   all_models_on_proxy,
   getDisplayModelName,
   setSelectedModelId,
+  teams,
 }) => {
   const [modelHealthStatuses, setModelHealthStatuses] = useState<{ [key: string]: HealthStatus }>({});
   const [selectedModelsForHealth, setSelectedModelsForHealth] = useState<string[]>([]);
@@ -574,6 +577,7 @@ const HealthCheckComponent: React.FC<HealthCheckComponentProps> = ({
             showErrorModal,
             showSuccessModal,
             setSelectedModelId,
+            teams,
           )}
           data={modelData.data.map((model: any) => {
             const modelName = model.model_name;

ui/litellm-dashboard/src/components/model_dashboard/health_check_columns.tsx

@@ -2,6 +2,7 @@ import { ColumnDef } from "@tanstack/react-table";
 import { Tooltip, Checkbox } from "antd";
 import { Text } from "@tremor/react";
 import { InformationCircleIcon, PlayIcon, RefreshIcon } from "@heroicons/react/outline";
+import { Team } from "@/components/key_team_helpers/key_list";
 
 interface HealthCheckData {
   model_name: string;
@@ -42,6 +43,7 @@ export const healthCheckColumns = (
   showErrorModal?: (modelName: string, cleanedError: string, fullError: string) => void,
   showSuccessModal?: (modelName: string, response: any) => void,
   setSelectedModelId?: (modelId: string) => void,
+  teams?: Team[] | null,
 ): ColumnDef<HealthCheckData>[] => [
   {
     header: () => (
@@ -100,6 +102,31 @@ export const healthCheckColumns = (
       );
     },
   },
+  {
+    header: "Team Alias",
+    accessorKey: "model_info.team_id",
+    enableSorting: true,
+    sortingFn: "alphanumeric",
+    cell: ({ row }) => {
+      const model = row.original;
+      const teamId = model.model_info?.team_id;
+
+      if (!teamId) {
+        return <span className="text-gray-400 text-sm">-</span>;
+      }
+
+      const team = teams?.find((t) => t.team_id === teamId);
+      const teamAlias = team?.team_alias || teamId;
+
+      return (
+        <div className="text-sm">
+          <Tooltip title={teamAlias}>
+            <div className="truncate max-w-[150px]">{teamAlias}</div>
+          </Tooltip>
+        </div>
+      );
+    },
+  },
   {
     header: "Health Status",
     accessorKey: "health_status",

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

@@ -63,7 +63,7 @@ export default function OnboardingModal({
   return (
     <Modal
       title={modalType === "invitation" ? "Invitation Link" : "Reset Password Link"}
-      visible={isInvitationLinkModalVisible}
+      open={isInvitationLinkModalVisible}
       width={800}
       footer={null}
       onOk={handleInvitationOk}

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

@@ -1342,7 +1342,7 @@ const CreateKey: React.FC<CreateKeyProps> = ({ team, teams, data, addKey }) => {
       {isCreateUserModalVisible && (
         <Modal
           title="Create New User"
-          visible={isCreateUserModalVisible}
+          open={isCreateUserModalVisible}
           onCancel={() => setIsCreateUserModalVisible(false)}
           footer={null}
           width={800}
@@ -1359,7 +1359,7 @@ const CreateKey: React.FC<CreateKeyProps> = ({ team, teams, data, addKey }) => {
       )}
 
       {apiKey && (
-        <Modal visible={isModalVisible} onOk={handleOk} onCancel={handleCancel} footer={null}>
+        <Modal open={isModalVisible} onOk={handleOk} onCancel={handleCancel} footer={null}>
           <Grid numItems={1} className="gap-2 w-full">
             <Title>Save your Key</Title>
             <Col numColSpan={1}>

ui/litellm-dashboard/src/components/organization/add_org_admin.tsx

@@ -45,7 +45,7 @@ const AddOrgAdmin: FC<AddOrgAdminProps> = ({ userRole, userID, selectedOrganizat
 
       <Modal
         title="Add member"
-        visible={isAddMemberModalVisible}
+        open={isAddMemberModalVisible}
         width={800}
         footer={null}
         onOk={handleMemberOk}

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

@@ -60,7 +60,7 @@ const RequestAccess: React.FC<RequestAccessProps> = ({ userModels, accessToken,
       </Button>
       <Modal
         title="Request Access"
-        visible={isModalVisible}
+        open={isModalVisible}
         width={800}
         footer={null}
         onOk={handleOk}

ui/litellm-dashboard/src/components/tag_management/components/CreateTagModal.tsx

@@ -36,7 +36,7 @@ const CreateTagModal: React.FC<CreateTagModalProps> = ({ visible, onCancel, onSu
   };
 
   return (
-    <Modal title="Create New Tag" visible={visible} width={800} footer={null} onCancel={handleCancel}>
+    <Modal title="Create New Tag" open={visible} width={800} footer={null} onCancel={handleCancel}>
       <Form form={form} onFinish={handleFinish} labelCol={{ span: 8 }} wrapperCol={{ span: 16 }} labelAlign="left">
         <Form.Item label="Tag Name" name="tag_name" rules={[{ required: true, message: "Please input a tag name" }]}>
           <TextInput />

ui/litellm-dashboard/src/components/templates/key_edit_view.test.tsx

@@ -1,24 +1,61 @@
-import { fireEvent, waitFor } from "@testing-library/react";
-import { describe, expect, it, vi } from "vitest";
+import { screen, waitFor } from "@testing-library/react";
+import userEvent from "@testing-library/user-event";
+import { beforeEach, describe, expect, it, vi } from "vitest";
 import { renderWithProviders } from "../../../tests/test-utils";
 import { KeyResponse } from "../key_team_helpers/key_list";
 import { KeyEditView } from "./key_edit_view";
 
-// Mock window.matchMedia
-Object.defineProperty(window, "matchMedia", {
-  writable: true,
-  value: vi.fn().mockImplementation((query) => ({
-    matches: false,
-    media: query,
-    onchange: null,
-    addListener: vi.fn(),
-    removeListener: vi.fn(),
-    addEventListener: vi.fn(),
-    removeEventListener: vi.fn(),
-    dispatchEvent: vi.fn(),
-  })),
+vi.mock("../networking", async () => {
+  const actual = await vi.importActual("../networking");
+  return {
+    ...actual,
+    getPromptsList: vi.fn().mockResolvedValue({
+      prompts: [{ prompt_id: "prompt-1" }, { prompt_id: "prompt-2" }],
+    }),
+    modelAvailableCall: vi.fn().mockResolvedValue({
+      data: [{ id: "gpt-4" }, { id: "gpt-3.5-turbo" }],
+    }),
+    tagListCall: vi.fn().mockResolvedValue({
+      tag1: { name: "tag1", description: "Test tag 1" },
+      tag2: { name: "tag2", description: "Test tag 2" },
+    }),
+    getGuardrailsList: vi.fn().mockResolvedValue({
+      guardrails: [{ guardrail_name: "guardrail-1" }],
+    }),
+    getPoliciesList: vi.fn().mockResolvedValue({
+      policies: [{ policy_name: "policy-1" }],
+    }),
+    getPassThroughEndpointsCall: vi.fn().mockResolvedValue({
+      endpoints: [],
+    }),
+    vectorStoreListCall: vi.fn().mockResolvedValue({
+      data: [],
+    }),
+    mcpToolsCall: vi.fn().mockResolvedValue({
+      data: [],
+    }),
+    agentListCall: vi.fn().mockResolvedValue({
+      data: [],
+    }),
+    fetchMCPServers: vi.fn().mockResolvedValue([]),
+    fetchMCPAccessGroups: vi.fn().mockResolvedValue([]),
+    listMCPTools: vi.fn().mockResolvedValue({
+      tools: [],
+      error: null,
+      message: null,
+      stack_trace: null,
+    }),
+    getAgentsList: vi.fn().mockResolvedValue({
+      agents: [],
+    }),
+    getAgentAccessGroups: vi.fn().mockResolvedValue([]),
+  };
 });
 
+vi.mock("../organisms/create_key_button", () => ({
+  fetchTeamModels: vi.fn().mockResolvedValue(["team-model-1", "team-model-2"]),
+}));
+
 describe("KeyEditView", () => {
   const MOCK_KEY_DATA: KeyResponse = {
     token: "test-token-123",
@@ -93,8 +130,8 @@ describe("KeyEditView", () => {
     const { getByText } = renderWithProviders(
       <KeyEditView
         keyData={MOCK_KEY_DATA}
-        onCancel={() => {}}
-        onSubmit={async () => {}}
+        onCancel={() => { }}
+        onSubmit={async () => { }}
         accessToken={""}
         userID={""}
         userRole={""}
@@ -111,8 +148,8 @@ describe("KeyEditView", () => {
     const { getByText } = renderWithProviders(
       <KeyEditView
         keyData={MOCK_KEY_DATA}
-        onCancel={() => {}}
-        onSubmit={async () => {}}
+        onCancel={() => { }}
+        onSubmit={async () => { }}
         accessToken={""}
         userID={""}
         userRole={""}
@@ -129,8 +166,8 @@ describe("KeyEditView", () => {
     const { getByLabelText } = renderWithProviders(
       <KeyEditView
         keyData={MOCK_KEY_DATA}
-        onCancel={() => {}}
-        onSubmit={async () => {}}
+        onCancel={() => { }}
+        onSubmit={async () => { }}
         accessToken={""}
         userID={""}
         userRole={""}
@@ -144,13 +181,17 @@ describe("KeyEditView", () => {
     });
   });
 
+  beforeEach(() => {
+    vi.clearAllMocks();
+  });
+
   it("should call onCancel when cancel button is clicked", async () => {
     const onCancelMock = vi.fn();
-    const { getByText } = renderWithProviders(
+    renderWithProviders(
       <KeyEditView
         keyData={MOCK_KEY_DATA}
         onCancel={onCancelMock}
-        onSubmit={async () => {}}
+        onSubmit={async () => { }}
         accessToken={""}
         userID={""}
         userRole={""}
@@ -159,12 +200,272 @@ describe("KeyEditView", () => {
     );
 
     await waitFor(() => {
-      expect(getByText("Cancel")).toBeInTheDocument();
+      expect(screen.getByRole("button", { name: /cancel/i })).toBeInTheDocument();
     });
 
-    const cancelButton = getByText("Cancel");
-    fireEvent.click(cancelButton);
+    const cancelButton = screen.getByRole("button", { name: /cancel/i });
+    await userEvent.click(cancelButton);
 
     expect(onCancelMock).toHaveBeenCalledTimes(1);
   });
+
+  it("should display key alias input field", async () => {
+    renderWithProviders(
+      <KeyEditView
+        keyData={MOCK_KEY_DATA}
+        onCancel={() => { }}
+        onSubmit={async () => { }}
+        accessToken={""}
+        userID={""}
+        userRole={""}
+        premiumUser={false}
+      />,
+    );
+
+    await waitFor(() => {
+      expect(screen.getByLabelText("Key Alias")).toBeInTheDocument();
+    });
+  });
+
+  it("should display models select field", async () => {
+    renderWithProviders(
+      <KeyEditView
+        keyData={MOCK_KEY_DATA}
+        onCancel={() => { }}
+        onSubmit={async () => { }}
+        accessToken={""}
+        userID={""}
+        userRole={""}
+        premiumUser={false}
+      />,
+    );
+
+    await waitFor(() => {
+      expect(screen.getByText("Models")).toBeInTheDocument();
+    });
+  });
+
+  it("should display max budget input field", async () => {
+    renderWithProviders(
+      <KeyEditView
+        keyData={MOCK_KEY_DATA}
+        onCancel={() => { }}
+        onSubmit={async () => { }}
+        accessToken={""}
+        userID={""}
+        userRole={""}
+        premiumUser={false}
+      />,
+    );
+
+    await waitFor(() => {
+      expect(screen.getByLabelText("Max Budget (USD)")).toBeInTheDocument();
+    });
+  });
+
+  it("should display allowed routes input field", async () => {
+    renderWithProviders(
+      <KeyEditView
+        keyData={MOCK_KEY_DATA}
+        onCancel={() => { }}
+        onSubmit={async () => { }}
+        accessToken={""}
+        userID={""}
+        userRole={""}
+        premiumUser={false}
+      />,
+    );
+
+    await waitFor(() => {
+      expect(screen.getByLabelText(/allowed routes/i)).toBeInTheDocument();
+    });
+  });
+
+  it("should call onSubmit with form values when form is submitted", async () => {
+    const onSubmitMock = vi.fn().mockResolvedValue(undefined);
+    renderWithProviders(
+      <KeyEditView
+        keyData={MOCK_KEY_DATA}
+        onCancel={() => { }}
+        onSubmit={onSubmitMock}
+        accessToken={"test-token"}
+        userID={"test-user"}
+        userRole={"admin"}
+        premiumUser={false}
+      />,
+    );
+
+    await waitFor(() => {
+      expect(screen.getByRole("button", { name: /save changes/i })).toBeInTheDocument();
+    });
+
+    const submitButton = screen.getByRole("button", { name: /save changes/i });
+    await userEvent.click(submitButton);
+
+    await waitFor(() => {
+      expect(onSubmitMock).toHaveBeenCalled();
+    });
+  });
+
+  it("should disable models field when management routes are selected", async () => {
+    const keyDataWithManagementRoutes = {
+      ...MOCK_KEY_DATA,
+      allowed_routes: ["management_routes"],
+    };
+
+    renderWithProviders(
+      <KeyEditView
+        keyData={keyDataWithManagementRoutes}
+        onCancel={() => { }}
+        onSubmit={async () => { }}
+        accessToken={""}
+        userID={""}
+        userRole={""}
+        premiumUser={false}
+      />,
+    );
+
+    await waitFor(() => {
+      expect(screen.getByText("Models field is disabled for this key type")).toBeInTheDocument();
+    });
+  });
+
+  it("should disable models field when info routes are selected", async () => {
+    const keyDataWithInfoRoutes = {
+      ...MOCK_KEY_DATA,
+      allowed_routes: ["info_routes"],
+    };
+
+    renderWithProviders(
+      <KeyEditView
+        keyData={keyDataWithInfoRoutes}
+        onCancel={() => { }}
+        onSubmit={async () => { }}
+        accessToken={""}
+        userID={""}
+        userRole={""}
+        premiumUser={false}
+      />,
+    );
+
+    await waitFor(() => {
+      expect(screen.getByText("Models field is disabled for this key type")).toBeInTheDocument();
+    });
+  });
+
+  it("should disable guardrails selector when user is not premium", async () => {
+    renderWithProviders(
+      <KeyEditView
+        keyData={MOCK_KEY_DATA}
+        onCancel={() => { }}
+        onSubmit={async () => { }}
+        accessToken={"test-token"}
+        userID={""}
+        userRole={""}
+        premiumUser={false}
+      />,
+    );
+
+    await waitFor(() => {
+      expect(screen.getByText("Guardrails")).toBeInTheDocument();
+    });
+  });
+
+  it("should parse comma-separated allowed routes on submit", async () => {
+    const onSubmitMock = vi.fn().mockResolvedValue(undefined);
+    renderWithProviders(
+      <KeyEditView
+        keyData={MOCK_KEY_DATA}
+        onCancel={() => { }}
+        onSubmit={onSubmitMock}
+        accessToken={"test-token"}
+        userID={"test-user"}
+        userRole={"admin"}
+        premiumUser={false}
+      />,
+    );
+
+    await waitFor(() => {
+      expect(screen.getByLabelText(/allowed routes/i)).toBeInTheDocument();
+    });
+
+    const allowedRoutesInput = screen.getByLabelText(/allowed routes/i);
+    await userEvent.clear(allowedRoutesInput);
+    await userEvent.type(allowedRoutesInput, "route1, route2, route3");
+
+    const submitButton = screen.getByRole("button", { name: /save changes/i });
+    await userEvent.click(submitButton);
+
+    await waitFor(() => {
+      expect(onSubmitMock).toHaveBeenCalled();
+      const callArgs = onSubmitMock.mock.calls[0][0];
+      expect(Array.isArray(callArgs.allowed_routes)).toBe(true);
+      expect(callArgs.allowed_routes).toEqual(["route1", "route2", "route3"]);
+    });
+  });
+
+  it("should handle empty allowed routes string on submit", async () => {
+    const onSubmitMock = vi.fn().mockResolvedValue(undefined);
+    renderWithProviders(
+      <KeyEditView
+        keyData={MOCK_KEY_DATA}
+        onCancel={() => { }}
+        onSubmit={onSubmitMock}
+        accessToken={"test-token"}
+        userID={"test-user"}
+        userRole={"admin"}
+        premiumUser={false}
+      />,
+    );
+
+    await waitFor(() => {
+      expect(screen.getByLabelText(/allowed routes/i)).toBeInTheDocument();
+    });
+
+    const allowedRoutesInput = screen.getByLabelText(/allowed routes/i);
+    await userEvent.clear(allowedRoutesInput);
+
+    const submitButton = screen.getByRole("button", { name: /save changes/i });
+    await userEvent.click(submitButton);
+
+    await waitFor(() => {
+      expect(onSubmitMock).toHaveBeenCalled();
+      const callArgs = onSubmitMock.mock.calls[0][0];
+      expect(callArgs.allowed_routes).toEqual([]);
+    });
+  });
+
+
+  it("should disable cancel button during submission", async () => {
+    const onSubmitMock = vi.fn(
+      () =>
+        new Promise<void>((resolve) => {
+          setTimeout(resolve, 100);
+        }),
+    );
+
+    renderWithProviders(
+      <KeyEditView
+        keyData={MOCK_KEY_DATA}
+        onCancel={() => { }}
+        onSubmit={onSubmitMock}
+        accessToken={"test-token"}
+        userID={"test-user"}
+        userRole={"admin"}
+        premiumUser={false}
+      />,
+    );
+
+    await waitFor(() => {
+      expect(screen.getByRole("button", { name: /cancel/i })).toBeInTheDocument();
+    });
+
+    const submitButton = screen.getByRole("button", { name: /save changes/i });
+    await userEvent.click(submitButton);
+
+    await waitFor(() => {
+      const cancelButton = screen.getByRole("button", { name: /cancel/i });
+      expect(cancelButton).toBeDisabled();
+    });
+  });
 });

ui/litellm-dashboard/src/components/templates/key_edit_view.tsx

@@ -35,7 +35,6 @@ interface KeyEditViewProps {
 // Add this helper function
 const getAvailableModelsForKey = (keyData: KeyResponse, teams: any[] | null): string[] => {
   // If no teams data is available, return empty array
-  console.log("getAvailableModelsForKey:", teams);
   if (!teams || !keyData.team_id) {
     return [];
   }
@@ -172,7 +171,9 @@ export function KeyEditView({
       : [],
     auto_rotate: keyData.auto_rotate || false,
     ...(keyData.rotation_interval && { rotation_interval: keyData.rotation_interval }),
-    allowed_routes: keyData.allowed_routes,
+    allowed_routes: Array.isArray(keyData.allowed_routes) && keyData.allowed_routes.length > 0
+      ? keyData.allowed_routes.join(", ")
+      : "",
   };
 
   useEffect(() => {
@@ -197,7 +198,9 @@ export function KeyEditView({
         : [],
       auto_rotate: keyData.auto_rotate || false,
       ...(keyData.rotation_interval && { rotation_interval: keyData.rotation_interval }),
-      allowed_routes: keyData.allowed_routes,
+      allowed_routes: Array.isArray(keyData.allowed_routes) && keyData.allowed_routes.length > 0
+        ? keyData.allowed_routes.join(", ")
+        : "",
     });
   }, [keyData, form]);
 
@@ -226,11 +229,24 @@ export function KeyEditView({
     fetchTags();
   }, [accessToken]);
 
-  console.log("premiumUser:", premiumUser);
-
   const handleSubmit = async (values: any) => {
     try {
       setIsKeySaving(true);
+
+      // Parse allowed_routes from comma-separated string to array
+      if (typeof values.allowed_routes === "string") {
+        const trimmedInput = values.allowed_routes.trim();
+        if (trimmedInput === "") {
+          values.allowed_routes = [];
+        } else {
+          values.allowed_routes = trimmedInput
+            .split(",")
+            .map((route: string) => route.trim())
+            .filter((route: string) => route.length > 0);
+        }
+      }
+      // If it's already an array (shouldn't happen, but handle it), keep as is
+
       await onSubmit(values);
     } finally {
       setIsKeySaving(false);
@@ -251,7 +267,11 @@ export function KeyEditView({
           }
         >
           {({ getFieldValue, setFieldValue }) => {
-            const allowedRoutes = getFieldValue("allowed_routes") || [];
+            const allowedRoutesValue = getFieldValue("allowed_routes") || "";
+            // Convert string to array for checking
+            const allowedRoutes = typeof allowedRoutesValue === "string" && allowedRoutesValue.trim() !== ""
+              ? allowedRoutesValue.split(",").map((r: string) => r.trim()).filter((r: string) => r.length > 0)
+              : [];
             const isDisabled = allowedRoutes.includes("management_routes") || allowedRoutes.includes("info_routes");
             const models = getFieldValue("models") || [];
 
@@ -290,7 +310,11 @@ export function KeyEditView({
           shouldUpdate={(prevValues, currentValues) => prevValues.allowed_routes !== currentValues.allowed_routes}
         >
           {({ getFieldValue, setFieldValue }) => {
-            const allowedRoutes = getFieldValue("allowed_routes");
+            const allowedRoutesValue = getFieldValue("allowed_routes") || "";
+            // Convert string to array for getKeyTypeFromRoutes
+            const allowedRoutes = typeof allowedRoutesValue === "string" && allowedRoutesValue.trim() !== ""
+              ? allowedRoutesValue.split(",").map((r: string) => r.trim()).filter((r: string) => r.length > 0)
+              : [];
             const keyTypeValue = getKeyTypeFromRoutes(allowedRoutes);
 
             return (
@@ -302,13 +326,13 @@ export function KeyEditView({
                 onChange={(value) => {
                   switch (value) {
                     case "default":
-                      setFieldValue("allowed_routes", []);
+                      setFieldValue("allowed_routes", "");
                       break;
                     case "llm_api":
-                      setFieldValue("allowed_routes", ["llm_api_routes"]);
+                      setFieldValue("allowed_routes", "llm_api_routes");
                       break;
                     case "management":
-                      setFieldValue("allowed_routes", ["management_routes"]);
+                      setFieldValue("allowed_routes", "management_routes");
                       setFieldValue("models", []);
                       break;
                   }
@@ -344,6 +368,22 @@ export function KeyEditView({
         </Form.Item>
       </Form.Item>
 
+      <Form.Item
+        label={
+          <span>
+            Allowed Routes{" "}
+            <Tooltip title="List of allowed routes for the key (comma-separated). Can be specific routes (e.g., '/chat/completions') or route patterns (e.g., 'llm_api_routes', 'management_routes', '/keys/*'). Leave empty to allow all routes.">
+              <InfoCircleOutlined style={{ marginLeft: "4px" }} />
+            </Tooltip>
+          </span>
+        }
+        name="allowed_routes"
+      >
+        <Input
+          placeholder="Enter allowed routes (comma-separated). Special values: llm_api_routes, management_routes. Examples: llm_api_routes, /chat/completions, /keys/*. Leave empty to allow all routes"
+        />
+      </Form.Item>
+
       <Form.Item label="Max Budget (USD)" name="max_budget">
         <NumericalInput step={0.01} style={{ width: "100%" }} placeholder="Enter a numerical value" />
       </Form.Item>
@@ -473,7 +513,7 @@ export function KeyEditView({
               !premiumUser
                 ? "Premium feature - Upgrade to set allowed pass through routes by key"
                 : Array.isArray(keyData.metadata?.allowed_passthrough_routes) &&
-                    keyData.metadata.allowed_passthrough_routes.length > 0
+                  keyData.metadata.allowed_passthrough_routes.length > 0
                   ? `Current: ${keyData.metadata.allowed_passthrough_routes.join(", ")}`
                   : "Select or enter allowed pass through routes"
             }
@@ -590,11 +630,6 @@ export function KeyEditView({
         <Input />
       </Form.Item>
 
-      {/* Hidden form field for allowed_routes */}
-      <Form.Item name="allowed_routes" hidden>
-        <Input />
-      </Form.Item>
-
       {/* Hidden form field for disabled callbacks */}
       <Form.Item name="disabled_callbacks" hidden>
         <Input />

ui/litellm-dashboard/src/components/templates/key_info_view.test.tsx

@@ -1,6 +1,7 @@
-import useTeams from "@/app/(dashboard)/hooks/useTeams";
 import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized";
+import useTeams from "@/app/(dashboard)/hooks/useTeams";
 import { render, screen, waitFor } from "@testing-library/react";
+import userEvent from "@testing-library/user-event";
 import { beforeEach, describe, expect, it, vi } from "vitest";
 import { KeyResponse, Team } from "../key_team_helpers/key_list";
 import KeyInfoView from "./key_info_view";
@@ -13,6 +14,21 @@ vi.mock("@/app/(dashboard)/hooks/useAuthorized", () => ({
   default: vi.fn(),
 }));
 
+vi.mock("../networking", () => ({
+  keyDeleteCall: vi.fn().mockResolvedValue({}),
+  keyUpdateCall: vi.fn().mockResolvedValue({}),
+  getPolicyInfoWithGuardrails: vi.fn().mockResolvedValue({
+    resolved_guardrails: ["guardrail-1", "guardrail-2"],
+  }),
+}));
+
+vi.mock("@/utils/dataUtils", () => ({
+  copyToClipboard: vi.fn().mockResolvedValue(true),
+  formatNumberWithCommas: vi.fn((value: number, decimals?: number) => {
+    return value.toFixed(decimals ?? 2);
+  }),
+}));
+
 describe("KeyInfoView", () => {
   beforeEach(() => {
     vi.mocked(useTeams).mockReturnValue({
@@ -105,34 +121,34 @@ describe("KeyInfoView", () => {
   it("should render tags", async () => {
     vi.mocked(useAuthorized).mockReturnValue(baseUseAuthorizedMock);
 
-    const { getByText } = render(
+    render(
       <KeyInfoView
         keyData={MOCK_KEY_DATA}
-        onClose={() => {}}
+        onClose={() => { }}
         keyId={"test-key-id"}
-        onKeyDataUpdate={() => {}}
+        onKeyDataUpdate={() => { }}
         teams={[]}
       />,
     );
     await waitFor(() => {
-      expect(getByText("test-tag")).toBeInTheDocument();
+      expect(screen.getByText("test-tag")).toBeInTheDocument();
     });
   });
 
   it("should not render tags in metadata textarea", async () => {
     vi.mocked(useAuthorized).mockReturnValue(baseUseAuthorizedMock);
 
-    const { container, getByText } = render(
+    const { container } = render(
       <KeyInfoView
         keyData={MOCK_KEY_DATA}
-        onClose={() => {}}
+        onClose={() => { }}
         keyId={"test-key-id"}
-        onKeyDataUpdate={() => {}}
+        onKeyDataUpdate={() => { }}
         teams={[]}
       />,
     );
     await waitFor(() => {
-      expect(getByText("Metadata")).toBeInTheDocument();
+      expect(screen.getByText("Metadata")).toBeInTheDocument();
       const metadataBlock = container.querySelector("pre");
       expect(metadataBlock).toBeInTheDocument();
       expect(metadataBlock?.textContent?.trim()).toBe("{}");
@@ -153,7 +169,7 @@ describe("KeyInfoView", () => {
 
     const keyData = { ...MOCK_KEY_DATA, user_id: "other-user-id" };
     render(
-      <KeyInfoView keyData={keyData} onClose={() => {}} keyId={"test-key-id"} onKeyDataUpdate={() => {}} teams={[]} />,
+      <KeyInfoView keyData={keyData} onClose={() => { }} keyId={"test-key-id"} onKeyDataUpdate={() => { }} teams={[]} />,
     );
 
     await waitFor(() => {
@@ -182,6 +198,7 @@ describe("KeyInfoView", () => {
           role: "admin",
         },
       ],
+      spend: 0,
     };
 
     vi.mocked(useTeams).mockReturnValue({
@@ -197,7 +214,7 @@ describe("KeyInfoView", () => {
 
     const keyData = { ...MOCK_KEY_DATA, team_id: teamId, user_id: "other-user-id" };
     render(
-      <KeyInfoView keyData={keyData} onClose={() => {}} keyId={"test-key-id"} onKeyDataUpdate={() => {}} teams={[]} />,
+      <KeyInfoView keyData={keyData} onClose={() => { }} keyId={"test-key-id"} onKeyDataUpdate={() => { }} teams={[]} />,
     );
 
     await waitFor(() => {
@@ -221,7 +238,7 @@ describe("KeyInfoView", () => {
     const ownerUserId = "owner-user-id";
     const keyData = { ...MOCK_KEY_DATA, user_id: ownerUserId };
     render(
-      <KeyInfoView keyData={keyData} onClose={() => {}} keyId={"test-key-id"} onKeyDataUpdate={() => {}} teams={[]} />,
+      <KeyInfoView keyData={keyData} onClose={() => { }} keyId={"test-key-id"} onKeyDataUpdate={() => { }} teams={[]} />,
     );
 
     await waitFor(() => {
@@ -244,7 +261,7 @@ describe("KeyInfoView", () => {
 
     const keyData = { ...MOCK_KEY_DATA, user_id: "owner-user-id" };
     render(
-      <KeyInfoView keyData={keyData} onClose={() => {}} keyId={"test-key-id"} onKeyDataUpdate={() => {}} teams={[]} />,
+      <KeyInfoView keyData={keyData} onClose={() => { }} keyId={"test-key-id"} onKeyDataUpdate={() => { }} teams={[]} />,
     );
 
     await waitFor(() => {
@@ -268,7 +285,7 @@ describe("KeyInfoView", () => {
     const ownerUserId = "internal-viewer-user-id";
     const keyData = { ...MOCK_KEY_DATA, user_id: ownerUserId };
     render(
-      <KeyInfoView keyData={keyData} onClose={() => {}} keyId={"test-key-id"} onKeyDataUpdate={() => {}} teams={[]} />,
+      <KeyInfoView keyData={keyData} onClose={() => { }} keyId={"test-key-id"} onKeyDataUpdate={() => { }} teams={[]} />,
     );
 
     await waitFor(() => {
@@ -296,6 +313,7 @@ describe("KeyInfoView", () => {
           role: "admin",
         },
       ],
+      spend: 0,
     };
 
     vi.mocked(useTeams).mockReturnValue({
@@ -309,10 +327,9 @@ describe("KeyInfoView", () => {
       userRole: "user",
     });
 
-    // Key has a different team_id that doesn't match any team in teamsData
     const keyData = { ...MOCK_KEY_DATA, team_id: "non-matching-team-id", user_id: "other-user-id" };
     render(
-      <KeyInfoView keyData={keyData} onClose={() => {}} keyId={"test-key-id"} onKeyDataUpdate={() => {}} teams={[]} />,
+      <KeyInfoView keyData={keyData} onClose={() => { }} keyId={"test-key-id"} onKeyDataUpdate={() => { }} teams={[]} />,
     );
 
     await waitFor(() => {
@@ -320,4 +337,129 @@ describe("KeyInfoView", () => {
       expect(screen.queryByText("Delete Key")).not.toBeInTheDocument();
     });
   });
+
+  it("should call onClose when back button is clicked", async () => {
+    vi.mocked(useAuthorized).mockReturnValue(baseUseAuthorizedMock);
+    const onCloseMock = vi.fn();
+
+    render(
+      <KeyInfoView
+        keyData={MOCK_KEY_DATA}
+        onClose={onCloseMock}
+        keyId={"test-key-id"}
+        onKeyDataUpdate={() => { }}
+        teams={[]}
+      />,
+    );
+
+    await waitFor(() => {
+      expect(screen.getByRole("button", { name: /back to keys/i })).toBeInTheDocument();
+    });
+
+    const backButton = screen.getByRole("button", { name: /back to keys/i });
+    await userEvent.click(backButton);
+
+    expect(onCloseMock).toHaveBeenCalledTimes(1);
+  });
+
+
+  it("should show edit button in settings tab when user has write access", async () => {
+    vi.mocked(useAuthorized).mockReturnValue({
+      ...baseUseAuthorizedMock,
+      userRole: "Admin",
+    });
+
+    render(
+      <KeyInfoView
+        keyData={MOCK_KEY_DATA}
+        onClose={() => { }}
+        keyId={"test-key-id"}
+        onKeyDataUpdate={() => { }}
+        teams={[]}
+      />,
+    );
+
+    await waitFor(() => {
+      const settingsTab = screen.getByRole("tab", { name: /settings/i });
+      expect(settingsTab).toBeInTheDocument();
+    });
+
+    const settingsTab = screen.getByRole("tab", { name: /settings/i });
+    await userEvent.click(settingsTab);
+
+    await waitFor(() => {
+      expect(screen.getByRole("button", { name: /edit settings/i })).toBeInTheDocument();
+    });
+  });
+
+
+  it("should display guardrails when present", async () => {
+    vi.mocked(useAuthorized).mockReturnValue(baseUseAuthorizedMock);
+
+    const keyDataWithGuardrails = {
+      ...MOCK_KEY_DATA,
+      metadata: {
+        ...MOCK_KEY_DATA.metadata,
+        guardrails: ["guardrail-1", "guardrail-2"],
+      },
+    };
+
+    render(
+      <KeyInfoView
+        keyData={keyDataWithGuardrails}
+        onClose={() => { }}
+        keyId={"test-key-id"}
+        onKeyDataUpdate={() => { }}
+        teams={[]}
+      />,
+    );
+
+    await waitFor(() => {
+      expect(screen.getByText("Guardrails")).toBeInTheDocument();
+    });
+  });
+
+  it("should display policies when present", async () => {
+    vi.mocked(useAuthorized).mockReturnValue(baseUseAuthorizedMock);
+
+    const keyDataWithPolicies = {
+      ...MOCK_KEY_DATA,
+      metadata: {
+        ...MOCK_KEY_DATA.metadata,
+        policies: ["policy-1"],
+      },
+    };
+
+    render(
+      <KeyInfoView
+        keyData={keyDataWithPolicies}
+        onClose={() => { }}
+        keyId={"test-key-id"}
+        onKeyDataUpdate={() => { }}
+        teams={[]}
+      />,
+    );
+
+    await waitFor(() => {
+      expect(screen.getByText("Policies")).toBeInTheDocument();
+    });
+  });
+
+  it("should display no key found message when keyData is undefined", async () => {
+    vi.mocked(useAuthorized).mockReturnValue(baseUseAuthorizedMock);
+
+    render(
+      <KeyInfoView
+        keyData={undefined}
+        onClose={() => { }}
+        keyId={"test-key-id"}
+        onKeyDataUpdate={() => { }}
+        teams={[]}
+      />,
+    );
+
+    await waitFor(() => {
+      expect(screen.getByText("Key not found")).toBeInTheDocument();
+    });
+  });
 });

ui/litellm-dashboard/src/components/templates/key_info_view.tsx

@@ -1,9 +1,10 @@
+import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized";
 import useTeams from "@/app/(dashboard)/hooks/useTeams";
 import { formatNumberWithCommas, copyToClipboard as utilCopyToClipboard } from "@/utils/dataUtils";
 import { mapEmptyStringToNull } from "@/utils/keyUpdateUtils";
 import { ArrowLeftIcon, RefreshIcon, TrashIcon } from "@heroicons/react/outline";
 import { Badge, Button, Card, Grid, Tab, TabGroup, TabList, TabPanel, TabPanels, Text, Title } from "@tremor/react";
-import { Button as AntdButton, Form, Tooltip } from "antd";
+import { Button as AntdButton, Form, Tag, Tooltip } from "antd";
 import { CheckIcon, CopyIcon } from "lucide-react";
 import { useEffect, useState } from "react";
 import { isProxyAdminRole, isUserTeamAdminForSingleTeam, rolesWithWriteAccess } from "../../utils/roles";
@@ -14,12 +15,11 @@ import { extractLoggingSettings, formatMetadataForDisplay, stripTagsFromMetadata
 import { KeyResponse } from "../key_team_helpers/key_list";
 import LoggingSettingsView from "../logging_settings_view";
 import NotificationManager from "../molecules/notifications_manager";
-import { keyDeleteCall, keyUpdateCall, getPolicyInfoWithGuardrails } from "../networking";
+import { getPolicyInfoWithGuardrails, keyDeleteCall, keyUpdateCall } from "../networking";
 import ObjectPermissionsView from "../object_permissions_view";
 import { RegenerateKeyModal } from "../organisms/regenerate_key_modal";
 import { parseErrorMessage } from "../shared/errorUtils";
 import { KeyEditView } from "./key_edit_view";
-import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized";
 
 interface KeyInfoViewProps {
   keyId: string;
@@ -206,8 +206,8 @@ export default function KeyInfoView({
             ...(formValues.logging_settings ? { logging: formValues.logging_settings } : {}),
             ...(formValues.disabled_callbacks?.length > 0
               ? {
-                  litellm_disabled_callbacks: mapDisplayToInternalNames(formValues.disabled_callbacks),
-                }
+                litellm_disabled_callbacks: mapDisplayToInternalNames(formValues.disabled_callbacks),
+              }
               : {}),
           };
         } catch (error) {
@@ -225,8 +225,8 @@ export default function KeyInfoView({
           ...(formValues.logging_settings ? { logging: formValues.logging_settings } : {}),
           ...(formValues.disabled_callbacks?.length > 0
             ? {
-                litellm_disabled_callbacks: mapDisplayToInternalNames(formValues.disabled_callbacks),
-              }
+              litellm_disabled_callbacks: mapDisplayToInternalNames(formValues.disabled_callbacks),
+            }
             : {}),
         };
       }
@@ -334,7 +334,6 @@ export default function KeyInfoView({
     });
     return `${dateStr} at ${timeStr}`;
   };
-  console.log("userRole", userRole);
 
   const canModifyKey =
     isProxyAdminRole(userRole || "") ||
@@ -364,11 +363,10 @@ export default function KeyInfoView({
               size="small"
               icon={copiedStates["key-id"] ? <CheckIcon size={12} /> : <CopyIcon size={12} />}
               onClick={() => copyToClipboard(currentKeyData.token_id || currentKeyData.token, "key-id")}
-              className={`ml-2 transition-all duration-200${
-                copiedStates["key-id"]
-                  ? "text-green-600 bg-green-50 border-green-200"
-                  : "text-gray-500 hover:text-gray-700 hover:bg-gray-100"
-              }`}
+              className={`ml-2 transition-all duration-200${copiedStates["key-id"]
+                ? "text-green-600 bg-green-50 border-green-200"
+                : "text-gray-500 hover:text-gray-700 hover:bg-gray-100"
+                }`}
             />
           </div>
 
@@ -691,10 +689,10 @@ export default function KeyInfoView({
                     <div className="flex flex-wrap gap-2 mt-1">
                       {Array.isArray(currentKeyData.metadata?.tags) && currentKeyData.metadata.tags.length > 0
                         ? currentKeyData.metadata.tags.map((tag, index) => (
-                            <span key={index} className="px-2 mr-2 py-1 bg-blue-100 rounded text-xs">
-                              {tag}
-                            </span>
-                          ))
+                          <span key={index} className="px-2 mr-2 py-1 bg-blue-100 rounded text-xs">
+                            {tag}
+                          </span>
+                        ))
                         : "No tags specified"}
                     </div>
                   </div>
@@ -704,24 +702,39 @@ export default function KeyInfoView({
                     <Text>
                       {Array.isArray(currentKeyData.metadata?.prompts) && currentKeyData.metadata.prompts.length > 0
                         ? currentKeyData.metadata.prompts.map((prompt, index) => (
-                            <span key={index} className="px-2 mr-2 py-1 bg-blue-100 rounded text-xs">
-                              {prompt}
-                            </span>
-                          ))
+                          <span key={index} className="px-2 mr-2 py-1 bg-blue-100 rounded text-xs">
+                            {prompt}
+                          </span>
+                        ))
                         : "No prompts specified"}
                     </Text>
                   </div>
 
+                  <div>
+                    <Text className="font-medium">Allowed Routes</Text>
+                    <div className="flex flex-wrap gap-2 mt-1">
+                      {Array.isArray(currentKeyData.allowed_routes) && currentKeyData.allowed_routes.length > 0 ? (
+                        currentKeyData.allowed_routes.map((route, index) => (
+                          <span key={index} className="px-2 py-1 bg-blue-100 rounded text-xs">
+                            {route}
+                          </span>
+                        ))
+                      ) : (
+                        <Tag color="green">All routes allowed</Tag>
+                      )}
+                    </div>
+                  </div>
+
                   <div>
                     <Text className="font-medium">Allowed Pass Through Routes</Text>
                     <Text>
                       {Array.isArray(currentKeyData.metadata?.allowed_passthrough_routes) &&
-                      currentKeyData.metadata.allowed_passthrough_routes.length > 0
+                        currentKeyData.metadata.allowed_passthrough_routes.length > 0
                         ? currentKeyData.metadata.allowed_passthrough_routes.map((route, index) => (
-                            <span key={index} className="px-2 mr-2 py-1 bg-blue-100 rounded text-xs">
-                              {route}
-                            </span>
-                          ))
+                          <span key={index} className="px-2 mr-2 py-1 bg-blue-100 rounded text-xs">
+                            {route}
+                          </span>
+                        ))
                         : "No pass through routes specified"}
                     </Text>
                   </div>

ui/litellm-dashboard/src/components/templates/model_dashboard.tsx

@@ -1368,6 +1368,7 @@ const OldModelDashboard: React.FC<ModelDashboardProps> = ({
                     all_models_on_proxy={all_models_on_proxy}
                     getDisplayModelName={getDisplayModelName}
                     setSelectedModelId={setSelectedModelId}
+                    teams={teams}
                   />
                 </TabPanel>
                 <TabPanel>

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

@@ -0,0 +1,504 @@
+import { screen, waitFor } from "@testing-library/react";
+import userEvent from "@testing-library/user-event";
+import { beforeEach, describe, expect, it, vi } from "vitest";
+import { renderWithProviders } from "../../tests/test-utils";
+import { UserEditView } from "./user_edit_view";
+
+vi.mock("./key_team_helpers/fetch_available_models_team_key", () => ({
+  getModelDisplayName: vi.fn((model: string) => model),
+}));
+
+vi.mock("../utils/roles", () => ({
+  all_admin_roles: ["Admin", "Admin Viewer", "proxy_admin", "proxy_admin_viewer", "org_admin"],
+}));
+
+vi.mock("antd", async (importOriginal) => {
+  const actual = await importOriginal<typeof import("antd")>();
+  const React = await import("react");
+  const SelectComponent = ({
+    value,
+    onChange,
+    mode,
+    children,
+    placeholder,
+    disabled,
+    style,
+    allowClear,
+    ...props
+  }: any) => {
+    const isMultiple = mode === "multiple";
+    const selectValue = isMultiple ? (Array.isArray(value) ? value : []) : value || "";
+    return React.createElement(
+      "select",
+      {
+        multiple: isMultiple,
+        value: selectValue,
+        onChange: (e: React.ChangeEvent<HTMLSelectElement>) => {
+          const selectedValues = Array.from(e.target.selectedOptions, (option) => option.value);
+          onChange(isMultiple ? selectedValues : selectedValues[0] || undefined);
+        },
+        disabled,
+        placeholder,
+        style,
+        "aria-label": placeholder || "Select",
+        role: "combobox",
+        ...props,
+      },
+      children,
+    );
+  };
+  SelectComponent.Option = ({ value: optionValue, children: optionChildren }: any) =>
+    React.createElement("option", { value: optionValue }, optionChildren);
+  return {
+    ...actual,
+    Select: SelectComponent,
+    Tooltip: ({ children }: { children?: React.ReactNode }) => React.createElement(React.Fragment, null, children),
+    Checkbox: ({ checked, onChange, children, ...props }: any) =>
+      React.createElement(
+        "label",
+        { style: { display: "flex", alignItems: "center", gap: "8px" } },
+        React.createElement("input", {
+          type: "checkbox",
+          checked: checked,
+          onChange: (e: React.ChangeEvent<HTMLInputElement>) => onChange({ target: { checked: e.target.checked } }),
+          ...props,
+        }),
+        children,
+      ),
+  };
+});
+
+vi.mock("@tremor/react", async (importOriginal) => {
+  const actual = await importOriginal<typeof import("@tremor/react")>();
+  const React = await import("react");
+  return {
+    ...actual,
+    SelectItem: ({ value, children, title }: any) => {
+      const childText = React.Children.toArray(children)
+        .map((child: any) => (typeof child === "string" ? child : child?.props?.children || ""))
+        .join(" ");
+      return React.createElement("option", { value, title }, childText || title || value);
+    },
+  };
+});
+
+describe("UserEditView", () => {
+  const MOCK_USER_DATA = {
+    user_id: "user-123",
+    user_info: {
+      user_email: "test@example.com",
+      user_alias: "Test User",
+      user_role: "proxy_admin",
+      models: ["gpt-4", "gpt-3.5-turbo"],
+      max_budget: 100.5,
+      budget_duration: "30d",
+      metadata: {
+        key1: "value1",
+        key2: "value2",
+      },
+    },
+  };
+
+  const MOCK_POSSIBLE_UI_ROLES = {
+    proxy_admin: {
+      ui_label: "Proxy Admin",
+      description: "Full access to proxy",
+    },
+    proxy_admin_viewer: {
+      ui_label: "Proxy Admin Viewer",
+      description: "Read-only access",
+    },
+    user: {
+      ui_label: "User",
+      description: "Standard user",
+    },
+  };
+
+  const defaultProps = {
+    userData: MOCK_USER_DATA,
+    onCancel: vi.fn(),
+    onSubmit: vi.fn(),
+    teams: null,
+    accessToken: "test-token",
+    userID: "current-user-1",
+    userRole: "Admin",
+    userModels: ["gpt-4", "gpt-3.5-turbo", "claude-3"],
+    possibleUIRoles: MOCK_POSSIBLE_UI_ROLES,
+    isBulkEdit: false,
+  };
+
+  beforeEach(() => {
+    vi.clearAllMocks();
+  });
+
+  it("should render", async () => {
+    renderWithProviders(<UserEditView {...defaultProps} />);
+
+    await waitFor(() => {
+      expect(screen.getByRole("button", { name: /save changes/i })).toBeInTheDocument();
+    });
+  });
+
+  it("should display user ID field when not in bulk edit mode", async () => {
+    renderWithProviders(<UserEditView {...defaultProps} />);
+
+    await waitFor(() => {
+      expect(screen.getByLabelText("User ID")).toBeInTheDocument();
+    });
+
+    const userIdInput = screen.getByLabelText("User ID");
+    expect(userIdInput).toBeDisabled();
+    expect(userIdInput).toHaveValue("user-123");
+  });
+
+  it("should not display user ID field when in bulk edit mode", async () => {
+    renderWithProviders(<UserEditView {...defaultProps} isBulkEdit={true} />);
+
+    await waitFor(() => {
+      expect(screen.getByRole("button", { name: /save changes/i })).toBeInTheDocument();
+    });
+
+    expect(screen.queryByLabelText("User ID")).not.toBeInTheDocument();
+  });
+
+  it("should display email field when not in bulk edit mode", async () => {
+    renderWithProviders(<UserEditView {...defaultProps} />);
+
+    await waitFor(() => {
+      expect(screen.getByLabelText("Email")).toBeInTheDocument();
+    });
+
+    const emailInput = screen.getByLabelText("Email");
+    expect(emailInput).toHaveValue("test@example.com");
+  });
+
+  it("should not display email field when in bulk edit mode", async () => {
+    renderWithProviders(<UserEditView {...defaultProps} isBulkEdit={true} />);
+
+    await waitFor(() => {
+      expect(screen.getByRole("button", { name: /save changes/i })).toBeInTheDocument();
+    });
+
+    expect(screen.queryByLabelText("Email")).not.toBeInTheDocument();
+  });
+
+  it("should display user alias field with initial value", async () => {
+    renderWithProviders(<UserEditView {...defaultProps} />);
+
+    await waitFor(() => {
+      expect(screen.getByLabelText("User Alias")).toBeInTheDocument();
+    });
+
+    const aliasInput = screen.getByLabelText("User Alias");
+    expect(aliasInput).toHaveValue("Test User");
+  });
+
+  it("should display personal models select with available models", async () => {
+    renderWithProviders(<UserEditView {...defaultProps} />);
+
+    await waitFor(() => {
+      expect(screen.getByText("Personal Models")).toBeInTheDocument();
+    });
+
+    const modelsSelect = screen.getByRole("combobox", { name: /select models/i });
+    expect(modelsSelect).toBeInTheDocument();
+  });
+
+  it("should disable models select when user role is not admin", async () => {
+    renderWithProviders(<UserEditView {...defaultProps} userRole="user" />);
+
+    await waitFor(() => {
+      const modelsSelect = screen.getByRole("combobox", { name: /select models/i });
+      expect(modelsSelect).toBeDisabled();
+    });
+  });
+
+  it("should enable models select when user role is admin", async () => {
+    renderWithProviders(<UserEditView {...defaultProps} userRole="Admin" />);
+
+    await waitFor(() => {
+      const modelsSelect = screen.getByRole("combobox", { name: /select models/i });
+      expect(modelsSelect).not.toBeDisabled();
+    });
+  });
+
+  it("should display max budget input field", async () => {
+    renderWithProviders(<UserEditView {...defaultProps} />);
+
+    await waitFor(() => {
+      expect(screen.getByText("Max Budget (USD)")).toBeInTheDocument();
+    });
+  });
+
+  it("should display unlimited budget checkbox", async () => {
+    renderWithProviders(<UserEditView {...defaultProps} />);
+
+    await waitFor(() => {
+      expect(screen.getByLabelText("Unlimited Budget")).toBeInTheDocument();
+    });
+  });
+
+  it("should set unlimited budget checkbox when max_budget is null", async () => {
+    const userDataWithNullBudget = {
+      ...MOCK_USER_DATA,
+      user_info: {
+        ...MOCK_USER_DATA.user_info,
+        max_budget: null,
+      },
+    };
+
+    renderWithProviders(<UserEditView {...defaultProps} userData={userDataWithNullBudget} />);
+
+    await waitFor(() => {
+      const checkbox = screen.getByLabelText("Unlimited Budget");
+      expect(checkbox).toBeChecked();
+    });
+  });
+
+  it("should disable budget input when unlimited budget is checked", async () => {
+    const userDataWithNullBudget = {
+      ...MOCK_USER_DATA,
+      user_info: {
+        ...MOCK_USER_DATA.user_info,
+        max_budget: null,
+      },
+    };
+
+    renderWithProviders(<UserEditView {...defaultProps} userData={userDataWithNullBudget} />);
+
+    await waitFor(() => {
+      const budgetInput = screen.getByRole("spinbutton", { name: /max budget/i });
+      expect(budgetInput).toBeDisabled();
+    });
+  });
+
+  it("should enable budget input when unlimited budget is unchecked", async () => {
+    renderWithProviders(<UserEditView {...defaultProps} />);
+
+    await waitFor(() => {
+      const budgetInput = screen.getByRole("spinbutton", { name: /max budget/i });
+      expect(budgetInput).not.toBeDisabled();
+    });
+  });
+
+  it("should clear budget value when unlimited budget is checked", async () => {
+    renderWithProviders(<UserEditView {...defaultProps} />);
+
+    await waitFor(() => {
+      expect(screen.getByLabelText("Unlimited Budget")).toBeInTheDocument();
+    });
+
+    const checkbox = screen.getByLabelText("Unlimited Budget");
+    await userEvent.click(checkbox);
+
+    await waitFor(() => {
+      const budgetInput = screen.getByRole("spinbutton", { name: /max budget/i });
+      expect(budgetInput).toHaveValue(null);
+    });
+  });
+
+  it("should display metadata textarea with formatted JSON", async () => {
+    renderWithProviders(<UserEditView {...defaultProps} />);
+
+    await waitFor(() => {
+      expect(screen.getByLabelText("Metadata")).toBeInTheDocument();
+    });
+
+    const metadataTextarea = screen.getByLabelText("Metadata");
+    const expectedJson = JSON.stringify(MOCK_USER_DATA.user_info.metadata, null, 2);
+    expect(metadataTextarea).toHaveValue(expectedJson);
+  });
+
+  it("should display empty metadata textarea when metadata is undefined", async () => {
+    const userDataWithoutMetadata = {
+      ...MOCK_USER_DATA,
+      user_info: {
+        ...MOCK_USER_DATA.user_info,
+        metadata: undefined,
+      },
+    };
+
+    renderWithProviders(<UserEditView {...defaultProps} userData={userDataWithoutMetadata} />);
+
+    await waitFor(() => {
+      const metadataTextarea = screen.getByLabelText("Metadata");
+      expect(metadataTextarea).toHaveValue("");
+    });
+  });
+
+  it("should call onCancel when cancel button is clicked", async () => {
+    const onCancelMock = vi.fn();
+    renderWithProviders(<UserEditView {...defaultProps} onCancel={onCancelMock} />);
+
+    await waitFor(() => {
+      expect(screen.getByRole("button", { name: /cancel/i })).toBeInTheDocument();
+    });
+
+    const cancelButton = screen.getByRole("button", { name: /cancel/i });
+    await userEvent.click(cancelButton);
+
+    expect(onCancelMock).toHaveBeenCalledTimes(1);
+  });
+
+  it("should call onSubmit with form values when form is submitted", async () => {
+    const onSubmitMock = vi.fn();
+    renderWithProviders(<UserEditView {...defaultProps} onSubmit={onSubmitMock} />);
+
+    await waitFor(() => {
+      expect(screen.getByRole("button", { name: /save changes/i })).toBeInTheDocument();
+    });
+
+    const submitButton = screen.getByRole("button", { name: /save changes/i });
+    await userEvent.click(submitButton);
+
+    await waitFor(() => {
+      expect(onSubmitMock).toHaveBeenCalled();
+    });
+
+    const callArgs = onSubmitMock.mock.calls[0][0];
+    expect(callArgs.user_id).toBe("user-123");
+    expect(callArgs.user_email).toBe("test@example.com");
+    expect(callArgs.user_alias).toBe("Test User");
+    expect(callArgs.user_role).toBe("proxy_admin");
+    expect(callArgs.models).toEqual(["gpt-4", "gpt-3.5-turbo"]);
+    expect(callArgs.max_budget).toBe(100.5);
+    expect(callArgs.budget_duration).toBe("30d");
+    expect(callArgs.metadata).toEqual(MOCK_USER_DATA.user_info.metadata);
+  });
+
+  it("should set max_budget to null when unlimited budget is checked on submit", async () => {
+    const onSubmitMock = vi.fn();
+    const userDataWithNullBudget = {
+      ...MOCK_USER_DATA,
+      user_info: {
+        ...MOCK_USER_DATA.user_info,
+        max_budget: null,
+      },
+    };
+
+    renderWithProviders(<UserEditView {...defaultProps} userData={userDataWithNullBudget} onSubmit={onSubmitMock} />);
+
+    await waitFor(() => {
+      expect(screen.getByRole("button", { name: /save changes/i })).toBeInTheDocument();
+    });
+
+    const submitButton = screen.getByRole("button", { name: /save changes/i });
+    await userEvent.click(submitButton);
+
+    await waitFor(() => {
+      expect(onSubmitMock).toHaveBeenCalled();
+    });
+
+    const callArgs = onSubmitMock.mock.calls[0][0];
+    expect(callArgs.max_budget).toBeNull();
+  });
+
+  it("should require budget when unlimited budget is not checked", async () => {
+    renderWithProviders(<UserEditView {...defaultProps} />);
+
+    await waitFor(() => {
+      expect(screen.getByText("Max Budget (USD)")).toBeInTheDocument();
+    });
+
+    const budgetInput = screen.getByRole("spinbutton", { name: /max budget/i });
+    await userEvent.clear(budgetInput);
+
+    const checkbox = screen.getByLabelText("Unlimited Budget");
+    expect(checkbox).not.toBeChecked();
+
+    const submitButton = screen.getByRole("button", { name: /save changes/i });
+    await userEvent.click(submitButton);
+
+    await waitFor(() => {
+      expect(screen.getByText("Please enter a budget or select Unlimited Budget")).toBeInTheDocument();
+    });
+  });
+
+  it("should allow submission when unlimited budget is checked even if budget is empty", async () => {
+    const onSubmitMock = vi.fn();
+    renderWithProviders(<UserEditView {...defaultProps} onSubmit={onSubmitMock} />);
+
+    await waitFor(() => {
+      expect(screen.getByLabelText("Unlimited Budget")).toBeInTheDocument();
+    });
+
+    const checkbox = screen.getByLabelText("Unlimited Budget");
+    await userEvent.click(checkbox);
+
+    await waitFor(() => {
+      expect(checkbox).toBeChecked();
+    });
+
+    const submitButton = screen.getByRole("button", { name: /save changes/i });
+    await userEvent.click(submitButton);
+
+    await waitFor(() => {
+      expect(onSubmitMock).toHaveBeenCalled();
+    });
+  });
+
+  it("should update form values when userData changes", async () => {
+    const { rerender } = renderWithProviders(<UserEditView {...defaultProps} />);
+
+    await waitFor(() => {
+      expect(screen.getByLabelText("User Alias")).toHaveValue("Test User");
+    });
+
+    const updatedUserData = {
+      ...MOCK_USER_DATA,
+      user_info: {
+        ...MOCK_USER_DATA.user_info,
+        user_alias: "Updated Alias",
+      },
+    };
+
+    rerender(<UserEditView {...defaultProps} userData={updatedUserData} />);
+
+    await waitFor(() => {
+      expect(screen.getByLabelText("User Alias")).toHaveValue("Updated Alias");
+    });
+  });
+
+  it("should handle user data with empty models array", async () => {
+    const userDataWithEmptyModels = {
+      ...MOCK_USER_DATA,
+      user_info: {
+        ...MOCK_USER_DATA.user_info,
+        models: [],
+      },
+    };
+
+    renderWithProviders(<UserEditView {...defaultProps} userData={userDataWithEmptyModels} />);
+
+    await waitFor(() => {
+      expect(screen.getByRole("button", { name: /save changes/i })).toBeInTheDocument();
+    });
+
+    const submitButton = screen.getByRole("button", { name: /save changes/i });
+    await userEvent.click(submitButton);
+
+    await waitFor(() => {
+      expect(defaultProps.onSubmit).toHaveBeenCalled();
+    });
+
+    const callArgs = defaultProps.onSubmit.mock.calls[0][0];
+    expect(callArgs.models).toEqual([]);
+  });
+
+  it("should handle user data with undefined max_budget", async () => {
+    const userDataWithUndefinedBudget = {
+      ...MOCK_USER_DATA,
+      user_info: {
+        ...MOCK_USER_DATA.user_info,
+        max_budget: undefined,
+      },
+    };
+
+    renderWithProviders(<UserEditView {...defaultProps} userData={userDataWithUndefinedBudget} />);
+
+    await waitFor(() => {
+      const checkbox = screen.getByLabelText("Unlimited Budget");
+      expect(checkbox).toBeChecked();
+    });
+  });
+});

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

@@ -1,12 +1,11 @@
-import React from "react";
-import { Form, Select, Tooltip } from "antd";
-import NumericalInput from "./shared/numerical_input";
-import { TextInput, Textarea, SelectItem } from "@tremor/react";
-import { Button } from "@tremor/react";
-import { getModelDisplayName } from "./key_team_helpers/fetch_available_models_team_key";
-import { all_admin_roles } from "../utils/roles";
 import { InfoCircleOutlined } from "@ant-design/icons";
+import { Button, SelectItem, TextInput, Textarea } from "@tremor/react";
+import { Checkbox, Form, Select, Tooltip } from "antd";
+import React, { useState } from "react";
+import { all_admin_roles } from "../utils/roles";
 import BudgetDurationDropdown from "./common_components/budget_duration_dropdown";
+import { getModelDisplayName } from "./key_team_helpers/fetch_available_models_team_key";
+import NumericalInput from "./shared/numerical_input";
 
 interface UserEditViewProps {
   userData: any;
@@ -34,21 +33,34 @@ export function UserEditView({
   isBulkEdit = false,
 }: UserEditViewProps) {
   const [form] = Form.useForm();
+  const [unlimitedBudget, setUnlimitedBudget] = useState(false);
 
   // Set initial form values
   React.useEffect(() => {
+    const maxBudget = userData.user_info?.max_budget;
+    const isUnlimited = maxBudget === null || maxBudget === undefined;
+    setUnlimitedBudget(isUnlimited);
+
     form.setFieldsValue({
       user_id: userData.user_id,
       user_email: userData.user_info?.user_email,
       user_alias: userData.user_info?.user_alias,
       user_role: userData.user_info?.user_role,
       models: userData.user_info?.models || [],
-      max_budget: userData.user_info?.max_budget,
+      max_budget: isUnlimited ? "" : maxBudget,
       budget_duration: userData.user_info?.budget_duration,
       metadata: userData.user_info?.metadata ? JSON.stringify(userData.user_info.metadata, null, 2) : undefined,
     });
   }, [userData, form]);
 
+  const handleUnlimitedBudgetChange = (e: any) => {
+    const checked = e.target.checked;
+    setUnlimitedBudget(checked);
+    if (checked) {
+      form.setFieldsValue({ max_budget: "" });
+    }
+  };
+
   const handleSubmit = (values: any) => {
     // Convert metadata back to an object if it exists and is a string
     if (values.metadata && typeof values.metadata === "string") {
@@ -60,6 +72,10 @@ export function UserEditView({
       }
     }
 
+    if (unlimitedBudget || values.max_budget === "" || values.max_budget === undefined) {
+      values.max_budget = null;
+    }
+
     onSubmit(values);
   };
 
@@ -138,8 +154,36 @@ export function UserEditView({
         </Select>
       </Form.Item>
 
-      <Form.Item label="Max Budget (USD)" name="max_budget">
-        <NumericalInput step={0.01} precision={2} style={{ width: "100%" }} />
+      <Form.Item
+        label={
+          <div style={{ display: "flex", alignItems: "center", gap: "12px" }}>
+            <span>Max Budget (USD)</span>
+            <Checkbox
+              checked={unlimitedBudget}
+              onChange={handleUnlimitedBudgetChange}
+            >
+              Unlimited Budget
+            </Checkbox>
+          </div>
+        }
+        name="max_budget"
+        rules={[
+          {
+            validator: (_, value) => {
+              if (!unlimitedBudget && (value === "" || value === null || value === undefined)) {
+                return Promise.reject(new Error("Please enter a budget or select Unlimited Budget"));
+              }
+              return Promise.resolve();
+            },
+          },
+        ]}
+      >
+        <NumericalInput
+          step={0.01}
+          precision={2}
+          style={{ width: "100%" }}
+          disabled={unlimitedBudget}
+        />
       </Form.Item>
 
       <Form.Item label="Reset Budget" name="budget_duration">

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

@@ -108,7 +108,7 @@ const VectorStoreForm: React.FC<VectorStoreFormProps> = ({
   };
 
   return (
-    <Modal title="Add New Vector Store" visible={isVisible} width={1000} footer={null} onCancel={handleCancel}>
+    <Modal title="Add New Vector Store" open={isVisible} width={1000} footer={null} onCancel={handleCancel}>
       <Form form={form} onFinish={handleCreate} labelCol={{ span: 8 }} wrapperCol={{ span: 16 }} labelAlign="left">
         <Form.Item
           label={