From 028e6871dd5f8611f84c1e2dc853f44e506e5a92 Mon Sep 17 00:00:00 2001 From: Sameer Kankute Date: Thu, 5 Mar 2026 14:27:51 +0530 Subject: [PATCH 01/11] feat(agents): add static_headers and extra_headers fields to schema and types Add two new fields to LiteLLM_AgentsTable: - static_headers (Json): admin-configured headers always sent to the backend agent - extra_headers (String[]): header names to extract from the client request and forward Extend AgentConfig, PatchAgentRequest, and AgentResponse with the same fields. Also remove duplicate spec_path field from LiteLLM_MCPServerTable. Co-Authored-By: Claude Sonnet 4.6 --- litellm/proxy/schema.prisma | 3 ++- litellm/types/agents.py | 6 ++++++ 2 files changed, 8 insertions(+), 1 deletion(-) diff --git a/litellm/proxy/schema.prisma b/litellm/proxy/schema.prisma index 43972724ecc..6f4ef0c24b6 100644 --- a/litellm/proxy/schema.prisma +++ b/litellm/proxy/schema.prisma @@ -63,6 +63,8 @@ model LiteLLM_AgentsTable { agent_name String @unique litellm_params Json? agent_card_params Json + static_headers Json? @default("{}") + extra_headers String[] @default([]) agent_access_groups String[] @default([]) object_permission_id String? object_permission LiteLLM_ObjectPermissionTable? @relation(fields: [object_permission_id], references: [object_permission_id]) @@ -305,7 +307,6 @@ model LiteLLM_MCPServerTable { registration_url String? allow_all_keys Boolean @default(false) available_on_public_internet Boolean @default(true) - spec_path String? is_byok Boolean @default(false) byok_description String[] @default([]) byok_api_key_help_url String? diff --git a/litellm/types/agents.py b/litellm/types/agents.py index 3ad898b1935..7879cae9ff6 100644 --- a/litellm/types/agents.py +++ b/litellm/types/agents.py @@ -179,6 +179,8 @@ class AgentConfig(TypedDict, total=False): agent_card_params: Required[AgentCard] litellm_params: Dict[str, Any] # allow for any future litellm params object_permission: AgentObjectPermission + static_headers: Optional[Dict[str, str]] + extra_headers: Optional[List[str]] class PatchAgentRequest(TypedDict, total=False): @@ -186,6 +188,8 @@ class PatchAgentRequest(TypedDict, total=False): agent_card_params: AgentCard litellm_params: Dict[str, Any] object_permission: AgentObjectPermission + static_headers: Optional[Dict[str, str]] + extra_headers: Optional[List[str]] # Request/Response models for CRUD endpoints @@ -197,6 +201,8 @@ class AgentResponse(BaseModel): litellm_params: Optional[Dict[str, Any]] = None agent_card_params: Dict[str, Any] object_permission: Optional[Dict[str, Any]] = None + static_headers: Optional[Dict[str, str]] = None + extra_headers: Optional[List[str]] = None created_at: Optional[datetime] = None updated_at: Optional[datetime] = None created_by: Optional[str] = None From 07ee1e9886f54b773f4d9de7e3c8181e90d30d6e Mon Sep 17 00:00:00 2001 From: Sameer Kankute Date: Thu, 5 Mar 2026 14:28:01 +0530 Subject: [PATCH 02/11] feat(agents): persist static_headers and extra_headers in agent registry Update add_agent_to_db, patch_agent_in_db, and update_agent_in_db to read and write the two new header fields when creating or updating agents. Co-Authored-By: Claude Sonnet 4.6 --- .../proxy/agent_endpoints/agent_registry.py | 31 +++++++++++++++++++ 1 file changed, 31 insertions(+) diff --git a/litellm/proxy/agent_endpoints/agent_registry.py b/litellm/proxy/agent_endpoints/agent_registry.py index 159c9fb93d9..550182f966f 100644 --- a/litellm/proxy/agent_endpoints/agent_registry.py +++ b/litellm/proxy/agent_endpoints/agent_registry.py @@ -128,6 +128,14 @@ class AgentRegistry: agent_copy, None, prisma_client ) + # Serialize static_headers + static_headers_obj = agent.get("static_headers") + static_headers_val: Optional[str] = ( + safe_dumps(dict(static_headers_obj)) if static_headers_obj else None + ) + + extra_headers_val: Optional[List[str]] = agent.get("extra_headers") + create_data: Dict[str, Any] = { "agent_name": agent_name, "litellm_params": litellm_params, @@ -137,6 +145,10 @@ class AgentRegistry: "created_at": datetime.now(timezone.utc), "updated_at": datetime.now(timezone.utc), } + if static_headers_val is not None: + create_data["static_headers"] = static_headers_val + if extra_headers_val is not None: + create_data["extra_headers"] = extra_headers_val if object_permission_id is not None: create_data["object_permission_id"] = object_permission_id @@ -214,6 +226,12 @@ class AgentRegistry: update_data["agent_card_params"] = safe_dumps( augment_agent.get("agent_card_params") ) + if agent.get("static_headers") is not None: + update_data["static_headers"] = safe_dumps( + dict(agent.get("static_headers")) # type: ignore + ) + if agent.get("extra_headers") is not None: + update_data["extra_headers"] = agent.get("extra_headers") if agent.get("object_permission") is not None: agent_copy = dict(augment_agent) existing_object_permission_id = existing_agent.get( @@ -281,6 +299,15 @@ class AgentRegistry: ) agent_card_params: str = safe_dumps(agent_card_params_dict) + # Serialize static_headers for update + static_headers_obj_u = agent.get("static_headers") + static_headers_val_u: Optional[str] = ( + safe_dumps(dict(static_headers_obj_u)) + if static_headers_obj_u is not None + else None + ) + extra_headers_val_u: Optional[List[str]] = agent.get("extra_headers") + update_data: Dict[str, Any] = { "agent_name": agent_name, "litellm_params": litellm_params, @@ -288,6 +315,10 @@ class AgentRegistry: "updated_by": updated_by, "updated_at": datetime.now(timezone.utc), } + if static_headers_val_u is not None: + update_data["static_headers"] = static_headers_val_u + if extra_headers_val_u is not None: + update_data["extra_headers"] = extra_headers_val_u if agent.get("object_permission") is not None: existing_agent = await prisma_client.db.litellm_agentstable.find_unique( where={"agent_id": agent_id} From 16a30b55f5493bbf0754aac0dc4ea4c54b681804 Mon Sep 17 00:00:00 2001 From: Sameer Kankute Date: Thu, 5 Mar 2026 14:28:11 +0530 Subject: [PATCH 03/11] feat(agents): add merge_agent_headers utility Mirrors merge_mcp_headers from the MCP server utils. Dynamic headers come first; static (admin-configured) headers overlay and win on conflict. Co-Authored-By: Claude Sonnet 4.6 --- litellm/proxy/agent_endpoints/utils.py | 27 ++++++++++++++++++++++++++ 1 file changed, 27 insertions(+) create mode 100644 litellm/proxy/agent_endpoints/utils.py diff --git a/litellm/proxy/agent_endpoints/utils.py b/litellm/proxy/agent_endpoints/utils.py new file mode 100644 index 00000000000..2b968de54be --- /dev/null +++ b/litellm/proxy/agent_endpoints/utils.py @@ -0,0 +1,27 @@ +"""Utility helpers for A2A agent endpoints.""" + +from typing import Dict, Mapping, Optional + + +def merge_agent_headers( + *, + dynamic_headers: Optional[Mapping[str, str]] = None, + static_headers: Optional[Mapping[str, str]] = None, +) -> Optional[Dict[str, str]]: + """Merge outbound HTTP headers for A2A agent calls. + + Merge rules: + - Start with ``dynamic_headers`` (values extracted from the incoming client request). + - Overlay ``static_headers`` (admin-configured per agent). + + If both contain the same key, ``static_headers`` wins. + """ + merged: Dict[str, str] = {} + + if dynamic_headers: + merged.update({str(k): str(v) for k, v in dynamic_headers.items()}) + + if static_headers: + merged.update({str(k): str(v) for k, v in static_headers.items()}) + + return merged or None From 20a4eea27e71cfc5933670b73747fb46d66dd41d Mon Sep 17 00:00:00 2001 From: Sameer Kankute Date: Thu, 5 Mar 2026 14:28:28 +0530 Subject: [PATCH 04/11] feat(agents): forward custom headers to backend A2A agents In invoke_agent_a2a: - Extract admin-configured extra_headers from client request by name - Extract convention-based headers (x-a2a-{agent_id/name}-{header}) from client request - Merge with static_headers (static wins on conflict) - Pass merged headers down to asend_message and _handle_stream_message In asend_message / asend_message_streaming: - Accept agent_extra_headers kwarg - Overlay onto LiteLLM internal headers before creating the httpx client Co-Authored-By: Claude Sonnet 4.6 --- litellm/a2a_protocol/main.py | 14 ++++++- .../proxy/agent_endpoints/a2a_endpoints.py | 37 ++++++++++++++++++- 2 files changed, 48 insertions(+), 3 deletions(-) diff --git a/litellm/a2a_protocol/main.py b/litellm/a2a_protocol/main.py index 485b57e311b..6ac88d3a430 100644 --- a/litellm/a2a_protocol/main.py +++ b/litellm/a2a_protocol/main.py @@ -169,6 +169,7 @@ async def asend_message( api_base: Optional[str] = None, litellm_params: Optional[Dict[str, Any]] = None, agent_id: Optional[str] = None, + agent_extra_headers: Optional[Dict[str, str]] = None, **kwargs: Any, ) -> LiteLLMSendMessageResponse: """ @@ -250,9 +251,12 @@ async def asend_message( "Either a2a_client or api_base is required for standard A2A flow" ) trace_id = trace_id or str(uuid.uuid4()) - extra_headers = {"X-LiteLLM-Trace-Id": trace_id} + extra_headers: Dict[str, str] = {"X-LiteLLM-Trace-Id": trace_id} if agent_id: extra_headers["X-LiteLLM-Agent-Id"] = agent_id + # Overlay agent-level headers (agent headers take precedence over LiteLLM internal ones) + if agent_extra_headers: + extra_headers.update(agent_extra_headers) a2a_client = await create_a2a_client( base_url=api_base, extra_headers=extra_headers ) @@ -426,6 +430,7 @@ async def asend_message_streaming( agent_id: Optional[str] = None, metadata: Optional[Dict[str, Any]] = None, proxy_server_request: Optional[Dict[str, Any]] = None, + agent_extra_headers: Optional[Dict[str, str]] = None, ) -> AsyncIterator[Any]: """ Async: Send a streaming message to an A2A agent. @@ -507,7 +512,12 @@ async def asend_message_streaming( raise ValueError( "Either a2a_client or api_base is required for standard A2A flow" ) - a2a_client = await create_a2a_client(base_url=api_base) + streaming_extra_headers: Optional[Dict[str, str]] = None + if agent_extra_headers: + streaming_extra_headers = dict(agent_extra_headers) + a2a_client = await create_a2a_client( + base_url=api_base, extra_headers=streaming_extra_headers + ) # Type assertion: a2a_client is guaranteed to be non-None here assert a2a_client is not None diff --git a/litellm/proxy/agent_endpoints/a2a_endpoints.py b/litellm/proxy/agent_endpoints/a2a_endpoints.py index 6bcee14f29e..344070d17fc 100644 --- a/litellm/proxy/agent_endpoints/a2a_endpoints.py +++ b/litellm/proxy/agent_endpoints/a2a_endpoints.py @@ -6,13 +6,14 @@ The A2A SDK can point to LiteLLM's URL and invoke agents registered with LiteLLM """ import json -from typing import Any, Optional +from typing import Any, Dict, Optional from fastapi import APIRouter, Depends, HTTPException, Request, Response from fastapi.responses import JSONResponse, StreamingResponse from litellm._logging import verbose_proxy_logger from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy.agent_endpoints.utils import merge_agent_headers from litellm.proxy.auth.user_api_key_auth import user_api_key_auth from litellm.types.utils import all_litellm_params @@ -55,6 +56,7 @@ async def _handle_stream_message( metadata: Optional[dict] = None, proxy_server_request: Optional[dict] = None, *, + agent_extra_headers: Optional[Dict[str, str]] = None, user_api_key_dict: Optional[UserAPIKeyAuth] = None, request_data: Optional[dict] = None, proxy_logging_obj: Optional[Any] = None, @@ -105,6 +107,7 @@ async def _handle_stream_message( agent_id=agent_id, metadata=metadata, proxy_server_request=proxy_server_request, + agent_extra_headers=agent_extra_headers, ) if ( @@ -385,6 +388,36 @@ async def invoke_agent_a2a( version=version, ) + # Build merged headers for the backend agent + static_headers: Dict[str, str] = dict(agent.static_headers or {}) + + raw_headers = dict(request.headers) + normalized = {k.lower(): v for k, v in raw_headers.items()} + + dynamic_headers: Dict[str, str] = {} + + # 1. Admin-configured extra_headers: forward named headers from client request + if agent.extra_headers: + for header_name in agent.extra_headers: + val = normalized.get(header_name.lower()) + if val is not None: + dynamic_headers[header_name] = val + + # 2. Convention-based forwarding: x-a2a-{agent_id_or_name}-{header_name} + # Matches both agent_id (UUID) and agent_name (alias), case-insensitive. + for alias in (agent.agent_id.lower(), agent.agent_name.lower()): + prefix = f"x-a2a-{alias}-" + for key, val in normalized.items(): + if key.startswith(prefix): + header_name = key[len(prefix) :] + if header_name: + dynamic_headers[header_name] = val + + agent_extra_headers = merge_agent_headers( + dynamic_headers=dynamic_headers or None, + static_headers=static_headers or None, + ) + # Route through SDK functions if method == "message/send": from a2a.types import MessageSendParams, SendMessageRequest @@ -401,6 +434,7 @@ async def invoke_agent_a2a( metadata=data.get("metadata", {}), proxy_server_request=data.get("proxy_server_request"), litellm_logging_obj=logging_obj, + agent_extra_headers=agent_extra_headers, ) response = await proxy_logging_obj.post_call_success_hook( @@ -425,6 +459,7 @@ async def invoke_agent_a2a( agent_id=agent.agent_id, metadata=data.get("metadata", {}), proxy_server_request=data.get("proxy_server_request"), + agent_extra_headers=agent_extra_headers, user_api_key_dict=user_api_key_dict, request_data=data, proxy_logging_obj=proxy_logging_obj, From 6e9c7c4a8dd8ddce1b911d77e2009aac3de5f9d3 Mon Sep 17 00:00:00 2001 From: Sameer Kankute Date: Thu, 5 Mar 2026 14:28:36 +0530 Subject: [PATCH 05/11] feat(agents): add Prisma migration for agent header columns ALTER TABLE LiteLLM_AgentsTable to add: - static_headers JSONB DEFAULT '{}' - extra_headers TEXT[] DEFAULT ARRAY[]::TEXT[] Co-Authored-By: Claude Sonnet 4.6 --- .../20260305000000_add_agent_headers/migration.sql | 5 +++++ 1 file changed, 5 insertions(+) create mode 100644 litellm-proxy-extras/litellm_proxy_extras/migrations/20260305000000_add_agent_headers/migration.sql diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260305000000_add_agent_headers/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260305000000_add_agent_headers/migration.sql new file mode 100644 index 00000000000..acb35baba96 --- /dev/null +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260305000000_add_agent_headers/migration.sql @@ -0,0 +1,5 @@ +-- Add static_headers and extra_headers to LiteLLM_AgentsTable + +ALTER TABLE "LiteLLM_AgentsTable" + ADD COLUMN IF NOT EXISTS "static_headers" JSONB DEFAULT '{}', + ADD COLUMN IF NOT EXISTS "extra_headers" TEXT[] DEFAULT ARRAY[]::TEXT[]; From fd53678898b71f6da4384e5984a4d1308f2ee060 Mon Sep 17 00:00:00 2001 From: Sameer Kankute Date: Thu, 5 Mar 2026 14:28:48 +0530 Subject: [PATCH 06/11] test(agents): add tests for A2A custom header forwarding Covers: - Static headers forwarded to backend - Dynamic headers extracted by name (extra_headers config) - Convention-based x-a2a-{agent_id/name}-{header} forwarding - Static headers win over dynamic on conflict - Unrelated x-a2a- prefixes are not forwarded - No-header case leaves existing behaviour unchanged - merge_agent_headers utility unit tests Co-Authored-By: Claude Sonnet 4.6 --- .../agent_endpoints/test_agent_headers.py | 339 ++++++++++++++++++ 1 file changed, 339 insertions(+) create mode 100644 tests/test_litellm/proxy/agent_endpoints/test_agent_headers.py diff --git a/tests/test_litellm/proxy/agent_endpoints/test_agent_headers.py b/tests/test_litellm/proxy/agent_endpoints/test_agent_headers.py new file mode 100644 index 00000000000..b52c0afb0c0 --- /dev/null +++ b/tests/test_litellm/proxy/agent_endpoints/test_agent_headers.py @@ -0,0 +1,339 @@ +""" +Unit tests for A2A agent custom header forwarding. + +Tests cover: +- Static headers forwarded to backend agent +- Dynamic headers extracted from client request and forwarded +- Static headers win over dynamic on conflict +- No headers configured — existing behavior unchanged +- merge_agent_headers utility +""" + +import sys +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest + + +# --------------------------------------------------------------------------- +# Helper: build a minimal mock agent +# --------------------------------------------------------------------------- + +def _make_mock_agent( + static_headers=None, + extra_headers=None, + url="http://backend-agent:10001", +): + mock_agent = MagicMock() + mock_agent.agent_id = "agent-123" + mock_agent.agent_card_params = {"url": url, "name": "Test Agent"} + mock_agent.litellm_params = {} + mock_agent.static_headers = static_headers or {} + mock_agent.extra_headers = extra_headers or [] + return mock_agent + + +def _make_mock_request(extra_headers=None, method="message/send"): + """Build a mock FastAPI Request with configurable headers.""" + mock_request = MagicMock() + headers = {"content-type": "application/json"} + if extra_headers: + headers.update(extra_headers) + mock_request.headers = headers + mock_request.json = AsyncMock( + return_value={ + "jsonrpc": "2.0", + "id": "test-id", + "method": method, + "params": { + "message": { + "role": "user", + "parts": [{"kind": "text", "text": "Hello"}], + "messageId": "msg-123", + } + }, + } + ) + return mock_request + + +def _make_a2a_types_module(): + """Return (module, MessageSendParams, SendMessageRequest, SendStreamingMessageRequest).""" + try: + from a2a.types import ( + MessageSendParams, + SendMessageRequest, + SendStreamingMessageRequest, + ) + mock_a2a_types = MagicMock() + mock_a2a_types.MessageSendParams = MessageSendParams + mock_a2a_types.SendMessageRequest = SendMessageRequest + mock_a2a_types.SendStreamingMessageRequest = SendStreamingMessageRequest + return mock_a2a_types + except ImportError: + pass + + def _make_cls(name): + class MockCls: + def __init__(self, **kwargs): + self.__dict__.update(kwargs) + self._kwargs = kwargs + + def model_dump(self, mode="json", exclude_none=False): + result = dict(self._kwargs) + if exclude_none: + result = {k: v for k, v in result.items() if v is not None} + return result + + MockCls.__name__ = name + return MockCls + + mock_a2a_types = MagicMock() + mock_a2a_types.MessageSendParams = _make_cls("MessageSendParams") + mock_a2a_types.SendMessageRequest = _make_cls("SendMessageRequest") + mock_a2a_types.SendStreamingMessageRequest = _make_cls( + "SendStreamingMessageRequest" + ) + return mock_a2a_types + + +async def _invoke(mock_agent, mock_request, mock_asend_message): + """Run invoke_agent_a2a with standard patches applied.""" + from litellm.proxy._types import UserAPIKeyAuth + + mock_user_api_key_dict = UserAPIKeyAuth(api_key="sk-test", user_id="u1") + mock_fastapi_response = MagicMock() + mock_a2a_types = _make_a2a_types_module() + + mock_response = MagicMock() + mock_response.model_dump.return_value = { + "jsonrpc": "2.0", + "id": "test-id", + "result": {"status": "success"}, + } + + with patch( + "litellm.proxy.agent_endpoints.a2a_endpoints._get_agent", + return_value=mock_agent, + ), patch( + "litellm.proxy.agent_endpoints.auth.agent_permission_handler.AgentRequestHandler.is_agent_allowed", + new_callable=AsyncMock, + return_value=True, + ), patch( + "litellm.proxy.common_request_processing.add_litellm_data_to_request", + side_effect=lambda data, **kw: data, + ), patch( + "litellm.a2a_protocol.asend_message", + new_callable=AsyncMock, + return_value=mock_response, + ) as mock_asend, patch( + "litellm.a2a_protocol.create_a2a_client", + new_callable=AsyncMock, + ), patch( + "litellm.proxy.proxy_server.general_settings", + {}, + ), patch( + "litellm.proxy.proxy_server.proxy_config", + MagicMock(), + ), patch( + "litellm.proxy.proxy_server.version", + "1.0.0", + ), patch.dict( + sys.modules, + {"a2a": MagicMock(), "a2a.types": mock_a2a_types}, + ), patch( + "litellm.a2a_protocol.main.A2A_SDK_AVAILABLE", + True, + ): + from litellm.proxy.agent_endpoints.a2a_endpoints import invoke_agent_a2a + + await invoke_agent_a2a( + agent_id="test-agent", + request=mock_request, + fastapi_response=mock_fastapi_response, + user_api_key_dict=mock_user_api_key_dict, + ) + return mock_asend + + +# --------------------------------------------------------------------------- +# Tests +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_static_headers_forwarded(): + """Static headers configured on the agent are passed to asend_message.""" + mock_agent = _make_mock_agent( + static_headers={"Authorization": "Bearer token123"} + ) + mock_request = _make_mock_request() + + mock_asend = await _invoke(mock_agent, mock_request, None) + + call_kwargs = mock_asend.call_args.kwargs + headers = call_kwargs.get("agent_extra_headers") + assert headers is not None, "agent_extra_headers should not be None" + assert headers.get("Authorization") == "Bearer token123" + + +@pytest.mark.asyncio +async def test_dynamic_headers_forwarded(): + """Dynamic headers listed in extra_headers are extracted from the client request.""" + mock_agent = _make_mock_agent(extra_headers=["x-api-key"]) + mock_request = _make_mock_request(extra_headers={"x-api-key": "secret"}) + + mock_asend = await _invoke(mock_agent, mock_request, None) + + call_kwargs = mock_asend.call_args.kwargs + headers = call_kwargs.get("agent_extra_headers") + assert headers is not None + assert headers.get("x-api-key") == "secret" + + +@pytest.mark.asyncio +async def test_static_overrides_dynamic(): + """When the same header appears in both static and dynamic, static wins.""" + mock_agent = _make_mock_agent( + static_headers={"Authorization": "Bearer static-token"}, + extra_headers=["Authorization"], + ) + # Client sends a different value for Authorization + mock_request = _make_mock_request( + extra_headers={"Authorization": "Bearer dynamic-token"} + ) + + mock_asend = await _invoke(mock_agent, mock_request, None) + + call_kwargs = mock_asend.call_args.kwargs + headers = call_kwargs.get("agent_extra_headers") + assert headers is not None + assert headers.get("Authorization") == "Bearer static-token" + + +@pytest.mark.asyncio +async def test_no_headers(): + """When no headers are configured, agent_extra_headers is None and behaviour is unchanged.""" + mock_agent = _make_mock_agent() # no static_headers or extra_headers + mock_request = _make_mock_request() + + mock_asend = await _invoke(mock_agent, mock_request, None) + + call_kwargs = mock_asend.call_args.kwargs + headers = call_kwargs.get("agent_extra_headers") + assert headers is None + + +# --------------------------------------------------------------------------- +# Convention-based x-a2a-{agent_id/name}-{header_name} tests +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_convention_header_by_agent_name(): + """x-a2a-{agent_name}-{header} is forwarded using the agent name alias.""" + mock_agent = _make_mock_agent() + mock_agent.agent_name = "my-agent" + mock_request = _make_mock_request( + extra_headers={"x-a2a-my-agent-authorization": "Bearer conv-token"} + ) + + mock_asend = await _invoke(mock_agent, mock_request, None) + + headers = mock_asend.call_args.kwargs.get("agent_extra_headers") + assert headers is not None + assert headers.get("authorization") == "Bearer conv-token" + + +@pytest.mark.asyncio +async def test_convention_header_by_agent_id(): + """x-a2a-{agent_id}-{header} is forwarded using the agent UUID.""" + mock_agent = _make_mock_agent() + mock_agent.agent_id = "abc-123" + mock_agent.agent_name = "other-name" + mock_request = _make_mock_request( + extra_headers={"x-a2a-abc-123-x-api-key": "id-secret"} + ) + + mock_asend = await _invoke(mock_agent, mock_request, None) + + headers = mock_asend.call_args.kwargs.get("agent_extra_headers") + assert headers is not None + assert headers.get("x-api-key") == "id-secret" + + +@pytest.mark.asyncio +async def test_convention_header_static_still_wins(): + """Static headers still override convention-based dynamic headers.""" + mock_agent = _make_mock_agent( + static_headers={"authorization": "Bearer static-wins"} + ) + mock_agent.agent_name = "my-agent" + mock_request = _make_mock_request( + extra_headers={"x-a2a-my-agent-authorization": "Bearer conv-value"} + ) + + mock_asend = await _invoke(mock_agent, mock_request, None) + + headers = mock_asend.call_args.kwargs.get("agent_extra_headers") + assert headers is not None + assert headers.get("authorization") == "Bearer static-wins" + + +@pytest.mark.asyncio +async def test_convention_unrelated_prefix_not_forwarded(): + """Headers that start with x-a2a- but target a different agent are ignored.""" + mock_agent = _make_mock_agent() + mock_agent.agent_id = "agent-abc" + mock_agent.agent_name = "my-agent" + mock_request = _make_mock_request( + extra_headers={"x-a2a-other-agent-authorization": "Bearer wrong"} + ) + + mock_asend = await _invoke(mock_agent, mock_request, None) + + headers = mock_asend.call_args.kwargs.get("agent_extra_headers") + assert headers is None + + +# --------------------------------------------------------------------------- +# Direct unit test for the merge utility +# --------------------------------------------------------------------------- + + +def test_merge_agent_headers_util_dynamic_only(): + from litellm.proxy.agent_endpoints.utils import merge_agent_headers + + result = merge_agent_headers(dynamic_headers={"x-key": "val"}) + assert result == {"x-key": "val"} + + +def test_merge_agent_headers_util_static_only(): + from litellm.proxy.agent_endpoints.utils import merge_agent_headers + + result = merge_agent_headers(static_headers={"Authorization": "Bearer tok"}) + assert result == {"Authorization": "Bearer tok"} + + +def test_merge_agent_headers_util_static_wins(): + from litellm.proxy.agent_endpoints.utils import merge_agent_headers + + result = merge_agent_headers( + dynamic_headers={"Authorization": "dynamic", "x-extra": "d"}, + static_headers={"Authorization": "static"}, + ) + assert result == {"Authorization": "static", "x-extra": "d"} + + +def test_merge_agent_headers_util_none_returns_none(): + from litellm.proxy.agent_endpoints.utils import merge_agent_headers + + result = merge_agent_headers() + assert result is None + + +def test_merge_agent_headers_util_empty_dicts_returns_none(): + from litellm.proxy.agent_endpoints.utils import merge_agent_headers + + result = merge_agent_headers(dynamic_headers={}, static_headers={}) + assert result is None From 36d279ab42c20185d435d502f18f895487249ab3 Mon Sep 17 00:00:00 2001 From: Sameer Kankute Date: Thu, 5 Mar 2026 14:34:11 +0530 Subject: [PATCH 07/11] feat(ui/agents): add Authentication Headers section to agent create/edit form Add a new "Authentication Headers" panel to AgentFormFields: - Static Headers: key-value Form.List (always sent to the backend agent, static wins on conflict with dynamic) - Forward Client Headers: Select[tags] of header names to extract from the client request and forward (extra_headers) Update buildAgentDataFromForm to serialize both fields for the API. Update parseAgentForForm to deserialize them back for editing. Covers both the create wizard (add_agent_form) and the edit view (agent_info). Co-Authored-By: Claude Sonnet 4.6 --- .../src/components/agents/agent_config.ts | 26 +++++++ .../components/agents/agent_form_fields.tsx | 70 ++++++++++++++++++- 2 files changed, 94 insertions(+), 2 deletions(-) diff --git a/ui/litellm-dashboard/src/components/agents/agent_config.ts b/ui/litellm-dashboard/src/components/agents/agent_config.ts index f85c4daac66..01041c5cee4 100644 --- a/ui/litellm-dashboard/src/components/agents/agent_config.ts +++ b/ui/litellm-dashboard/src/components/agents/agent_config.ts @@ -269,6 +269,23 @@ export const buildAgentDataFromForm = (values: any, existingAgent?: any) => { agentData.litellm_params = params; } + // static_headers: convert [{header, value}, ...] → {header: value, ...} + if (Array.isArray(values.static_headers) && values.static_headers.length > 0) { + const staticHeaders: Record = {}; + values.static_headers.forEach((entry: { header?: string; value?: string }) => { + const key = entry?.header?.trim(); + if (key) staticHeaders[key] = entry?.value ?? ""; + }); + if (Object.keys(staticHeaders).length > 0) { + agentData.static_headers = staticHeaders; + } + } + + // extra_headers: already an array of strings from Select tags + if (Array.isArray(values.extra_headers) && values.extra_headers.length > 0) { + agentData.extra_headers = values.extra_headers; + } + return agentData; }; @@ -302,5 +319,14 @@ export const parseAgentForForm = (agent: any) => { cost_per_query: agent.litellm_params?.cost_per_query, input_cost_per_token: agent.litellm_params?.input_cost_per_token, output_cost_per_token: agent.litellm_params?.output_cost_per_token, + // static_headers: {key: value} → [{header, value}, ...] + static_headers: agent.static_headers + ? Object.entries(agent.static_headers as Record).map(([header, value]) => ({ + header, + value, + })) + : [], + // extra_headers: already an array of strings + extra_headers: agent.extra_headers ?? [], }; }; diff --git a/ui/litellm-dashboard/src/components/agents/agent_form_fields.tsx b/ui/litellm-dashboard/src/components/agents/agent_form_fields.tsx index d5429d2a3b5..42e55b8c56f 100644 --- a/ui/litellm-dashboard/src/components/agents/agent_form_fields.tsx +++ b/ui/litellm-dashboard/src/components/agents/agent_form_fields.tsx @@ -1,7 +1,7 @@ import React from "react"; -import { Form, Input, Switch, Collapse } from "antd"; +import { Form, Input, Switch, Collapse, Select, Space, Tooltip } from "antd"; import { Button as AntButton } from "antd"; -import { PlusOutlined, MinusCircleOutlined } from "@ant-design/icons"; +import { PlusOutlined, MinusCircleOutlined, InfoCircleOutlined } from "@ant-design/icons"; import { AGENT_FORM_CONFIG, SKILL_FIELD_CONFIG } from "./agent_config"; import CostConfigFields from "./cost_config_fields"; @@ -188,6 +188,72 @@ const AgentFormFields: React.FC = ({ showAgentName = true, ))} )} + + {/* Authentication Headers */} + {shouldShow("auth_headers") && ( + + {/* Static Headers */} + + Static Headers{" "} + + + + + } + > + + {(fields, { add, remove }) => ( + <> + {fields.map(({ key, name, ...restField }) => ( + + + + + + + + remove(name)} style={{ color: "#ff4d4f" }} /> + + ))} + add()} icon={} style={{ width: "100%" }}> + Add Static Header + + + )} + + + + {/* Extra Headers (dynamic forwarding) */} + + Forward Client Headers{" "} + + + + + } + name="extra_headers" + > +