feat: a2a extensions API and async agent card caching; fix task propagation & streaming

Adds initial extensions API (with registry temporarily no-op), introduces aiocache for async caching, ensures reference task IDs propagate correctly, fixes streamed response model handling, updates streaming tests, and regenerates lockfiles.

Author: greysonlalonde

lib/crewai/pyproject.toml

@@ -95,6 +95,7 @@ a2a = [
      "a2a-sdk~=0.3.10",
      "httpx-auth~=0.23.1",
      "httpx-sse~=0.4.0",
+     "aiocache[redis,memcached]~=0.12.3",
 ]
 
 

lib/crewai/src/crewai/a2a/extensions/__init__.py

@@ -0,0 +1,4 @@
+"""A2A Protocol Extensions for CrewAI.
+
+This module contains extensions to the A2A (Agent-to-Agent) protocol.
+"""

lib/crewai/src/crewai/a2a/extensions/base.py

@@ -0,0 +1,193 @@
+"""Base extension interface for A2A wrapper integrations.
+
+This module defines the protocol for extending A2A wrapper functionality
+with custom logic for conversation processing, prompt augmentation, and
+agent response handling.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Sequence
+from typing import TYPE_CHECKING, Any, Protocol
+
+
+if TYPE_CHECKING:
+    from a2a.types import Message
+
+    from crewai.agent.core import Agent
+
+
+class ConversationState(Protocol):
+    """Protocol for extension-specific conversation state.
+
+    Extensions can define their own state classes that implement this protocol
+    to track conversation-specific data extracted from message history.
+    """
+
+    def is_ready(self) -> bool:
+        """Check if the state indicates readiness for some action.
+
+        Returns:
+            True if the state is ready, False otherwise.
+        """
+        ...
+
+
+class A2AExtension(Protocol):
+    """Protocol for A2A wrapper extensions.
+
+    Extensions can implement this protocol to inject custom logic into
+    the A2A conversation flow at various integration points.
+    """
+
+    def inject_tools(self, agent: Agent) -> None:
+        """Inject extension-specific tools into the agent.
+
+        Called when an agent is wrapped with A2A capabilities. Extensions
+        can add tools that enable extension-specific functionality.
+
+        Args:
+            agent: The agent instance to inject tools into.
+        """
+        ...
+
+    def extract_state_from_history(
+        self, conversation_history: Sequence[Message]
+    ) -> ConversationState | None:
+        """Extract extension-specific state from conversation history.
+
+        Called during prompt augmentation to allow extensions to analyze
+        the conversation history and extract relevant state information.
+
+        Args:
+            conversation_history: The sequence of A2A messages exchanged.
+
+        Returns:
+            Extension-specific conversation state, or None if no relevant state.
+        """
+        ...
+
+    def augment_prompt(
+        self,
+        base_prompt: str,
+        conversation_state: ConversationState | None,
+    ) -> str:
+        """Augment the task prompt with extension-specific instructions.
+
+        Called during prompt augmentation to allow extensions to add
+        custom instructions based on conversation state.
+
+        Args:
+            base_prompt: The base prompt to augment.
+            conversation_state: Extension-specific state from extract_state_from_history.
+
+        Returns:
+            The augmented prompt with extension-specific instructions.
+        """
+        ...
+
+    def process_response(
+        self,
+        agent_response: Any,
+        conversation_state: ConversationState | None,
+    ) -> Any:
+        """Process and potentially modify the agent response.
+
+        Called after parsing the agent's response, allowing extensions to
+        enhance or modify the response based on conversation state.
+
+        Args:
+            agent_response: The parsed agent response.
+            conversation_state: Extension-specific state from extract_state_from_history.
+
+        Returns:
+            The processed agent response (may be modified or original).
+        """
+        ...
+
+
+class ExtensionRegistry:
+    """Registry for managing A2A extensions.
+
+    Maintains a collection of extensions and provides methods to invoke
+    their hooks at various integration points.
+    """
+
+    def __init__(self) -> None:
+        """Initialize the extension registry."""
+        self._extensions: list[A2AExtension] = []
+
+    def register(self, extension: A2AExtension) -> None:
+        """Register an extension.
+
+        Args:
+            extension: The extension to register.
+        """
+        self._extensions.append(extension)
+
+    def inject_all_tools(self, agent: Agent) -> None:
+        """Inject tools from all registered extensions.
+
+        Args:
+            agent: The agent instance to inject tools into.
+        """
+        for extension in self._extensions:
+            extension.inject_tools(agent)
+
+    def extract_all_states(
+        self, conversation_history: Sequence[Message]
+    ) -> dict[type[A2AExtension], ConversationState]:
+        """Extract conversation states from all registered extensions.
+
+        Args:
+            conversation_history: The sequence of A2A messages exchanged.
+
+        Returns:
+            Mapping of extension types to their conversation states.
+        """
+        states: dict[type[A2AExtension], ConversationState] = {}
+        for extension in self._extensions:
+            state = extension.extract_state_from_history(conversation_history)
+            if state is not None:
+                states[type(extension)] = state
+        return states
+
+    def augment_prompt_with_all(
+        self,
+        base_prompt: str,
+        extension_states: dict[type[A2AExtension], ConversationState],
+    ) -> str:
+        """Augment prompt with instructions from all registered extensions.
+
+        Args:
+            base_prompt: The base prompt to augment.
+            extension_states: Mapping of extension types to conversation states.
+
+        Returns:
+            The fully augmented prompt.
+        """
+        augmented = base_prompt
+        for extension in self._extensions:
+            state = extension_states.get(type(extension))
+            augmented = extension.augment_prompt(augmented, state)
+        return augmented
+
+    def process_response_with_all(
+        self,
+        agent_response: Any,
+        extension_states: dict[type[A2AExtension], ConversationState],
+    ) -> Any:
+        """Process response through all registered extensions.
+
+        Args:
+            agent_response: The parsed agent response.
+            extension_states: Mapping of extension types to conversation states.
+
+        Returns:
+            The processed agent response.
+        """
+        processed = agent_response
+        for extension in self._extensions:
+            state = extension_states.get(type(extension))
+            processed = extension.process_response(processed, state)
+        return processed

lib/crewai/src/crewai/a2a/extensions/registry.py

@@ -0,0 +1,34 @@
+"""Extension registry factory for A2A configurations.
+
+This module provides utilities for creating extension registries from A2A configurations.
+"""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+from crewai.a2a.extensions.base import ExtensionRegistry
+
+
+if TYPE_CHECKING:
+    from crewai.a2a.config import A2AConfig
+
+
+def create_extension_registry_from_config(
+    a2a_config: list[A2AConfig] | A2AConfig,
+) -> ExtensionRegistry:
+    """Create an extension registry from A2A configuration.
+
+    Args:
+        a2a_config: A2A configuration (single or list)
+
+    Returns:
+        Configured extension registry with all applicable extensions
+    """
+    registry = ExtensionRegistry()
+    configs = a2a_config if isinstance(a2a_config, list) else [a2a_config]
+
+    for _ in configs:
+        pass
+
+    return registry

lib/crewai/src/crewai/a2a/utils.py

@@ -23,6 +23,8 @@
     TextPart,
     TransportProtocol,
 )
+from aiocache import cached  # type: ignore[import-untyped]
+from aiocache.serializers import PickleSerializer  # type: ignore[import-untyped]
 import httpx
 from pydantic import BaseModel, Field, create_model
 
@@ -65,7 +67,7 @@ def _fetch_agent_card_cached(
         endpoint: A2A agent endpoint URL
         auth_hash: Hash of the auth object
         timeout: Request timeout
-        _ttl_hash: Time-based hash for cache invalidation (unused in body)
+        _ttl_hash: Time-based hash for cache invalidation
 
     Returns:
         Cached AgentCard
@@ -106,7 +108,18 @@ def fetch_agent_card(
         A2AClientHTTPError: If authentication fails
     """
     if use_cache:
