a2a agent Header-Based Context Propagation (#19504)

* a2a agent Header-Based Context Propagation

* docs: add guide for A2A context header forwarding

docs/my-website/docs/a2a.md

@@ -193,6 +193,120 @@ The logs show:
   style={{width: '100%', display: 'block', margin: '2rem auto'}}
 />
 
+
+## Forwarding LiteLLM Context Headers
+
+When LiteLLM invokes your A2A agent, it sends special headers that enable:
+- **Trace Grouping**: All LLM calls from the same agent execution appear under one trace
+- **Agent Spend Tracking**: Costs are attributed to the specific agent
+
+| Header | Purpose |
+|--------|---------|
+| `X-LiteLLM-Trace-Id` | Links all LLM calls to the same execution flow |
+| `X-LiteLLM-Agent-Id` | Attributes spend to the correct agent |
+
+
+To enable these features, your A2A server must **forward these headers** to any LLM calls it makes back to LiteLLM.
+
+### Implementation Steps
+
+**Step 1: Extract headers from incoming A2A request**
+```python def get_litellm_headers(request) -> dict:
+    """Extract X-LiteLLM-* headers from incoming A2A request."""
+    all_headers = request.call_context.state.get('headers', {})
+    return {
+        k: v for k, v in all_headers.items() 
+        if k.lower().startswith('x-litellm-')
+    }
+```
+
+**Step 2: Forward headers to your LLM calls**
+Pass the extracted headers when making calls back to LiteLLM:
+<Tabs>
+<TabItem value="openai" label="OpenAI SDK" default>
+
+```python from openai import OpenAI
+
+headers = get_litellm_headers(request)
+
+client = OpenAI(
+    api_key="sk-your-litellm-key",
+    base_url="http://localhost:4000",
+    default_headers=headers,  # Forward headers
+)
+
+response = client.chat.completions.create(
+    model="gpt-4o",
+    messages=[{"role": "user", "content": "Hello"}]
+)
+```
+</TabItem>
+
+<TabItem value="langchain" label="LangChain">
+
+```python
+from langchain_openai import ChatOpenAI
+
+headers = get_litellm_headers(request)
+
+llm = ChatOpenAI(
+    model="gpt-4o",
+    openai_api_key="sk-your-litellm-key",
+    base_url="http://localhost:4000",
+    default_headers=headers,  # Forward headers
+)
+```
+</TabItem>
+<TabItem value="litellm" label="LiteLLM SDK">
+
+```python
+import litellm
+
+headers = get_litellm_headers(request)
+
+response = litellm.completion(
+    model="gpt-4o",
+    messages=[{"role": "user", "content": "Hello"}],
+    api_base="http://localhost:4000",
+    extra_headers=headers,  # Forward headers
+)
+```
+</TabItem>
+<TabItem value="requests" label="HTTP (requests/httpx)">
+
+```python
+import httpx
+
+headers = get_litellm_headers(request)
+headers["Authorization"] = "Bearer sk-your-litellm-key"
+
+response = httpx.post(
+    "http://localhost:4000/v1/chat/completions",
+    headers=headers,
+    json={"model": "gpt-4o", "messages": [{"role": "user", "content": "Hello"}]}
+)
+```
+</TabItem>
+</Tabs>
+
+### Result
+
+With header forwarding enabled, you'll see:
+
+**Trace Grouping in Langfuse:**
+
+<Image
+  img={require('../img/a2a_trace_grouping.png')}
+  style={{width: '80%', display: 'block', margin: '0', borderRadius: '8px'}}
+/>
+
+**Agent Spend Attribution:**
+
+<Image
+  img={require('../img/a2a_agent_spend.png')}
+  style={{width: '80%', display: 'block', margin: '0', borderRadius: '8px'}}
+/>
+
 ## API Reference
 
 ### Endpoint

litellm/a2a_protocol/main.py

@@ -9,7 +9,7 @@ import datetime
 from typing import TYPE_CHECKING, Any, AsyncIterator, Coroutine, Dict, Optional, Union
 
 import litellm
-from litellm._logging import verbose_logger
+from litellm._logging import verbose_logger, verbose_proxy_logger
 from litellm.a2a_protocol.streaming_iterator import A2AStreamingIterator
 from litellm.a2a_protocol.utils import A2ARequestUtils
 from litellm.constants import DEFAULT_A2A_AGENT_TIMEOUT
@@ -20,6 +20,7 @@ from litellm.llms.custom_httpx.http_handler import (
 )
 from litellm.types.agents import LiteLLMSendMessageResponse
 from litellm.utils import client
+import uuid
 
 if TYPE_CHECKING:
     from a2a.client import A2AClient as A2AClientType
@@ -225,7 +226,11 @@ async def asend_message(
             raise ValueError(
                 "Either a2a_client or api_base is required for standard A2A flow"
             )
-        a2a_client = await create_a2a_client(base_url=api_base)
+        trace_id = str(uuid.uuid4())
+        extra_headers = {"X-LiteLLM-Trace-Id": trace_id}
+        if agent_id:
+            extra_headers["X-LiteLLM-Agent-Id"] = agent_id
+        a2a_client = await create_a2a_client(base_url=api_base, extra_headers=extra_headers)
 
     # Type assertion: a2a_client is guaranteed to be non-None here
     assert a2a_client is not None
@@ -490,6 +495,10 @@ async def create_a2a_client(
     )
     httpx_client = http_handler.client
 
+    if extra_headers:
+        httpx_client.headers.update(extra_headers)
+        verbose_proxy_logger.debug(f"A2A client created with extra_headers={extra_headers}")
+
     # Resolve agent card
     resolver = A2ACardResolver(
         httpx_client=httpx_client,

litellm/proxy/_types.py

@@ -3421,6 +3421,8 @@ class LitellmMetadataFromRequestHeaders(TypedDict, total=False):
     """
 
     spend_logs_metadata: Optional[dict]
+    agent_id: Optional[str]
+    trace_id: Optional[str]
 
 
 class JWTKeyItem(TypedDict, total=False):

litellm/proxy/litellm_pre_call_utils.py

@@ -558,6 +558,16 @@ class LiteLLMProxyRequestSetup:
         #########################################################################################
         # Finally update the requests metadata with the `metadata_from_headers`
         #########################################################################################
+        agent_id_from_header = headers.get("x-litellm-agent-id")
+        trace_id_from_header = headers.get("x-litellm-trace-id")
+        if agent_id_from_header:
+            metadata_from_headers["agent_id"] = agent_id_from_header
+            verbose_proxy_logger.debug(f"Extracted agent_id from header: {agent_id_from_header}")
+        
+        if trace_id_from_header:
+            metadata_from_headers["trace_id"] = trace_id_from_header
+            verbose_proxy_logger.debug(f"Extracted trace_id from header: {trace_id_from_header}")
+
         if isinstance(data[_metadata_variable_name], dict):
             data[_metadata_variable_name].update(metadata_from_headers)
         return data

litellm/proxy/spend_tracking/spend_tracking_utils.py

@@ -396,7 +396,7 @@ def get_logging_payload(  # noqa: PLR0915
         )
 
     # Extract agent_id for A2A requests (set directly on model_call_details)
-    agent_id: Optional[str] = kwargs.get("agent_id")
+    agent_id: Optional[str] = kwargs.get("agent_id") or metadata.get("agent_id")
     custom_llm_provider = kwargs.get("custom_llm_provider")
     raw_model = cast(str, kwargs.get("model") or "")
     model_name = reconstruct_model_name(raw_model, custom_llm_provider, metadata or {})