feat: add easy access to headers

Author: prassanna-ravishankar

src/agentex/lib/sdk/fastacp/base/base_acp_server.py

@@ -128,9 +128,17 @@ async def _handle_jsonrpc(self, request: Request):
                     ),
                 )
 
-            # Parse params into appropriate model based on method
+            # Extract custom headers (already filtered by AgentEx based on agent manifest)
+            custom_headers = {
+                key: value for key, value in request.headers.items()
+                if key.lower().startswith('x-')  # Custom headers typically start with 'x-'
+            }
+            
+            # Parse params into appropriate model based on method and include headers
             params_model = PARAMS_MODEL_BY_METHOD[method]
-            params = params_model.model_validate(rpc_request.params)
+            params_data = dict(rpc_request.params) if rpc_request.params else {}
+            params_data['extra_headers'] = custom_headers if custom_headers else None
+            params = params_model.model_validate(params_data)
 
             if method in RPC_SYNC_METHODS:
                 handler = self._handlers[method]

src/agentex/lib/types/acp.py

@@ -25,6 +25,7 @@ class CreateTaskParams(BaseModel):
         agent: The agent that the task was sent to.
         task: The task to be created.
         params: The parameters for the task as inputted by the user.
+        extra_headers: Custom headers forwarded to this agent (filtered by manifest allowlist).
     """
 
     agent: Agent = Field(..., description="The agent that the task was sent to")
@@ -33,6 +34,10 @@ class CreateTaskParams(BaseModel):
         None,
         description="The parameters for the task as inputted by the user",
     )
+    extra_headers: dict[str, str] | None = Field(
+        default=None,
+        description="Custom headers forwarded to this agent (filtered by manifest allowlist)",
+    )
 
 
 class SendMessageParams(BaseModel):
@@ -43,6 +48,7 @@ class SendMessageParams(BaseModel):
         task: The task that the message was sent to.
         content: The message that was sent to the agent.
         stream: Whether to stream the message back to the agentex server from the agent.
+        extra_headers: Custom headers forwarded to this agent (filtered by manifest allowlist).
     """
 
     agent: Agent = Field(..., description="The agent that the message was sent to")
@@ -54,6 +60,10 @@ class SendMessageParams(BaseModel):
         False,
         description="Whether to stream the message back to the agentex server from the agent",
     )
+    extra_headers: dict[str, str] | None = Field(
+        default=None,
+        description="Custom headers forwarded to this agent (filtered by manifest allowlist)",
+    )
 
 
 class SendEventParams(BaseModel):
@@ -63,11 +73,16 @@ class SendEventParams(BaseModel):
         agent: The agent that the event was sent to.
         task: The task that the message was sent to.
         event: The event that was sent to the agent.
+        extra_headers: Custom headers forwarded to this agent (filtered by manifest allowlist).
     """
 
     agent: Agent = Field(..., description="The agent that the event was sent to")
     task: Task = Field(..., description="The task that the message was sent to")
     event: Event = Field(..., description="The event that was sent to the agent")
+    extra_headers: dict[str, str] | None = Field(
+        default=None,
+        description="Custom headers forwarded to this agent (filtered by manifest allowlist)",
+    )
 
 
 class CancelTaskParams(BaseModel):
@@ -76,10 +91,15 @@ class CancelTaskParams(BaseModel):
     Attributes:
         agent: The agent that the task was sent to.
         task: The task that was cancelled.
+        extra_headers: Custom headers forwarded to this agent (filtered by manifest allowlist).
     """
 
     agent: Agent = Field(..., description="The agent that the task was sent to")
     task: Task = Field(..., description="The task that was cancelled")
+    extra_headers: dict[str, str] | None = Field(
+        default=None,
+        description="Custom headers forwarded to this agent (filtered by manifest allowlist)",
+    )
 
 
 RPC_SYNC_METHODS = [

src/agentex/lib/types/agent_configs.py

@@ -1,5 +1,5 @@
 from typing import Literal
-from pydantic import BaseModel, Field, model_validator, validator
+from pydantic import BaseModel, Field, field_validator, model_validator
 
 
 class TemporalWorkflowConfig(BaseModel):
@@ -101,9 +101,9 @@ class TemporalConfig(BaseModel):
         description="List of temporal workflow configurations. Used when enabled=true.",
     )
 