-        auth_hash = hash((type(auth).__name__, id(auth))) if auth else 0
+        if auth:
+            auth_data = auth.model_dump_json(
+                exclude={
+                    "_access_token",
+                    "_token_expires_at",
+                    "_refresh_token",
+                    "_authorization_callback",
+                }
+            )
+            auth_hash = hash((type(auth).__name__, auth_data))
+        else:
+            auth_hash = 0
         _auth_store[auth_hash] = auth
         ttl_hash = int(time.time() // cache_ttl)
         return _fetch_agent_card_cached(endpoint, auth_hash, timeout, ttl_hash)
@@ -121,6 +134,26 @@ def fetch_agent_card(
         loop.close()
 
 
+@cached(ttl=300, serializer=PickleSerializer())  # type: ignore[untyped-decorator]
+async def _fetch_agent_card_async_cached(
+    endpoint: str,
+    auth_hash: int,
+    timeout: int,
+) -> AgentCard:
+    """Cached async implementation of AgentCard fetching.
+
+    Args:
+        endpoint: A2A agent endpoint URL
+        auth_hash: Hash of the auth object
+        timeout: Request timeout in seconds
+
+    Returns:
+        Cached AgentCard object
+    """
+    auth = _auth_store.get(auth_hash)
+    return await _fetch_agent_card_async(endpoint=endpoint, auth=auth, timeout=timeout)
+
+
 async def _fetch_agent_card_async(
     endpoint: str,
     auth: AuthScheme | None,
@@ -339,7 +372,22 @@ async def _execute_a2a_delegation_async(
     Returns:
         Dictionary with status, result/error, and new history
     """
-    agent_card = await _fetch_agent_card_async(endpoint, auth, timeout)
+    if auth:
+        auth_data = auth.model_dump_json(
+            exclude={
+                "_access_token",
+                "_token_expires_at",
+                "_refresh_token",
+                "_authorization_callback",
+            }
+        )
+        auth_hash = hash((type(auth).__name__, auth_data))
+    else:
+        auth_hash = 0
+    _auth_store[auth_hash] = auth
+    agent_card = await _fetch_agent_card_async_cached(
+        endpoint=endpoint, auth_hash=auth_hash, timeout=timeout
+    )
 
     validate_auth_against_agent_card(agent_card, auth)
 
@@ -556,6 +604,34 @@ async def _execute_a2a_delegation_async(
                         }
                         break
         except Exception as e:
+            if isinstance(e, A2AClientHTTPError):
+                error_msg = f"HTTP Error {e.status_code}: {e!s}"
+
+                error_message = Message(
+                    role=Role.agent,
+                    message_id=str(uuid.uuid4()),
+                    parts=[Part(root=TextPart(text=error_msg))],
+                    context_id=context_id,
+                    task_id=task_id,
+                )
+                new_messages.append(error_message)
+
+                crewai_event_bus.emit(
+                    None,
+                    A2AResponseReceivedEvent(
+                        response=error_msg,
+                        turn_number=turn_number,
+                        is_multiturn=is_multiturn,
+                        status="failed",
+                        agent_role=agent_role,
+                    ),
+                )
+                return {
+                    "status": "failed",
+                    "error": error_msg,
+                    "history": new_messages,
+                }
+
             current_exception: Exception | BaseException | None = e
             while current_exception:
                 if hasattr(current_exception, "response"):
@@ -752,4 +828,5 @@ def get_a2a_agents_and_response_model(
         Tuple of A2A agent IDs and response model
     """
     a2a_agents, agent_ids = extract_a2a_agent_ids_from_config(a2a_config=a2a_config)
+
     return a2a_agents, create_agent_response_model(agent_ids)