feat: add header forwarding logic

Author: prassanna-ravishankar

src/agentex/lib/adk/_modules/acp.py

@@ -59,6 +59,7 @@ async def create_task(
         start_to_close_timeout: timedelta = timedelta(seconds=5),
         heartbeat_timeout: timedelta = timedelta(seconds=5),
         retry_policy: RetryPolicy = DEFAULT_RETRY_POLICY,
+        extra_headers: dict[str, str] | None = None,
     ) -> Task:
         """
         Create a new task.
@@ -71,6 +72,7 @@ async def create_task(
             start_to_close_timeout: The start to close timeout for the task.
             heartbeat_timeout: The heartbeat timeout for the task.
             retry_policy: The retry policy for the task.
+            extra_headers: Additional HTTP headers to forward to the agent (filtered by agent's allowlist).
 
         Returns:
             The task entry.
@@ -85,6 +87,7 @@ async def create_task(
                     params=params,
                     trace_id=trace_id,
                     parent_span_id=parent_span_id,
+                    extra_headers=extra_headers,
                 ),
                 response_type=Task,
                 start_to_close_timeout=start_to_close_timeout,
@@ -99,6 +102,7 @@ async def create_task(
                 params=params,
                 trace_id=trace_id,
                 parent_span_id=parent_span_id,
+                extra_headers=extra_headers,
             )
 
     async def send_event(
@@ -112,15 +116,22 @@ async def send_event(
         start_to_close_timeout: timedelta = timedelta(seconds=5),
         heartbeat_timeout: timedelta = timedelta(seconds=5),
         retry_policy: RetryPolicy = DEFAULT_RETRY_POLICY,
+        extra_headers: dict[str, str] | None = None,
     ) -> Event:
         """
         Send an event to a task.
 
         Args:
             task_id: The ID of the task to send the event to.
-            data: The data to send to the event.
+            content: The content to send to the event.
             agent_id: The ID of the agent to send the event to.
             agent_name: The name of the agent to send the event to.