-    @validator("workflows")
+    @field_validator("workflows")
     @classmethod
-    def validate_workflows_not_empty(cls, v):
+    def validate_workflows_not_empty(cls, v: list[TemporalWorkflowConfig]) -> list[TemporalWorkflowConfig]:
         """Ensure workflows list is not empty when provided"""
         if v is not None and len(v) == 0:
             raise ValueError("workflows list cannot be empty when provided")

tests/test_header_filtering.py

@@ -1,26 +1,16 @@
 """
 Tests for custom header filtering functionality.
 """
+import fnmatch
 import pytest
-from unittest.mock import AsyncMock, MagicMock
 
 from agentex.lib.types.agent_configs import CustomHeadersConfig
-from agentex.lib.core.services.adk.acp.acp import ACPService
 
 
-class TestHeaderFiltering:
-    """Test header filtering functionality in ACPService."""
-    
-    def setup_method(self):
-        """Set up test fixtures."""
-        self.mock_client = MagicMock()
-        self.mock_tracer = MagicMock()
-        self.acp_service = ACPService(
-            agentex_client=self.mock_client,
-            tracer=self.mock_tracer
-        )
+class TestCustomHeadersConfig:
+    """Test CustomHeadersConfig functionality."""
 
-    def test_custom_headers_config_creation(self):
+    def test_custom_headers_config_creation(self) -> None:
         """Test that CustomHeadersConfig can be created with defaults."""
         config = CustomHeadersConfig()
 
@@ -29,7 +19,7 @@ def test_custom_headers_config_creation(self):
         assert config.max_header_size == 8192
         assert config.max_headers_count == 50
 
