From fe4cff8595622d68eb97efbd4d5276eca0ebbf0c Mon Sep 17 00:00:00 2001
From: Cursor Agent <cursoragent@cursor.com>
Date: Sat, 7 Mar 2026 05:38:20 +0000
Subject: [PATCH] fix: use get_async_httpx_client in a2a_protocol and add
 master_key monkeypatch to files tests

- Replace httpx.AsyncClient() with get_async_httpx_client() in a2a_protocol/main.py
  to satisfy the ensure_async_clients_test CI check
- Add httpxSpecialProvider.A2AProvider enum value
- Add master_key=None monkeypatch to test_managed_files_with_loadbalancing

Co-authored-by: Ishaan Jaff <ishaan-jaff@users.noreply.github.com>
---
 litellm/a2a_protocol/main.py                  | 23 ++++++++++++-------
 litellm/types/llms/custom_http.py             |  1 +
 .../test_files_endpoint.py                    |  2 ++
 3 files changed, 18 insertions(+), 8 deletions(-)

diff --git a/litellm/a2a_protocol/main.py b/litellm/a2a_protocol/main.py
index 47a6d8be558..ed4c9c36c2f 100644
--- a/litellm/a2a_protocol/main.py
+++ b/litellm/a2a_protocol/main.py
@@ -653,15 +653,22 @@ async def create_a2a_client(
 
     verbose_logger.info(f"Creating A2A client for {base_url}")
 
-    # Always create a fresh httpx client per A2A call so that per-agent auth
-    # headers (extra_headers) are never shared across agents or requests.
-    # Mutating a cached shared client would cause headers from one agent to
-    # bleed into requests made to a different agent.
-    httpx_client = httpx.AsyncClient(
-        timeout=httpx.Timeout(timeout),
-        headers=extra_headers or {},
-    )
+    # Use get_async_httpx_client with per-agent params so that different agents
+    # (with different extra_headers) get separate cached clients.  The params
+    # dict is hashed into the cache key, keeping agent auth isolated while
+    # still reusing connections within the same agent.
+    _client_params: dict = {"timeout": timeout}
     if extra_headers:
+        # Include sorted header keys in params so each unique header set
+        # produces a distinct cache key.
+        _client_params["extra_headers"] = str(sorted(extra_headers.items()))
+    _async_handler = get_async_httpx_client(
+        llm_provider=httpxSpecialProvider.A2AProvider,
+        params=_client_params,
+    )
+    httpx_client = _async_handler.client
+    if extra_headers:
+        httpx_client.headers.update(extra_headers)
         verbose_proxy_logger.debug(
             f"A2A client created with extra_headers={list(extra_headers.keys())}"
         )
diff --git a/litellm/types/llms/custom_http.py b/litellm/types/llms/custom_http.py
index 32f3dc2efaf..8f192d876c4 100644
--- a/litellm/types/llms/custom_http.py
+++ b/litellm/types/llms/custom_http.py
@@ -24,6 +24,7 @@ class httpxSpecialProvider(str, Enum):
     Search = "search"
     MCP = "mcp"
     RAG = "rag"
+    A2AProvider = "a2a_provider"
     A2A = "a2a"
     PromptManagement = "prompt_management"
     UI = "ui"
diff --git a/tests/test_litellm/proxy/openai_files_endpoint/test_files_endpoint.py b/tests/test_litellm/proxy/openai_files_endpoint/test_files_endpoint.py
index bfbc4888295..b1420a48497 100644
--- a/tests/test_litellm/proxy/openai_files_endpoint/test_files_endpoint.py
+++ b/tests/test_litellm/proxy/openai_files_endpoint/test_files_endpoint.py
@@ -962,6 +962,8 @@ def test_managed_files_with_loadbalancing(mocker: MockerFixture, monkeypatch, ll
     monkeypatch.setattr(
         "litellm.proxy.proxy_server.proxy_logging_obj", proxy_logging_obj
     )
+    # Disable auth so the test doesn't depend on master_key state from other tests
+    monkeypatch.setattr("litellm.proxy.proxy_server.master_key", None)
     
     # Create batch file content
     test_file_content = b'{"custom_id": "request-1", "method": "POST", "url": "/v1/chat/completions", "body": {"model": "gpt-3.5-turbo", "messages": [{"role": "user", "content": "Hello"}]}}'