+            trace_id: The trace ID for the event.
+            parent_span_id: The parent span ID for the event.
+            start_to_close_timeout: The start to close timeout for the event.
+            heartbeat_timeout: The heartbeat timeout for the event.
+            retry_policy: The retry policy for the event.
+            extra_headers: Additional HTTP headers to forward to the agent (filtered by agent's allowlist).
 
         Returns:
             The event entry.
@@ -135,6 +146,7 @@ async def send_event(
                     content=content,
                     trace_id=trace_id,
                     parent_span_id=parent_span_id,
+                    extra_headers=extra_headers,
                 ),
                 response_type=None,
                 start_to_close_timeout=start_to_close_timeout,
@@ -149,6 +161,7 @@ async def send_event(
                 content=content,
                 trace_id=trace_id,
                 parent_span_id=parent_span_id,
+                extra_headers=extra_headers,
             )
 
     async def send_message(
@@ -162,15 +175,22 @@ async def send_message(
         start_to_close_timeout: timedelta = timedelta(seconds=5),
         heartbeat_timeout: timedelta = timedelta(seconds=5),
         retry_policy: RetryPolicy = DEFAULT_RETRY_POLICY,
+        extra_headers: dict[str, str] | None = None,
     ) -> List[TaskMessage]:
         """
         Send a message to a task.
 
         Args:
-            task_id: The ID of the task to send the message to.
             content: The task message content to send to the task.
+            task_id: The ID of the task to send the message to.
             agent_id: The ID of the agent to send the message to.
             agent_name: The name of the agent to send the message to.
+            trace_id: The trace ID for the message.
+            parent_span_id: The parent span ID for the message.
+            start_to_close_timeout: The start to close timeout for the message.
+            heartbeat_timeout: The heartbeat timeout for the message.
+            retry_policy: The retry policy for the message.
+            extra_headers: Additional HTTP headers to forward to the agent (filtered by agent's allowlist).
 
         Returns:
             The message entry.
@@ -185,6 +205,7 @@ async def send_message(
                     content=content,
                     trace_id=trace_id,
                     parent_span_id=parent_span_id,
+                    extra_headers=extra_headers,
                 ),
                 response_type=TaskMessage,
                 start_to_close_timeout=start_to_close_timeout,
@@ -199,6 +220,7 @@ async def send_message(
                 content=content,
                 trace_id=trace_id,
                 parent_span_id=parent_span_id,
+                extra_headers=extra_headers,
             )
 
     async def cancel_task(
@@ -210,6 +232,7 @@ async def cancel_task(
         start_to_close_timeout: timedelta = timedelta(seconds=5),
         heartbeat_timeout: timedelta = timedelta(seconds=5),
         retry_policy: RetryPolicy = DEFAULT_RETRY_POLICY,
+        extra_headers: dict[str, str] | None = None,
     ) -> Task:
         """
         Cancel a task.
@@ -222,6 +245,7 @@ async def cancel_task(
             start_to_close_timeout: The start to close timeout for the task.
             heartbeat_timeout: The heartbeat timeout for the task.
             retry_policy: The retry policy for the task.
+            extra_headers: Additional HTTP headers to forward to the agent (filtered by agent's allowlist).
 
         Returns:
             The task entry.
@@ -234,6 +258,7 @@ async def cancel_task(
                     task_name=task_name,
                     trace_id=trace_id,
                     parent_span_id=parent_span_id,
+                    extra_headers=extra_headers,
                 ),
                 response_type=None,
                 start_to_close_timeout=start_to_close_timeout,
@@ -246,4 +271,5 @@ async def cancel_task(
                 task_name=task_name,
                 trace_id=trace_id,
                 parent_span_id=parent_span_id,
+                extra_headers=extra_headers,
             )

src/agentex/lib/core/services/adk/acp/acp.py

@@ -1,9 +1,11 @@
 from typing import Any, List, cast
+import fnmatch
 
 from agentex import AsyncAgentex
 from agentex.lib.core.tracing.tracer import AsyncTracer
 from agentex.lib.utils.logging import make_logger
 from agentex.lib.utils.temporal import heartbeat_if_in_workflow
+from agentex.lib.types.agent_configs import CustomHeadersConfig
 from agentex.types.event import Event
 from agentex.types.task import Task
 from agentex.types.task_message import TaskMessage
@@ -22,6 +24,73 @@ def __init__(
         self._agentex_client = agentex_client
         self._tracer = tracer
 
+    async def _get_agent_header_config(
+        self, agent_name: str | None, agent_id: str | None
+    ) -> CustomHeadersConfig:
+        """
+        Get agent's header configuration from manifest.
+        
+        For now, we'll return a default empty config since we need to implement
+        manifest loading. This provides secure-by-default behavior.
+        
+        TODO: Implement actual manifest loading to get agent's custom_headers config
+        """
+        # Default configuration - no headers allowed (secure by default)
+        # This will be replaced with actual manifest loading
+        return CustomHeadersConfig(
+            strategy="allowlist",
+            allowed_headers=[],
+            max_header_size=8192,
+            max_headers_count=50
+        )
+
+    def _filter_headers(
+        self,
+        extra_headers: dict[str, str] | None,
+        config: CustomHeadersConfig
+    ) -> dict[str, str]:
+        """
+        Filter headers based on agent's configuration.
+        
+        Args:
+            extra_headers: Headers to filter
+            config: Agent's header configuration
+            
+        Returns:
+            Filtered headers dictionary containing only allowed headers
+        """
+        if not extra_headers or not config.allowed_headers:
+            return {}
+
+        filtered = {}
+        headers_added = 0
+
+        for header_name, header_value in extra_headers.items():
+            # Check if we've hit the max headers limit
+            if headers_added >= config.max_headers_count:
+                logger.warning("Reached maximum header count limit (%s), ignoring remaining headers", config.max_headers_count)
+                break
+
+            # Check against allowlist patterns
+            header_allowed = False
+            for pattern in config.allowed_headers:
+                if fnmatch.fnmatch(header_name.lower(), pattern.lower()):
+                    header_allowed = True
+                    break
+
+            if header_allowed:
+                # Apply size limits
+                if len(header_value) <= config.max_header_size:
+                    filtered[header_name] = header_value
+                    headers_added += 1
+                    logger.debug("Allowed header: %s", header_name)
+                else:
+                    logger.warning("Header '%s' exceeds size limit (%s > %s), ignoring", header_name, len(header_value), config.max_header_size)
+            else:
+                logger.debug("Header '%s' not in allowlist, ignoring", header_name)
+
+        return filtered
+
     async def task_create(
         self,
         name: str | None = None,
@@ -30,6 +99,7 @@ async def task_create(
         params: dict[str, Any] | None = None,
         trace_id: str | None = None,
         parent_span_id: str | None = None,
+        extra_headers: dict[str, str] | None = None,
     ) -> Task:
         trace = self._tracer.trace(trace_id=trace_id)
         async with trace.span(
@@ -43,6 +113,11 @@ async def task_create(
             },
         ) as span:
             heartbeat_if_in_workflow("task create")
+            
+            # Get agent's header configuration and filter headers
+            header_config = await self._get_agent_header_config(agent_name, agent_id)
+            filtered_headers = self._filter_headers(extra_headers, header_config)
+            
             if agent_name:
                 json_rpc_response = await self._agentex_client.agents.rpc_by_name(
                     agent_name=agent_name,
@@ -51,6 +126,7 @@ async def task_create(
                         "name": name,
                         "params": params,
                     },
+                    extra_headers=filtered_headers,
                 )
             elif agent_id:
                 json_rpc_response = await self._agentex_client.agents.rpc(
@@ -60,6 +136,7 @@ async def task_create(
                         "name": name,
                         "params": params,
                     },
+                    extra_headers=filtered_headers,
                 )
             else:
                 raise ValueError("Either agent_name or agent_id must be provided")
@@ -78,6 +155,7 @@ async def message_send(
         task_name: str | None = None,
         trace_id: str | None = None,
         parent_span_id: str | None = None,
+        extra_headers: dict[str, str] | None = None,
     ) -> List[TaskMessage]:
         trace = self._tracer.trace(trace_id=trace_id)
         async with trace.span(
@@ -92,6 +170,11 @@ async def message_send(
             },
         ) as span:
             heartbeat_if_in_workflow("message send")
+            
+            # Get agent's header configuration and filter headers
+            header_config = await self._get_agent_header_config(agent_name, agent_id)
+            filtered_headers = self._filter_headers(extra_headers, header_config)
+            
             if agent_name:
                 json_rpc_response = await self._agentex_client.agents.rpc_by_name(
                     agent_name=agent_name,
@@ -101,6 +184,7 @@ async def message_send(
                         "content": cast(TaskMessageContentParam, content.model_dump()),
                         "stream": False,
                     },
+                    extra_headers=filtered_headers,
                 )
             elif agent_id:
                 json_rpc_response = await self._agentex_client.agents.rpc(
@@ -111,12 +195,13 @@ async def message_send(
                         "content": cast(TaskMessageContentParam, content.model_dump()),
                         "stream": False,
                     },
+                    extra_headers=filtered_headers,
                 )
             else:
                 raise ValueError("Either agent_name or agent_id must be provided")
 
             task_messages: List[TaskMessage] = []
-            logger.info(f"json_rpc_response: {json_rpc_response}")
+            logger.info("json_rpc_response: %s", json_rpc_response)
             if isinstance(json_rpc_response.result, list):
                 for message in json_rpc_response.result:
                     task_message = TaskMessage.model_validate(message)
@@ -137,6 +222,7 @@ async def event_send(
         task_name: str | None = None,
         trace_id: str | None = None,
         parent_span_id: str | None = None,
+        extra_headers: dict[str, str] | None = None,
     ) -> Event:
         trace = self._tracer.trace(trace_id=trace_id)
         async with trace.span(
@@ -150,6 +236,11 @@ async def event_send(
             },
         ) as span:
             heartbeat_if_in_workflow("event send")
+            
+            # Get agent's header configuration and filter headers
+            header_config = await self._get_agent_header_config(agent_name, agent_id)
+            filtered_headers = self._filter_headers(extra_headers, header_config)
+            
             if agent_name:
                 json_rpc_response = await self._agentex_client.agents.rpc_by_name(
                     agent_name=agent_name,
@@ -158,6 +249,7 @@ async def event_send(
                         "task_id": task_id,
                         "content": cast(TaskMessageContentParam, content.model_dump()),
                     },
+                    extra_headers=filtered_headers,
                 )
             elif agent_id:
                 json_rpc_response = await self._agentex_client.agents.rpc(
@@ -167,6 +259,7 @@ async def event_send(
                         "task_id": task_id,
                         "content": cast(TaskMessageContentParam, content.model_dump()),
                     },
+                    extra_headers=filtered_headers,
                 )
             else:
                 raise ValueError("Either agent_name or agent_id must be provided")
@@ -182,6 +275,7 @@ async def task_cancel(
         task_name: str | None = None,
         trace_id: str | None = None,
         parent_span_id: str | None = None,
+        extra_headers: dict[str, str] | None = None,
     ) -> Task:
         trace = self._tracer.trace(trace_id=trace_id)
         async with trace.span(
@@ -193,13 +287,23 @@ async def task_cancel(
             },
         ) as span:
             heartbeat_if_in_workflow("task cancel")
+            
+            # Note: The original implementation seems to treat task_name as agent_name and task_id as agent_id
+            # This maintains backward compatibility while adding header support
+            # Get agent's header configuration and filter headers
+            agent_name = task_name if task_name else None
+            agent_id = task_id if task_id else None
+            header_config = await self._get_agent_header_config(agent_name, agent_id)
+            filtered_headers = self._filter_headers(extra_headers, header_config)
+            
             if task_name:
                 json_rpc_response = await self._agentex_client.agents.rpc_by_name(
                     agent_name=task_name,
                     method="task/cancel",
                     params={
                         "task_name": task_name,
                     },
+                    extra_headers=filtered_headers,
                 )
             elif task_id:
                 json_rpc_response = await self._agentex_client.agents.rpc(
@@ -208,6 +312,7 @@ async def task_cancel(
                     params={
                         "task_id": task_id,
                     },
+                    extra_headers=filtered_headers,
                 )
             else:
                 raise ValueError("Either task_name or task_id must be provided")