-    def test_custom_headers_config_with_headers(self):
+    def test_custom_headers_config_with_headers(self) -> None:
         """Test CustomHeadersConfig with allowed headers."""
         config = CustomHeadersConfig(
             allowed_headers=["x-user-email", "x-tenant-id"],
@@ -42,7 +32,7 @@ def test_custom_headers_config_with_headers(self):
         assert config.max_header_size == 4096
         assert config.max_headers_count == 10
 
-    def test_custom_headers_config_validation(self):
+    def test_custom_headers_config_validation(self) -> None:
         """Test CustomHeadersConfig validation for positive values."""
         # Test negative max_header_size
         with pytest.raises(ValueError, match="max_header_size must be greater than 0"):
@@ -52,23 +42,57 @@ def test_custom_headers_config_validation(self):
         with pytest.raises(ValueError, match="max_headers_count must be greater than 0"):
             CustomHeadersConfig(max_headers_count=0)
 
-    def test_filter_headers_no_headers(self):
+
+def filter_headers_standalone(
+    headers: dict[str, str] | None,
+    config: CustomHeadersConfig
+) -> dict[str, str]:
+    """Standalone header filtering function for testing."""
+    if not headers:
+        return {}
+    
+    if not config.allowed_headers:
+        return {}
+    
+    filtered = {}
+    for header_name, header_value in headers.items():
+        # Check size limit
+        if len(header_value) > config.max_header_size:
+            continue
+        
+        # Check if header matches any allowed pattern (case-insensitive)
+        for allowed_pattern in config.allowed_headers:
+            if fnmatch.fnmatch(header_name.lower(), allowed_pattern.lower()):
+                filtered[header_name] = header_value
+                break
+        
+        # Check count limit
+        if len(filtered) >= config.max_headers_count:
+            break
+    
+    return filtered
+
+
+class TestHeaderFiltering:
+    """Test header filtering logic."""
+    
+    def test_filter_headers_no_headers(self) -> None:
         """Test header filtering with no input headers."""
         config = CustomHeadersConfig(allowed_headers=["x-user-email"])
-        result = self.acp_service._filter_headers(None, config)
+        result = filter_headers_standalone(None, config)
         assert result == {}
 
-        result = self.acp_service._filter_headers({}, config)
+        result = filter_headers_standalone({}, config)
         assert result == {}
 
-    def test_filter_headers_no_allowed_headers(self):
+    def test_filter_headers_no_allowed_headers(self) -> None:
         """Test header filtering with no allowed headers (secure by default)."""
         config = CustomHeadersConfig(allowed_headers=[])
         headers = {"x-user-email": "[email protected]", "x-admin-token": "secret"}
-        result = self.acp_service._filter_headers(headers, config)
+        result = filter_headers_standalone(headers, config)
         assert result == {}
 
-    def test_filter_headers_allowed_headers(self):
+    def test_filter_headers_allowed_headers(self) -> None:
         """Test header filtering with allowed headers."""
         config = CustomHeadersConfig(allowed_headers=["x-user-email", "x-tenant-id"])
         headers = {
@@ -77,15 +101,15 @@ def test_filter_headers_allowed_headers(self):
             "x-admin-token": "secret",  # Should be filtered out
             "content-type": "application/json"  # Should be filtered out
         }
-        result = self.acp_service._filter_headers(headers, config)
+        result = filter_headers_standalone(headers, config)
 
         expected = {
             "x-user-email": "[email protected]",
             "x-tenant-id": "tenant123"
         }
         assert result == expected
 
-    def test_filter_headers_case_insensitive_patterns(self):
+    def test_filter_headers_case_insensitive_patterns(self) -> None:
         """Test header filtering with case-insensitive pattern matching."""
         config = CustomHeadersConfig(allowed_headers=["X-User-Email", "x-tenant-*"])
         headers = {
@@ -94,7 +118,7 @@ def test_filter_headers_case_insensitive_patterns(self):
             "x-tenant-name": "acme",  # Should match x-tenant-*
             "x-admin-token": "secret"  # Should be filtered out
         }
-        result = self.acp_service._filter_headers(headers, config)
+        result = filter_headers_standalone(headers, config)
 
         expected = {
             "x-user-email": "[email protected]",
@@ -103,29 +127,22 @@ def test_filter_headers_case_insensitive_patterns(self):
         }
         assert result == expected
 
-    def test_filter_headers_size_limit(self):
+    def test_filter_headers_size_limit(self) -> None:
         """Test header filtering with size limits."""
         config = CustomHeadersConfig(
-            allowed_headers=["x-data"],
+            allowed_headers=["x-data", "x-large-data"],
             max_header_size=10  # Very small limit for testing
         )
         headers = {
             "x-data": "small",  # 5 chars - should pass
-            "x-data-large": "this is a very long header value that exceeds the limit"  # Should be rejected
-        }
-        # Note: x-data-large would be rejected anyway since it's not in allowed_headers
-        # Let's test with allowed header that's too long
-        config.allowed_headers = ["x-data", "x-large-data"]
-        headers = {
-            "x-data": "small",  # Should pass
             "x-large-data": "this header value is way too long for the configured limit"  # Should be rejected due to size
         }
-        result = self.acp_service._filter_headers(headers, config)
+        result = filter_headers_standalone(headers, config)
 
         expected = {"x-data": "small"}
         assert result == expected
 
-    def test_filter_headers_count_limit(self):
+    def test_filter_headers_count_limit(self) -> None:
         """Test header filtering with count limits."""
         config = CustomHeadersConfig(
             allowed_headers=["x-header-*"],
@@ -137,24 +154,13 @@ def test_filter_headers_count_limit(self):
             "x-header-3": "value3",  # Should be ignored due to count limit
             "x-header-4": "value4"   # Should be ignored due to count limit
         }
-        result = self.acp_service._filter_headers(headers, config)
+        result = filter_headers_standalone(headers, config)
 
         # Should only get first 2 headers that match
         assert len(result) == 2
         assert "x-header-1" in result
         assert "x-header-2" in result
 
-    @pytest.mark.asyncio
-    async def test_get_agent_header_config_default(self):
-        """Test getting default agent header configuration."""
-        config = await self.acp_service._get_agent_header_config("test-agent", None)
-        
-        # Should return secure default (no headers allowed)
-        assert config.strategy == "allowlist"
-        assert config.allowed_headers == []
-        assert config.max_header_size == 8192
-        assert config.max_headers_count == 50
-
 
 if __name__ == "__main__":
     pytest.main([__file__])