From d65d97d609b51e98fca8cbd38ec752c2790158e8 Mon Sep 17 00:00:00 2001
From: Ishaan Jaffer <ishaanjaffer0324@gmail.com>
Date: Sun, 14 Dec 2025 14:06:35 -0800
Subject: [PATCH] fix: a2a gateway - azure foundry invoke

---
 .../litellm_completion_bridge/handler.py      | 46 ++++++++++++-------
 1 file changed, 30 insertions(+), 16 deletions(-)

diff --git a/litellm/a2a_protocol/litellm_completion_bridge/handler.py b/litellm/a2a_protocol/litellm_completion_bridge/handler.py
index 2eab2551833..cc222d3ee18 100644
--- a/litellm/a2a_protocol/litellm_completion_bridge/handler.py
+++ b/litellm/a2a_protocol/litellm_completion_bridge/handler.py
@@ -55,7 +55,6 @@ class A2ACompletionBridgeHandler:
         # Get completion params
         custom_llm_provider = litellm_params.get("custom_llm_provider")
         model = litellm_params.get("model", "agent")
-        api_key = litellm_params.get("api_key")
 
         # Build full model string if provider specified
         # Skip prepending if model already starts with the provider prefix
@@ -68,14 +67,22 @@ class A2ACompletionBridgeHandler:
             f"A2A completion bridge: model={full_model}, api_base={api_base}"
         )
 
+        # Build completion params dict
+        completion_params = {
+            "model": full_model,
+            "messages": openai_messages,
+            "api_base": api_base,
+            "stream": False,
+        }
+        # Add litellm_params (contains api_key, client_id, client_secret, tenant_id, etc.)
+        litellm_params_to_add = {
+            k: v for k, v in litellm_params.items()
+            if k not in ("model", "custom_llm_provider")
+        }
+        completion_params.update(litellm_params_to_add)
+
         # Call litellm.acompletion
-        response = await litellm.acompletion(
-            model=full_model,
-            messages=openai_messages,
-            api_base=api_base,
-            api_key=api_key,
-            stream=False,
-        )
+        response = await litellm.acompletion(**completion_params)
 
         # Transform response to A2A format
         a2a_response = A2ACompletionBridgeTransformation.openai_response_to_a2a_response(
@@ -129,7 +136,6 @@ class A2ACompletionBridgeHandler:
         # Get completion params
         custom_llm_provider = litellm_params.get("custom_llm_provider")
         model = litellm_params.get("model", "agent")
-        api_key = litellm_params.get("api_key")
 
         # Build full model string if provider specified
         # Skip prepending if model already starts with the provider prefix
@@ -142,6 +148,20 @@ class A2ACompletionBridgeHandler:
             f"A2A completion bridge streaming: model={full_model}, api_base={api_base}"
         )
 
+        # Build completion params dict
+        completion_params = {
+            "model": full_model,
+            "messages": openai_messages,
+            "api_base": api_base,
+            "stream": True,
+        }
+        # Add litellm_params (contains api_key, client_id, client_secret, tenant_id, etc.)
+        litellm_params_to_add = {
+            k: v for k, v in litellm_params.items()
+            if k not in ("model", "custom_llm_provider")
+        }
+        completion_params.update(litellm_params_to_add)
+
         # 1. Emit initial task event (kind: "task", status: "submitted")
         task_event = A2ACompletionBridgeTransformation.create_task_event(ctx)
         yield task_event
@@ -156,13 +176,7 @@ class A2ACompletionBridgeHandler:
         yield working_event
 
         # Call litellm.acompletion with streaming
-        response = await litellm.acompletion(
-            model=full_model,
-            messages=openai_messages,
-            api_base=api_base,
-            api_key=api_key,
-            stream=True,
-        )
+        response = await litellm.acompletion(**completion_params)
 
         # 3. Accumulate content and emit artifact update
         accumulated_text = ""