aip-agents-binary 0.0.0__py3-none-any.whl

aip_agents/agent/base_langgraph_agent.py

@@ -0,0 +1,4048 @@
+"""Base class for LangGraph-based agent implementations.
+
+This class provides the core LangGraph machinery including graph compilation,
+state handling, and I/O mapping for LangGraph agents.
+
+Authors:
+    Christian Trisno Sen Long Chen (christian.t.s.l.chen@gdplabs.id)
+    Putu Ravindra Wiguna (putu.r.wiguna@gdplabs.id)
+    Fachriza Adhiatma (fachriza.d.adhiatma@gdplabs.id)
+    Raymond Christopher (raymond.christopher@gdplabs.id)
+    Reinhart Linanda (reinhart.linanda@gdplabs.id)
+"""
+
+import asyncio
+import copy
+import hashlib
+import inspect
+import json
+import uuid
+from abc import abstractmethod
+from collections.abc import AsyncGenerator, Sequence
+from contextlib import suppress
+from contextvars import ContextVar
+from dataclasses import dataclass
+from typing import Any, cast
+
+from a2a.types import AgentCard
+from aiostream import stream as astream
+from gllm_core.event import EventEmitter  # type: ignore[import-untyped]
+from gllm_core.event.handler import StreamEventHandler  # type: ignore[import-untyped]
+from gllm_core.schema import Chunk, Event  # type: ignore[import-untyped]
+from langchain_core.messages import AIMessage, BaseMessage, HumanMessage, SystemMessage, ToolMessage
+from langchain_core.tools import BaseTool
+from langgraph.graph import StateGraph
+from langgraph.graph.state import CompiledStateGraph
+from langgraph.types import Checkpointer, StreamWriter
+from pydantic import ValidationError
+
+from aip_agents.agent.base_agent import BaseAgent
+from aip_agents.agent.graph_runtime import CONFIGURABLE_KEY as _CONFIGURABLE_KEY
+from aip_agents.agent.graph_runtime import GraphRuntime
+from aip_agents.agent.system_instruction_context import get_current_date_context
+from aip_agents.context.prompt import PromptBuildContext
+from aip_agents.mcp.client import LangchainMCPClient
+from aip_agents.mcp.client.errors import MCPServerStartupError
+from aip_agents.memory import BaseMemory, MemoryFactory, MemoryMethod
+from aip_agents.memory.constants import MemoryDefaults
+from aip_agents.middleware.manager import MiddlewareManager
+from aip_agents.schema.agent import StreamMode
+from aip_agents.schema.hitl import HitlMetadata
+from aip_agents.schema.step_limit import StepLimitConfig
+from aip_agents.skills.errors import SkillInstallError
+from aip_agents.tools.tool_config_injector import (
+    CONFIG_SCHEMA_ATTR,
+    TOOL_CONFIG_SCHEMA_ATTR,
+    inject_config_methods_into_tool,
+)
+from aip_agents.types import A2AEvent, A2AStreamEventType
+from aip_agents.utils import augment_query_with_file_paths, build_input_attachment_registry, validate_references
+from aip_agents.utils.attachment_helpers import attach_attachments_to_state, merge_input_attachments
+from aip_agents.utils.langgraph.converter import (
+    THINKINGS_METADATA_KEY,
+    convert_langchain_messages_to_gllm_messages,
+    convert_lm_output_to_langchain_message,
+)
+from aip_agents.utils.langgraph.delegation_context import is_delegated_agent_execution
+from aip_agents.utils.langgraph.tool_managers.a2a_tool_manager import A2AToolManager
+from aip_agents.utils.langgraph.tool_managers.delegation_tool_manager import (
+    DelegationToolManager,
+)
+from aip_agents.utils.logger import get_logger
+from aip_agents.utils.metadata.activity_metadata_helper import create_tool_activity_info
+from aip_agents.utils.metadata_helper import (
+    DefaultStepMessages,
+    Kind,
+    MetadataFieldKeys,
+    Status,
+    end_step_counter_scope,
+    get_next_step_number,
+    start_step_counter_scope,
+)
+from aip_agents.utils.pii import deanonymize_final_response_content
+from aip_agents.utils.sse_chunk_transformer import SSEChunkTransformer
+from aip_agents.utils.step_limit_manager import _STEP_LIMIT_CONFIG_CVAR
+from aip_agents.utils.token_usage_helper import (
+    MODEL_USAGE_DELTA_KEY,
+    MODEL_USAGE_KEY,
+    STEP_USAGE_KEY,
+    TOTAL_USAGE_KEY,
+    USAGE_METADATA_KEY,
+)
+
+logger = get_logger(__name__)
+
+# Fixed bounded timeout (in seconds) for cancellation-safe auto-cleanup.
+# This value is intentionally not configurable via env-var or per-call override
+# in this iteration; see spec for rationale.
+_AUTO_CLEANUP_TIMEOUT_S: float = 1.0
+
+# Context variable to access current thread_id during streaming callbacks
+_THREAD_ID_CVAR: ContextVar[str | None] = ContextVar("aip_agents_thread_id", default=None)
+
+# Context variable to track operation mode for dependency tracking
+# "parallel" = include all completed steps (default for backward compatibility)
+# "sequential" = include only the most recent completed step
+_OPERATION_MODE_CVAR: ContextVar[str] = ContextVar("aip_agents_operation_mode", default="parallel")
+
+
+@dataclass
+class _StreamingContext:
+    """Context object for managing streaming state and configuration."""
+
+    original_query: str
+    graph_input: dict[str, Any]
+    config: dict[str, Any]
+    memory_user_id: str | None
+    current_thread_id: str
+    token: Any
+    enable_token_streaming: bool
+    a2ui_stream_filter: Any | None = None
+
+    # Streaming state
+    final_event_yielded: bool = False
+    pending_artifacts: list | None = None
+    seen_artifact_hashes: set | None = None
+    processed_message_count: int = 0
+    final_state: dict[str, Any] | None = None
+    last_final_content: str | None = None
+    saved_memory: bool = False
+    is_token_streaming: bool = False
+
+    def __post_init__(self):
+        """Initialize mutable defaults."""
+        if self.pending_artifacts is None:
+            self.pending_artifacts = []
+        if self.seen_artifact_hashes is None:
+            self.seen_artifact_hashes = set()
+        if self.final_state is None:
+            self.final_state = {}
+
+
+MCP_INIT_TIMEOUT = 60.0
+
+
+class BaseLangGraphAgent(BaseAgent):
+    """Base class for LangGraph-based agents with unified tool approach.
+
+    Provides core LangGraph functionality including:
+    - Graph compilation and execution
+    - State schema management
+    - I/O mapping between user inputs and graph states
+    - Event emission support
+    - Tool resolution and handling
+    - A2A communication capabilities via tools
+    - Agent delegation capabilities via tools
+    - MCP server integration via tools
+    - Enhanced output extraction from various state formats
+    - Automatic resource cleanup after each run (arun/arun_stream/arun_a2a_stream)
+
+    Tool Management:
+    - regular_tools: Standard LangChain tools provided during initialization
+    - mcp_tools: Tools retrieved from MCP servers
+    - resolved_tools: Combined collection of all tools for graph execution
+
+    Subclasses must implement:
+    - define_graph(): Define the specific graph structure
+    - _prepare_graph_input(): Convert user input to graph state
+    - _format_graph_output(): Convert final graph state to user output
+    """
+
+    def __init__(  # noqa: PLR0913
+        self,
+        name: str,
+        instruction: str,
+        description: str | None = None,
+        model: Any | None = None,
+        tools: Sequence[BaseTool] | None = None,
+        state_schema: type | None = None,
+        thread_id_key: str = "thread_id",
+        event_emitter: EventEmitter | None = None,
+        checkpointer: Checkpointer | None = None,
+        enable_a2a_token_streaming: bool = False,
+        recursion_limit: int | None = None,
+        graph_config: dict[str, Any] | None = None,
+        **kwargs: Any,
+    ):
+        """Initialize the BaseLangGraphAgent.
+
+        Args:
+            name: The name of the agent.
+            instruction: The system instruction for the agent.
+            description: Human-readable description of the agent.
+            model: The model to use (lm_invoker, LangChain model, string, etc.).
+            tools: Sequence of regular LangChain tools (not A2A or delegation tools).
+            state_schema: The state schema for the LangGraph. Defaults to basic message state.
+            thread_id_key: Key for thread ID in configuration.
+            event_emitter: Optional event emitter for streaming updates.
+            checkpointer: Optional checkpointer for conversation persistence.
+            enable_a2a_token_streaming: Enable token-level streaming for A2A responses.
+                - False (default): Stream message-level events only
+                - True: Stream individual tokens plus message-level events
+            recursion_limit: Optional maximum number of graph steps before LangGraph raises
+                ``GraphRecursionError``.  Must be a positive integer when provided.
+                ``None`` (default) auto-derives ``StepLimitConfig().max_steps * 2``
+                (200 by default) via ``_effective_recursion_limit``; subclasses may
+                override ``_effective_recursion_limit`` to use a different source (e.g.
+                ``LangGraphReactAgent`` uses ``step_limit_config.max_steps * 2``).
+                A per-call value passed to ``_create_graph_config`` takes precedence.
+            graph_config: Optional baseline ``RunnableConfig`` keys (excluding
+                ``"configurable"``) applied to every graph invocation.  Keys are merged at
+                the lowest priority — the named ``recursion_limit`` param and any per-call
+                kwargs override them.  Passing ``"configurable"`` raises ``ValueError``
+                immediately because thread state is managed internally.
+            **kwargs: Additional keyword arguments passed to BaseAgent (including tool_configs and memory settings).
+                Memory settings include:
+                - memory_backend: Memory backend (e.g., "mem0")
+                - agent_id: Agent identifier for memory scoping
+                - memory_namespace: Memory namespace
+                - save_interaction_to_memory: Whether to save interactions (default True)
+
+        Raises:
+            ValueError: If ``recursion_limit`` is not ``None`` and is less than or equal to 0.
+            ValueError: If ``graph_config`` contains the key ``"configurable"``.
+        """
+        if recursion_limit is not None and recursion_limit <= 0:
+            raise ValueError(f"recursion_limit must be a positive integer, got {recursion_limit}")
+        if _CONFIGURABLE_KEY in (graph_config or {}):
+            raise ValueError(
+                "'configurable' cannot be set via graph_config. "
+                "Thread ID and checkpoint state are managed internally by the agent."
+            )
+
+        super().__init__(
+            name=name,
+            instruction=instruction,
+            description=description,
+            model=model,
+            tools=list(tools) if tools else [],
+            **kwargs,
+        )
+
+        self._prompt_build_context = PromptBuildContext(
+            base_instruction=instruction,
+            date_time_context=get_current_date_context(),
+        )
+
+        self.state_schema = state_schema
+        self.thread_id_key = thread_id_key
+        self.enable_a2a_token_streaming = enable_a2a_token_streaming
+        self.event_emitter = event_emitter
+        self.checkpointer = checkpointer
+        self.recursion_limit: int | None = recursion_limit
+        self.graph_config: dict[str, Any] = graph_config or {}
+        self._graph_runtime = GraphRuntime(
+            agent_name=lambda: self.name,
+            thread_id_key=lambda: self.thread_id_key,
+            graph_config=lambda: self.graph_config,
+            recursion_limit=lambda: self.recursion_limit,
+            effective_recursion_limit=lambda: self._effective_recursion_limit,
+            logger_=logger,
+        )
+        self.tool_output_manager = None
+
+        self._mem0_client: Any | None = None
+        self.memory: BaseMemory | None = None
+        self._initialize_memory_from_kwargs(name, kwargs)
+
+        self.a2a_tool_manager = A2AToolManager()
+        self.delegation_tool_manager = DelegationToolManager(parent_agent=self)
+
+        # Always initialize middleware manager (subclasses can override via _setup_middleware)
+        self._middleware_manager = MiddlewareManager([])
+
+        self.regular_tools: list[BaseTool] = self._resolve_and_validate_tools()
+        self.mcp_tools: list[BaseTool] = []
+        self.resolved_tools: list[BaseTool] = self.regular_tools.copy()
+
+        self._compiled_graph = self._build_and_compile_graph()
+
+        self._tool_parent_map_by_thread: dict[str, dict[str, str]] = {}
+        self._completed_tool_steps_by_thread: dict[str, list[str]] = {}
+        self._last_status_step_id_by_thread: dict[str, str] = {}
+        self._coordinator_completed_tool_steps_by_thread: dict[str, list[str]] = {}
+        self._emitted_tool_calls_by_thread: dict[str, set[str]] = {}
+        self._last_compaction_preparation: Any = None
+
+    def _create_default_event_emitter(self) -> EventEmitter:
+        """Create default event emitter for token streaming.
+
+        Returns:
+            EventEmitter with StreamEventHandler configured for token streaming.
+        """
+        stream_handler = StreamEventHandler(name=f"{self.name}_A2AStreamHandler")
+        logger.info(f"Agent '{self.name}': Auto-created event emitter for token streaming")
+        return EventEmitter(handlers=[stream_handler])
+
+    def _log_streaming_event_debug(self, source: str, event: dict[str, Any]) -> None:
+        """Log the raw streaming event for debugging purposes.
+
+        Args:
+            source: A short label describing where the event originated.
+            event: The event payload emitted by the streaming pipeline.
+        """
+        try:
+            logger.info("Streaming event (%s): %s", source, event)
+        except Exception as exc:  # noqa: BLE001
+            logger.debug("Failed to log streaming event: %s", exc, exc_info=True)
+
+    def _record_emitted_tool_calls(self, tool_calls_details: list[dict[str, Any]]) -> None:
+        """Track tool call IDs that have already been emitted to avoid duplicates.
+
+        Args:
+            tool_calls_details: Tool call metadata emitted by the tool_call event.
+        """
+        thread_id = _THREAD_ID_CVAR.get()
+        if not thread_id or not tool_calls_details:
+            return
+
+        emitted = self._emitted_tool_calls_by_thread.setdefault(thread_id, set())
+        for details in tool_calls_details:
+            call_id = details.get("id")
+            if isinstance(call_id, str) and call_id:
+                emitted.add(call_id)
+                logger.info(
+                    "Registered tool call event: agent=%s thread=%s call_id=%s payload=%s",
+                    self.name,
+                    thread_id,
+                    call_id,
+                    details,
+                )
+
+    def _discard_emitted_tool_call(self, tool_call_id: str | None) -> None:
+        """Remove a tool call ID from the emitted tracker.
+
+        Args:
+            tool_call_id: Identifier of the tool call to remove from cache.
+        """
+        if not tool_call_id:
+            return
+        thread_id = _THREAD_ID_CVAR.get()
+        if not thread_id:
+            return
+        emitted = self._emitted_tool_calls_by_thread.get(thread_id)
+        if emitted:
+            emitted.discard(tool_call_id)
+            logger.info(
+                "Cleared recorded tool call: agent=%s thread=%s call_id=%s",
+                self.name,
+                thread_id,
+                tool_call_id,
+            )
+
+    def _get_stream_handler(self) -> StreamEventHandler | None:
+        """Get StreamEventHandler from event_emitter if available.
+
+        Returns:
+            StreamEventHandler instance if found, None otherwise.
+        """
+        if not self.event_emitter or not self.event_emitter.handlers:
+            return None
+
+        for handler in self.event_emitter.handlers:
+            if isinstance(handler, StreamEventHandler):
+                return handler
+        return None
+
+    def set_operation_mode(self, mode: str) -> None:
+        """Set the operation mode for dependency tracking.
+
+        Args:
+            mode: Operation mode - "parallel" (default) or "sequential"
+        """
+        if mode not in ["parallel", "sequential"]:
+            raise ValueError(f"Invalid operation mode: {mode}. Must be 'parallel' or 'sequential'")
+        _OPERATION_MODE_CVAR.set(mode)
+
+    def _default_memory_agent_id(self, name: str) -> str:
+        """Create a stable identifier for memory scoping.
+
+        Args:
+            name: The agent's human-readable name.
+
+        Returns:
+            str: A deterministic ID derived from the class and name, suitable for scoping memory per agent.
+        """
+        base = f"{self.__class__.__name__}:{name}"
+        return f"{MemoryDefaults.AGENT_ID_PREFIX}{hashlib.sha256(base.encode()).hexdigest()}"
+
+    @staticmethod
+    def _parse_bool_value(value: Any) -> bool:
+        """Parse a value to boolean with string handling for "true"/"false".
+
+        Treats string "false" as False, "true" as True.
+        For other values, uses standard bool() conversion.
+
+        Args:
+            value: The value to parse.
+
+        Returns:
+            bool: The parsed boolean value.
+        """
+        if isinstance(value, str):
+            lower_value = value.lower().strip()
+            if lower_value == "false":
+                return False
+            elif lower_value == "true":
+                return True
+        return bool(value)
+
+    def _memory_enabled(self) -> bool:
+        """Check whether memory is enabled for this agent.
+
+        Returns:
+            bool: True when a memory adapter is set.
+        """
+        return self.memory is not None
+
+    def _has_lm_invoker(self) -> bool:
+        """Check whether lm_invoker is available for this agent.
+
+        Returns:
+            bool: True when lm_invoker attribute exists and is not None.
+        """
+        return self.lm_invoker is not None
+
+    def _memory_search(self, query: str, override_user_id: str | None = None) -> list[dict[str, Any]]:
+        """Search for relevant memories using the configured adapter.
+
+        Args:
+            query: The user query to retrieve relevant memories for.
+            override_user_id: Optional per-call override for the memory scope.
+
+        Returns:
+            list[dict[str, Any]]: Memory hits; empty list on failure or when disabled.
+        """
+        if not (self._memory_enabled() and isinstance(query, str)):
+            return []
+        try:
+            user_id = override_user_id or self.memory_agent_id
+            if hasattr(self.memory, MemoryMethod.SEARCH):
+                results = self.memory.search(  # type: ignore[attr-defined]
+                    query=query,
+                    user_id=user_id,
+                    limit=self.memory_retrieval_limit,
+                )
+                return results if isinstance(results, list) else []
+            return []
+        except Exception as e:  # noqa: BLE001
+            logger.debug(f"Memory: search ignored error: {e}")
+        return []
+
+    def _resolve_and_validate_tools(self) -> list[BaseTool]:
+        """Resolve and validate regular tools for LangGraph usage.
+
+        Also configures tools with injected configuration capabilities
+        from agent-level tool_configs.
+
+        Returns:
+            List of resolved LangChain BaseTool instances.
+        """
+        resolved = []
+        for tool in self.tools:
+            if isinstance(tool, BaseTool):
+                self._configure_injected_tool(tool)
+                resolved.append(tool)
+            else:
+                logger.warning(f"Agent '{self.name}': Tool {tool} is not a LangChain BaseTool, skipping")
+
+        logger.info(f"Agent '{self.name}': Resolved {len(resolved)} regular tools for LangGraph")
+        return resolved
+
+    def _initialize_memory_from_kwargs(self, agent_name: str, kwargs: dict[str, Any]) -> None:
+        """Initialize memory-related settings and adapter.
+
+        Extracts known memory kwargs, sets defaults, and initializes the adapter when enabled.
+        Keeps ``__init__`` concise and improves DX.
+
+        Args:
+            agent_name: The name of the agent, used to derive a default memory id.
+            kwargs: Keyword arguments from the agent constructor; consumed keys are removed.
+                Supported memory kwargs:
+                - memory_backend: str - Memory backend to use (e.g., "mem0")
+                - agent_id: str - Agent identifier for memory scoping
+                - memory_namespace: str - Memory namespace
+                - memory_retrieval_limit: int - Max memories to retrieve
+                - memory_max_chars: int - Max characters per memory
+                - save_interaction_to_memory: bool (default True) - Whether to save interactions to memory
+        """
+        # Initialize memory configuration settings
+        self.memory_backend: str | None = kwargs.pop("memory_backend", None)
+        self.memory_agent_id: str = str(kwargs.pop("agent_id", self._default_memory_agent_id(agent_name)))
+        self.memory_namespace: str | None = kwargs.pop("memory_namespace", None)
+        self.memory_retrieval_limit: int = int(kwargs.pop("memory_retrieval_limit", MemoryDefaults.RETRIEVAL_LIMIT))
+        self.memory_max_chars: int = int(kwargs.pop("memory_max_chars", MemoryDefaults.MAX_CHARS))
+
+        # Initialize memory interaction saving flag with proper bool conversion
+        save_raw = kwargs.pop("save_interaction_to_memory", True)
+        self.save_interaction_to_memory: bool = self._parse_bool_value(save_raw)
+
+        if self.memory_backend:
+            memory_kwargs = {
+                "limit": self.memory_retrieval_limit,
+                "max_chars": self.memory_max_chars,
+                "agent_id": self.memory_agent_id,
+            }
+            if self.memory_namespace:
+                memory_kwargs["namespace"] = self.memory_namespace
+
+            self._mem0_client = MemoryFactory.create(
+                self.memory_backend,
+                **memory_kwargs,
+            )
+            self.memory = self._mem0_client
+
+    def _configure_injected_tool(self, tool: BaseTool) -> None:
+        """Configure a tool with automatic configuration injection using agent-level defaults.
+
+        Args:
+            tool: The tool instance to configure.
+        """
+        if self._should_auto_inject_config(tool):
+            self._auto_inject_config_capabilities(tool)
+            self._apply_agent_config_to_tool(tool)
+
+    def _should_auto_inject_config(self, tool: BaseTool) -> bool:
+        """Check if tool needs auto-injection of configuration capabilities.
+
+        Args:
+            tool: The tool instance to check.
+
+        Returns:
+            True if tool needs auto-injection of configuration capabilities, False otherwise.
+        """
+        return (
+            hasattr(tool, TOOL_CONFIG_SCHEMA_ATTR)
+            and getattr(tool, TOOL_CONFIG_SCHEMA_ATTR) is not None
+            and not hasattr(tool, CONFIG_SCHEMA_ATTR)
+        )
+
+    def _auto_inject_config_capabilities(self, tool: BaseTool) -> None:
+        """Inject configuration capabilities into a tool with tool_config_schema.
+
+        Args:
+            tool: The tool instance to configure.
+        """
+        try:
+            tool_config_schema = getattr(tool, TOOL_CONFIG_SCHEMA_ATTR)
+            inject_config_methods_into_tool(tool, tool_config_schema)
+            logger.info(f"Agent '{self.name}': Auto-injected config capabilities into tool '{tool.name}'")
+        except Exception as e:
+            logger.warning(f"Agent '{self.name}': Failed to auto-inject config into tool '{tool.name}': {e}")
+
+    def _apply_agent_config_to_tool(self, tool: BaseTool) -> None:
+        """Apply agent-level configuration to a tool.
+
+        Args:
+            tool: The tool instance to configure.
+        """
+        tool_config_data = self._get_agent_config_for_tool(tool.name)
+
+        if tool_config_data is None:
+            return
+
+        try:
+            set_tool_config = getattr(tool, "set_tool_config", None)
+            if callable(set_tool_config):
+                set_tool_config(tool_config_data)
+            else:
+                raise AttributeError("set_tool_config not available")
+            logger.info(f"Agent '{self.name}': Configured tool '{tool.name}' with agent defaults: {tool_config_data}")
+        except Exception as e:
+            logger.warning(f"Agent '{self.name}': Failed to configure tool '{tool.name}': {e}")
+
+    def _get_agent_config_for_tool(self, tool_name: str) -> dict[str, Any] | None:
+        """Get agent-level configuration data for a specific tool.
+
+        This method intentionally returns only per-tool configuration and does NOT include
+        global agent configuration. Global configuration merging is handled separately
+        in the metadata resolution process during tool execution.
+
+        Args:
+            tool_name: The name of the tool to get configuration for.
+
+        Returns:
+            The configuration data for the tool, or None if no configuration is found.
+        """
+        if not isinstance(self.tool_configs, dict):
+            return None
+
+        return self.tool_configs.get(tool_name)
+
+    def _sanitize_tool_names(self) -> None:
+        """Correct resolved_tools' names that will be used for the agent according to the model provider's rules."""
+        for tool in self.resolved_tools:
+            sanitized_name = self.name_preprocessor.sanitize_tool_name(tool.name)
+            tool.name = sanitized_name
+
+    def _build_and_compile_graph(self) -> CompiledStateGraph:
+        """Build and compile the LangGraph while ensuring tool names are valid.
+
+        Returns:
+            Compiled LangGraph ready for execution.
+        """
+        self._sanitize_tool_names()
+        try:
+            compiled_graph = self._graph_runtime.build_and_compile_graph(self.state_schema, self.define_graph)
+            logger.info(
+                f"Agent '{self.name}': Successfully compiled LangGraph with {len(self.resolved_tools)} total tools"
+            )
+            return compiled_graph
+
+        except Exception as e:
+            logger.error(f"Agent '{self.name}': Failed to build LangGraph: {e}")
+            raise RuntimeError(f"Failed to build LangGraph for agent '{self.name}': {e}") from e
+
+    @abstractmethod
+    def define_graph(self, graph_builder: StateGraph) -> CompiledStateGraph:
+        """Define the specific graph structure for this agent type.
+
+        Subclasses must implement this method to:
+        1. Add nodes to the graph_builder
+        2. Add edges and conditional edges
+        3. Set entry points
+        4. Return the compiled graph
+
+        Args:
+            graph_builder: The StateGraph builder to define nodes and edges on.
+
+        Returns:
+            The compiled graph ready for execution.
+        """
+        raise NotImplementedError(f"Agent '{self.name}': Subclasses must implement define_graph method")
+
+    @abstractmethod
+    def _prepare_graph_input(self, input_data: Any, **kwargs: Any) -> dict[str, Any]:
+        """Convert user input to graph state format.
+
+        Args:
+            input_data: The user's input (query string, structured data, etc.).
+            **kwargs: Additional keyword arguments from the user.
+
+        Returns:
+            Dictionary representing the initial graph state.
+        """
+        raise NotImplementedError(f"Agent '{self.name}': Subclasses must implement _prepare_graph_input method")
+
+    @abstractmethod
+    def _format_graph_output(self, final_state_result: dict[str, Any]) -> Any:
+        """Convert final graph state to user-friendly output.
+
+        Args:
+            final_state_result: The final state from graph execution.
+
+        Returns:
+            Formatted output for the user.
+        """
+        raise NotImplementedError(f"Agent '{self.name}': Subclasses must implement _format_graph_output method")
+
+    def _extract_metadata_from_kwargs(self, **kwargs: Any) -> dict[str, Any]:
+        """Extract metadata from kwargs for agent implementations.
+
+        Supports both flat and mixed metadata schemas:
+        - Flat dict (legacy): all keys applied to all tools and model calls
+        - Mixed dict (new): top-level keys applied to all tools, 'tool_configs' section per-tool
+
+        Args:
+            **kwargs: Keyword arguments that may contain metadata.
+
+        Returns:
+            dict[str, Any]: The metadata dictionary, or an empty dict if no metadata was provided.
+        """
+        return kwargs.get("metadata", {})
+
+    def _extract_output_from_dict_state(self, dict_state: dict[str, Any]) -> str | None:
+        """Extract output from a dictionary state (migrated from BaseLangChainAgent).
+
+        Args:
+            dict_state: A dictionary containing agent state information.
+
+        Returns:
+            The extracted output string or None if no valid output found.
+        """
+        output_content: str | None = None
+        messages = dict_state.get("messages")
+        if messages and isinstance(messages, list) and messages:
+            last_message = messages[-1]
+            if isinstance(last_message, AIMessage):
+                candidate_content = getattr(last_message, "content", None)
+                if not candidate_content:
+                    output_content = ""
+                else:
+                    output_content = candidate_content
+            elif hasattr(last_message, "content"):
+                output_content = getattr(last_message, "content", None)
+        if output_content is None:
+            candidate_output_from_key = dict_state.get("output")
+            if isinstance(candidate_output_from_key, str):
+                output_content = candidate_output_from_key
+        return output_content
+
+    def _extract_output_from_list_state(self, list_state: list[Any]) -> str | None:
+        """Extract output from a list state (migrated from BaseLangChainAgent).
+
+        Args:
+            list_state: A list containing agent state information.
+
+        Returns:
+            The extracted output string or None if no valid output found.
+        """
+        output_content: str | None = None
+        if not list_state:
+            return None
+        last_item = list_state[-1]
+        if isinstance(last_item, AIMessage) and getattr(last_item, "content", None) is not None:
+            output_content = self._normalize_event_content(last_item.content)
+        elif isinstance(last_item, str):
+            output_content = last_item
+        return output_content
+
+    def _extract_output_from_final_state(self, final_state_result: Any) -> str:
+        """Enhanced output extraction from final state (migrated from BaseLangChainAgent).
+
+        Args:
+            final_state_result: The final state from graph execution.
+
+        Returns:
+            Extracted text content.
+        """
+        output_content: str | None = None
+        if isinstance(final_state_result, dict):
+            output_content = self._extract_output_from_dict_state(final_state_result)
+        elif isinstance(final_state_result, str):
+            output_content = final_state_result
+        elif isinstance(final_state_result, list):
+            output_content = self._extract_output_from_list_state(final_state_result)
+
+        if output_content is None:
+            return "Error: Could not extract output from agent's final state."
+        return output_content
+
+    def register_a2a_agents(self, agent_cards: list[AgentCard]) -> None:
+        """Register A2A communication capabilities using the A2A tool manager.
+
+        Args:
+            agent_cards (list[AgentCard]): List of AgentCard instances for external communication.
+        """
+        if not agent_cards:
+            logger.info(f"Agent '{self.name}': No A2A agents to register")
+            return
+
+        a2a_tools = self.a2a_tool_manager.register_resources(agent_cards)
+        self.resolved_tools.extend(a2a_tools)
+
+        logger.info(f"Agent '{self.name}': Registered {len(agent_cards)} A2A agents as tools")
+        self._rebuild_graph()
+
+    def register_delegation_agents(self, agents: list[BaseAgent]) -> None:
+        """Register internal agent delegation capabilities using the delegation tool manager.
+
+        Args:
+            agents: List of BaseAgent instances for internal task delegation.
+        """
+        if not agents:
+            logger.info(f"Agent '{self.name}': No delegation agents to register")
+            return
+
+        delegation_tools = self.delegation_tool_manager.register_resources(agents)
+        self.resolved_tools.extend(delegation_tools)
+        logger.info(f"Agent '{self.name}': Registered {len(agents)} delegation agents as streaming tools")
+
+        self._rebuild_graph()
+
+    def update_regular_tools(self, new_tools: list[BaseTool], rebuild_graph: bool | None = None) -> None:
+        """Update regular tools (not capabilities).
+
+        Args:
+            new_tools: New list of regular tools to use.
+            rebuild_graph: Whether to rebuild graph. If None, uses auto_rebuild_graph setting.
+        """
+        logger.info(f"Agent '{self.name}': Updating regular tools from {len(self.tools)} to {len(new_tools)}")
+
+        self.tools = list(new_tools)
+        old_resolved_count = len(self.regular_tools)
+        self.regular_tools = self._resolve_and_validate_tools()
+
+        logger.info(
+            f"Agent '{self.name}': Regular tools changed from {old_resolved_count} to {len(self.regular_tools)}"
+        )
+
+        self._rebuild_resolved_tools()
+
+        should_rebuild = rebuild_graph if rebuild_graph is not None else True
+        if should_rebuild:
+            try:
+                logger.info(f"Agent '{self.name}': Rebuilding graph with updated tools")
+                self._compiled_graph = self._build_and_compile_graph()
+            except Exception as e:
+                logger.error(f"Agent '{self.name}': Failed to rebuild graph after tool update: {e}")
+                raise
+
+    def _rebuild_resolved_tools(self) -> None:
+        """Rebuild resolved tools combining regular tools with capability tools."""
+        self.resolved_tools = self.regular_tools.copy()
+
+        if self.a2a_tool_manager:
+            a2a_tools = self.a2a_tool_manager.get_tools()
+            self.resolved_tools.extend(a2a_tools)
+            logger.info(f"Agent '{self.name}': Added {len(a2a_tools)} A2A tools")
+
+        if self.delegation_tool_manager:
+            delegation_tools = self.delegation_tool_manager.get_tools()
+            self.resolved_tools.extend(delegation_tools)
+            logger.info(f"Agent '{self.name}': Added {len(delegation_tools)} delegation tools")
+
+        if self.mcp_tools:
+            self.resolved_tools.extend(self.mcp_tools)
+            logger.info(f"Agent '{self.name}': Added {len(self.mcp_tools)} MCP tools")
+
+        logger.info(f"Agent '{self.name}': Rebuilt resolved tools: {len(self.resolved_tools)} total tools")
+
+    def _rebuild_graph(self) -> None:
+        """Rebuilds and recompiles the graph using the current set of tools.
+
+        Raises:
+            RuntimeError: If the graph rebuilding or compilation process fails.
+        """
+        try:
+            self._rebuild_resolved_tools()
+            self._compiled_graph = self._build_and_compile_graph()
+            logger.info(f"Agent '{self.name}': Successfully rebuilt graph")
+        except Exception as e:
+            logger.error(f"Agent '{self.name}': Failed to rebuild graph: {e}")
+            raise RuntimeError(f"Failed to rebuild graph for agent '{self.name}': {e}") from e
+
+    def run(self, query: str, **kwargs: Any) -> dict[str, Any]:
+        """Synchronously run the LangGraph agent.
+
+        Args:
+            query: The input query for the agent.
+            **kwargs: Additional keyword arguments.
+
+        Returns:
+            Dictionary containing the agent's response.
+        """
+        try:
+            return asyncio.run(self.arun(query, **kwargs))
+        except RuntimeError as e:
+            raise RuntimeError(f"Agent '{self.name}': Error in sync 'run'. Original: {e}") from e
+
+    async def _auto_cleanup(self, entrypoint: str) -> None:
+        """Best-effort auto-cleanup with stable warning marker.
+
+        This internal helper performs cleanup with the following guarantees:
+        - Cleanup exceptions are logged and suppressed (never raised)
+        - Cleanup failures are logged at WARNING level with stable marker format
+        - The marker format is: 'auto-cleanup failed entrypoint=<method>'
+        - Cancellation-safe: uses bounded 1.0s timeout via asyncio.wait_for
+          so cleanup does not hang indefinitely; CancelledError is re-raised
+
+        Args:
+            entrypoint: The public method name that initiated the cleanup (e.g., 'arun').
+        """
+        try:
+            await asyncio.wait_for(self.cleanup(), timeout=_AUTO_CLEANUP_TIMEOUT_S)
+        except asyncio.CancelledError:
+            logger.warning(f"Agent '{self.name}': auto-cleanup failed entrypoint={entrypoint}: cleanup cancelled")
+            raise
+        except TimeoutError:
+            logger.warning(
+                f"Agent '{self.name}': auto-cleanup failed entrypoint={entrypoint}: cleanup cancelled or timed out"
+            )
+        except Exception as e:
+            logger.warning(f"Agent '{self.name}': auto-cleanup failed entrypoint={entrypoint}: {e}")
+
+    async def arun(self, query: str, **kwargs: Any) -> dict[str, Any]:
+        """Asynchronously run the LangGraph agent with lazy MCP initialization.
+
+        This method automatically handles cleanup after execution completes,
+        fails, or is cancelled. Cleanup is best-effort and any cleanup errors
+        are logged but do not override the primary result/error.
+
+        Args:
+            query: The input query for the agent.
+            **kwargs: Additional keyword arguments including configurable for LangGraph.
+
+        Returns:
+            Dictionary containing the agent's response and full final state.
+        """
+        try:
+            await self._ensure_mcp_tools_initialized()
+            return await self._arun(query, **kwargs)
+        finally:
+            await self._auto_cleanup("arun")
+
+    async def _arun(self, query: str, **kwargs: Any) -> dict[str, Any]:
+        """Internal implementation of arun without MCP handling.
+
+        Args:
+            query: The input query for the agent.
+            **kwargs: Additional keyword arguments including configurable for LangGraph.
+
+        Returns:
+            Dictionary containing the agent's response and full final state.
+        """
+        files = kwargs.pop("files", [])
+        if files is None:
+            files = []
+        query, input_attachments = self._prepare_file_runtime_context(query, files, kwargs)
+
+        # Create config first to ensure thread_id is generated
+        config = self._create_graph_config(**kwargs)
+        thread_id = self._get_thread_id_from_config(config)
+
+        graph_input = self._prepare_graph_input(query, thread_id=thread_id, **kwargs)
+        self._attach_input_attachments_to_graph_input(graph_input, input_attachments)
+
+        try:
+            before_updates = await self._middleware_manager.abefore_run(graph_input, config)
+            if before_updates:
+                self._apply_before_run_updates(graph_input, before_updates)
+        except SkillInstallError:
+            raise
+        except Exception:  # pragma: no cover
+            # Best-effort: middleware hooks must not block agent execution.
+            pass  # pragma: no cover
+
+        try:
+            final_state_result = await self._graph_runtime.invoke(self._compiled_graph, graph_input, config)
+            formatted_output = self._format_graph_output(final_state_result)
+
+            try:
+                await self._middleware_manager.aafter_run(
+                    final_state=final_state_result,
+                    output=formatted_output,
+                    config=config,
+                    error=None,
+                )
+            except Exception:  # pragma: no cover
+                # Best-effort: middleware hooks must not block agent execution.
+                pass  # pragma: no cover
+
+            return {"output": formatted_output, "full_final_state": final_state_result}
+
+        except Exception as e:
+            try:
+                await self._middleware_manager.aafter_run(
+                    final_state={},
+                    output=None,
+                    config=config,
+                    error=e,
+                )
+            except Exception:  # pragma: no cover
+                pass  # pragma: no cover
+            logger.error(f"Agent '{self.name}': Error during graph execution: {e}")
+            raise RuntimeError(f"Agent '{self.name}': Graph execution failed: {e}") from e
+
+    @staticmethod
+    def _apply_before_run_updates(graph_input: dict[str, Any], before_updates: dict[str, Any]) -> None:
+        """Merge run-level before updates while preserving existing message history."""
+        message_updates = before_updates.get("messages")
+        if isinstance(message_updates, list) and message_updates:
+            existing_messages = graph_input.get("messages")
+            if isinstance(existing_messages, list):
+                graph_input["messages"] = [*existing_messages, *message_updates]
+            else:
+                graph_input["messages"] = list(message_updates)
+
+        non_message_updates = {k: v for k, v in before_updates.items() if k != "messages"}
+        if non_message_updates:
+            graph_input.update(non_message_updates)
+
+        source_errors = graph_input.get("skill_source_errors")
+        if isinstance(source_errors, list) and source_errors:
+            raise SkillInstallError(
+                "Skill source readiness failed before model execution.", source_errors=source_errors
+            )
+
+    def _prepare_file_runtime_context(
+        self,
+        query: str,
+        files: list[str | dict[str, object]],
+        kwargs: dict[str, Any],
+    ) -> tuple[str, list[dict[str, Any]]]:
+        """Augment query and metadata with safe image attachment references."""
+        augmented_query = augment_query_with_file_paths(query=query, files=files)
+        registry = build_input_attachment_registry(files)
+        attachments = registry.as_list()
+        if not attachments:
+            return augmented_query, []
+
+        kwargs["metadata"] = merge_input_attachments(kwargs.get("metadata"), attachments)
+        return augmented_query, attachments
+
+    @staticmethod
+    def _attach_input_attachments_to_graph_input(
+        graph_input: dict[str, Any], input_attachments: list[dict[str, Any]]
+    ) -> None:
+        """Attach input image metadata to runtime state when present."""
+        attach_attachments_to_state(graph_input, input_attachments)
+
+    async def _stream_with_lm_invoker(self, query: str, **kwargs: Any) -> AsyncGenerator[str | dict[str, Any], None]:
+        """Handle streaming for LM Invoker using StreamEventHandler.
+
+        Args:
+            query: The input query for the agent.
+            **kwargs: Additional keyword arguments.
+
+        Yields:
+            Chunks of output (strings or dicts) from the streaming response.
+        """
+        stream_handler = StreamEventHandler(name=f"{self.name}_StreamHandler")
+        event_emitter = EventEmitter(handlers=[stream_handler])
+
+        async def run_and_cleanup() -> None:
+            """Runs the agent and ensures event emitter cleanup."""
+            try:
+                await self.arun(
+                    query=query,
+                    event_emitter=event_emitter,
+                    **kwargs,
+                )
+            finally:
+                await event_emitter.close()
+
+        execution_task = asyncio.create_task(run_and_cleanup())
+
+        try:
+            async for event in stream_handler.stream():
+                chunk_data = json.loads(event)
+                chunk_value = chunk_data.get("value", "")
+                if not chunk_value:
+                    continue
+                if isinstance(chunk_value, str) or isinstance(chunk_value, dict):
+                    yield chunk_value
+
+            await execution_task
+
+        except asyncio.CancelledError:
+            execution_task.cancel()
+            await event_emitter.close()
+            with suppress(asyncio.CancelledError):
+                await execution_task
+            raise
+        except Exception as e:
+            execution_task.cancel()
+            await event_emitter.close()
+            with suppress(asyncio.CancelledError):
+                await execution_task
+            logger.error(f"Agent '{self.name}': Error during LM Invoker streaming: {e}")
+            yield {"error": f"Streaming failed: {e}"}
+
+    @property
+    def _effective_recursion_limit(self) -> int:
+        """Return the recursion limit to apply to every graph invocation.
+
+        When ``self.recursion_limit`` is ``None``, returns
+        ``StepLimitConfig().max_steps * 2`` (200 by default) so that LangGraph's
+        hard ``GraphRecursionError`` never fires before the AIP graceful error path.
+
+        Subclasses may override this to derive the limit from additional runtime
+        parameters (e.g. ``step_limit_config.max_steps``).
+
+        Returns:
+            The recursion limit to apply.  Always returns a positive integer.
+        """
+        if self.recursion_limit is not None:
+            return self.recursion_limit
+        return StepLimitConfig().max_steps * 2
+
+    def _resolve_configurable(self, kwargs: dict[str, Any]) -> dict[str, Any]:
+        """Return the configurable dict with a guaranteed thread ID entry.
+
+        Copies the incoming ``configurable`` value so the caller's dict is never
+        mutated. Generates and logs a new thread ID when the key is absent.
+
+        Args:
+            kwargs: The raw keyword arguments passed to ``_create_graph_config``.
+
+        Returns:
+            Configurable dict with a thread ID guaranteed to be present.
+        """
+        return self._graph_runtime.resolve_configurable(kwargs)  # pragma: no cover
+
+    def _build_config_metadata(self, kwargs: dict[str, Any]) -> dict[str, Any] | None:
+        """Build the ``metadata`` entry from kwargs and graph_config baseline.
+
+        Per the spec, per-call ``metadata`` **replaces** the ``graph_config`` baseline
+        entirely for that invocation. Internal keys (``pii_mapping``, ``memory_user_id``)
+        are then applied on top and always win.
+
+        Returns ``None`` when no metadata sources are present.
+
+        Args:
+            kwargs: The raw keyword arguments passed to ``_create_graph_config``.
+
+        Returns:
+            Merged metadata dict, or ``None`` if no metadata-related sources.
+        """
+        return self._graph_runtime.build_config_metadata(kwargs)  # pragma: no cover
+
+    def _apply_config_layers(self, config: dict[str, Any], kwargs: dict[str, Any]) -> None:
+        """Apply construction-time and per-call config layers onto ``config`` in place.
+
+        Layers are applied in priority order (each layer wins over the previous):
+
+        - **Layer 1** — ``self.graph_config`` baseline (excludes ``configurable``).
+        - **Layer 2** — ``_effective_recursion_limit``; only applied when
+          ``self.recursion_limit`` is explicitly set (explicit param wins over
+          ``graph_config``).  When not explicitly set, the auto-derived value
+          (``step_limit_config.max_steps * 2``, 200 if no ``step_limit_config``) is
+          written only when ``graph_config`` did not already supply
+          ``"recursion_limit"``.
+        - **Layer 3** — per-call ``_RUNNABLE_CONFIG_KEYS`` kwargs; ``metadata`` is
+          intentionally skipped because ``_build_config_metadata`` handles it with
+          internal-key guards that a plain overwrite would silently discard.
+
+        Args:
+            config: The partially-assembled config dict; mutated in place.
+            kwargs: The raw keyword arguments passed to ``_create_graph_config``.
+        """
+        self._graph_runtime.apply_config_layers(config, kwargs)  # pragma: no cover
+
+    def _create_graph_config(self, **kwargs: Any) -> dict[str, Any]:
+        """Assemble the final ``RunnableConfig``-compatible dict for a graph invocation.
+
+        Args:
+            **kwargs: Per-call overrides; accepted keys are ``configurable``,
+                ``metadata``, ``pii_mapping``, ``memory_user_id``, and any key in
+                ``_RUNNABLE_CONFIG_KEYS``.
+
+        Returns:
+            Config dict with a guaranteed thread ID and all applicable layers merged.
+        """
+        return self._graph_runtime.create_graph_config(**kwargs)
+
+    def _get_thread_id_from_config(self, config: dict[str, Any]) -> str | None:
+        """Extract thread_id from graph configuration.
+
+        Args:
+            config: Graph configuration dict with 'configurable' key.
+
+        Returns:
+            The thread_id value or None if not found.
+        """
+        return self._graph_runtime.get_thread_id_from_config(config)
+
+    def _process_langgraph_event(self, event: Any) -> str | dict[str, Any] | A2AEvent | None:
+        """Process a single LangGraph streaming event.
+
+        Args:
+            event: Event from LangGraph's astream_events.
+
+        Returns:
+            Processed output or None if event should be skipped.
+        """
+        event_type = event.get("event")
+        event_data = event.get("data")
+
+        if event_type == "on_chat_model_stream" and event_data:
+            chunk = event_data.get("chunk")
+            if chunk and hasattr(chunk, "content") and chunk.content:
+                return chunk.content
+
+        elif event_type == "on_tool_end" and event_data:
+            output = event_data.get("output")
+            if output:
+                return {"tool_output": str(output)}
+
+        return None
+
+    def _should_yield_a2a_event(self, event_data: A2AEvent) -> bool:
+        """Check if A2A event should be yielded based on event type.
+
+        Args:
+            event_data: A2AEvent with semantic type information.
+
+        Returns:
+            True if event should be yielded, False otherwise.
+        """
+        event_type = event_data.get("event_type")
+
+        if event_type in {
+            A2AStreamEventType.TOOL_CALL,
+            A2AStreamEventType.TOOL_RESULT,
+            A2AStreamEventType.CONTENT_CHUNK,
+            A2AStreamEventType.FINAL_RESPONSE,
+            A2AStreamEventType.ERROR,
+        }:
+            return True
+
+        if event_type == A2AStreamEventType.STATUS_UPDATE:
+            content = event_data.get("content", "")
+            return bool(content.strip())
+
+        return True
+
+    @staticmethod
+    def _normalize_custom_stream_event(event: A2AEvent) -> A2AEvent:
+        """Normalize CUSTOM stream mode events that lack required A2AEvent fields.
+
+        CUSTOM mode events bypass the _create_a2a_event factory and may arrive as
+        raw dicts without event_type or metadata. This method wraps those events
+        into a valid STATUS_UPDATE event with the available content preserved.
+
+        Args:
+            event: Incoming CUSTOM stream event, potentially missing A2AEvent fields.
+
+        Returns:
+            Normalized A2AEvent with guaranteed event_type and metadata fields.
+        """
+        normalized_event = dict(event)
+        existing_metadata = normalized_event.get("metadata")
+        metadata = dict(existing_metadata) if isinstance(existing_metadata, dict) else {}
+
+        if "event_type" in normalized_event:
+            normalized_event["metadata"] = metadata
+            return normalized_event
+
+        logger.warning(
+            "CUSTOM stream event missing A2AEvent fields (has keys: %s); wrapping as STATUS_UPDATE",
+            list(normalized_event.keys()),
+        )
+
+        metadata["original_event_keys"] = list(normalized_event.keys())
+
+        content = normalized_event.get("content") or normalized_event.get("status") or str(normalized_event)
+        return {
+            "event_type": A2AStreamEventType.STATUS_UPDATE,
+            "content": BaseLangGraphAgent._normalize_event_content(content),
+            "metadata": metadata,
+            "tool_info": normalized_event.get("tool_info"),
+            "is_final": normalized_event.get("is_final", False),
+            "artifacts": normalized_event.get("artifacts"),
+            "references": normalized_event.get("references"),
+            "step_usage": normalized_event.get("step_usage"),
+            "total_usage": normalized_event.get("total_usage"),
+            "model_usage_delta": normalized_event.get("model_usage_delta"),
+            "model_usage": normalized_event.get("model_usage"),
+            "thinking_and_activity_info": normalized_event.get("thinking_and_activity_info"),
+        }
+
+    @staticmethod
+    def _normalize_usage_metadata(usage: Any | None) -> dict[str, Any] | None:
+        """Normalize usage metadata to a dictionary when possible.
+
+        Args:
+            usage: Usage metadata from LangChain messages.
+
+        Returns:
+            A dictionary copy when usage is available, otherwise None.
+        """
+        if usage is None:
+            return None
+        if isinstance(usage, dict):
+            return dict(usage)
+        return cast(dict[str, Any], usage)
+
+    @staticmethod
+    def _normalize_event_content(content: Any) -> str:
+        """Normalize event content to a string payload.
+
+        Args:
+            content: Raw content payload from LangChain/LangGraph.
+
+        Returns:
+            String representation suitable for A2A events.
+        """
+        if isinstance(content, str):
+            return content
+        try:
+            return json.dumps(content)
+        except (TypeError, ValueError):
+            return str(content)
+
+    async def _stream_with_langgraph(self, query: str, **kwargs: Any) -> AsyncGenerator[str | dict[str, Any], None]:
+        """Handle streaming for LangChain models using LangGraph's native streaming.
+
+        Args:
+            query: The input query for the agent.
+            **kwargs: Additional keyword arguments.
+
+        Yields:
+            Chunks of output (strings or dicts) from the streaming response.
+        """
+        # Create config first to ensure thread_id is generated
+        config = self._create_graph_config(**kwargs)
+        thread_id = self._get_thread_id_from_config(config)
+
+        graph_input = self._prepare_graph_input(query, thread_id=thread_id, **kwargs)
+
+        try:
+            async for event in self._graph_runtime.stream_events(self._compiled_graph, graph_input, config):
+                processed_output = self._process_langgraph_event(event)
+                if processed_output is not None:
+                    yield processed_output
+
+        except Exception as e:
+            logger.error(f"Agent '{self.name}': Error during graph streaming: {e}")
+            yield {"error": f"Streaming failed: {e}"}
+
+    async def arun_stream(self, query: str, **kwargs: Any) -> AsyncGenerator[str | dict[str, Any], None]:
+        """Asynchronously stream the LangGraph agent's response.
+
+        If MCP configuration exists, connects to the MCP server and registers tools before streaming.
+        This method properly handles both LM Invoker and LangChain model streaming:
+        - For LM Invoker: Uses StreamEventHandler to capture streaming events
+        - For LangChain models: Uses LangGraph's native streaming implementation
+
+        This method automatically handles cleanup after streaming completes,
+        fails, is cancelled, or the consumer stops early (break/aclose).
+
+        Args:
+            query: The input query for the agent.
+            **kwargs: Additional keyword arguments.
+
+        Yields:
+            Chunks of output (strings or dicts) from the streaming response.
+        """
+        try:
+            await self._ensure_mcp_tools_initialized()
+            async for chunk in self._arun_stream(query, **kwargs):
+                yield chunk
+        finally:
+            await self._auto_cleanup("arun_stream")
+
+    async def _arun_stream(self, query: str, **kwargs: Any) -> AsyncGenerator[str | dict[str, Any], None]:
+        """Internal implementation of arun_stream without MCP handling.
+
+        This method properly handles both LM Invoker and LangChain model streaming:
+        - For LM Invoker: Uses StreamEventHandler to capture streaming events
+        - For LangChain models: Uses LangGraph's native streaming implementation
+
+        Args:
+            query: The input query for the agent.
+            **kwargs: Additional keyword arguments.
+
+        Yields:
+            Chunks of output (strings or dicts) from the streaming response.
+        """
+        if self._has_lm_invoker():
+            async for chunk in self._stream_with_lm_invoker(query, **kwargs):
+                yield chunk
+        else:
+            async for chunk in self._stream_with_langgraph(query, **kwargs):
+                yield chunk
+
+    def _initialize_mcp_client(self) -> None:
+        """Initialize/recreate MCP client with current config safely disposing previous.
+
+        This method creates a new LangchainMCPClient if MCP configuration exists,
+        and safely disposes of any existing client before setting the new one.
+        """
+        new_client = LangchainMCPClient(self.mcp_config) if self.mcp_config else None
+        self._set_mcp_client_safely(new_client)
+
+    async def _register_mcp_tools(self) -> None:
+        """Initialize MCP tools once during agent setup using persistent sessions.
+
+        This method connects to MCP servers, retrieves available tools, and integrates
+        them into the agent's tool collection. It includes timeout handling to prevent
+        hanging operations.
+
+        Raises:
+            RuntimeError: If MCP initialization times out after 30 seconds.
+            Exception: If MCP tool initialization fails for other reasons.
+        """
+        try:
+            logger.info(f"Agent '{self.name}': Initializing MCP tools with persistent sessions.")
+
+            # Add timeout for initialization to prevent hanging
+            mcp_client = self.mcp_client
+            if mcp_client is None:
+                return
+
+            await asyncio.wait_for(mcp_client.initialize(), timeout=MCP_INIT_TIMEOUT)
+
+            mcp_tools = await mcp_client.get_tools()
+
+            if not mcp_tools:
+                logger.warning(f"Agent '{self.name}': No MCP tools retrieved from configured servers.")
+                return
+
+            self.mcp_tools.extend(mcp_tools)
+            logger.info(f"Agent '{self.name}': Added {len(mcp_tools)} persistent MCP tools to graph.")
+            self._rebuild_graph()
+
+        except TimeoutError as err:
+            logger.error(f"Agent '{self.name}': MCP initialization timed out")
+            raise RuntimeError(f"Agent '{self.name}': MCP initialization timed out after 30 seconds") from err
+        except MCPServerStartupError:
+            raise
+        except Exception as e:
+            logger.error(f"Agent '{self.name}': Failed to initialize persistent MCP tools: {e}", exc_info=True)
+            raise
+
+    async def cleanup(self) -> None:
+        """Cleanup MCP resources including persistent sessions.
+
+        This method is called automatically after each ``arun``, ``arun_stream``,
+        and ``arun_a2a_stream`` invocation (including on error or cancellation).
+        Explicit calls are safe but redundant for one-shot usage patterns.
+
+        This method performs best-effort cleanup of MCP client resources
+        and resets lifecycle state so the next run can lazily re-initialize.
+
+        Post-cleanup invariants:
+        - _mcp_tools_initialized is False (allows lazy re-init on next call)
+        - mcp_tools list is cleared (prevents duplicate accumulation)
+        - resolved_tools and compiled graph are rebuilt without stale MCP tools
+
+        Errors during cleanup are logged but do not raise exceptions to ensure
+        the cleanup process completes gracefully. This method is idempotent.
+        """
+        if hasattr(self, "mcp_client") and self.mcp_client:
+            try:
+                await self.mcp_client.cleanup()
+                logger.debug(f"Agent '{self.name}': MCP client cleanup completed")
+            except Exception as e:
+                logger.warning(f"Agent '{self.name}': Error during MCP client cleanup: {e}")
+                # Don't re-raise - cleanup should be best-effort
+
+        # Reset MCP lifecycle state for safe re-initialization on next run
+        self._mcp_tools_initialized = False
+        if hasattr(self, "mcp_tools") and self.mcp_tools:
+            self.mcp_tools.clear()
+            try:
+                self._rebuild_resolved_tools()
+                self._compiled_graph = self._build_and_compile_graph()
+            except Exception as e:
+                logger.warning(f"Agent '{self.name}': Error rebuilding tools/graph during cleanup: {e}")
+                # Don't re-raise - cleanup should be best-effort
+
+    async def arun_a2a_stream(self, query: str, **kwargs: Any) -> AsyncGenerator[A2AEvent, None]:
+        """Asynchronously streams the agent's response in A2A format.
+
+        This method automatically handles cleanup after streaming completes,
+        fails, is cancelled, or the consumer stops early (break/aclose).
+
+        Args:
+            query: The input query for the agent.
+            **kwargs: Additional keyword arguments.
+
+        Yields:
+            Dictionaries with "status" and "content" keys.
+            Possible statuses: "working", "completed", "failed", "canceled".
+        """
+        try:
+            await self._ensure_mcp_tools_initialized()
+            async for chunk in self._arun_a2a_stream(query, **kwargs):
+                yield chunk
+        finally:
+            await self._auto_cleanup("arun_a2a_stream")
+
+    async def arun_sse_stream(
+        self,
+        query: str,
+        task_id: str | None = None,
+        context_id: str | None = None,
+        **kwargs: Any,
+    ) -> AsyncGenerator[A2AEvent, None]:
+        """Stream agent response as SSE-compatible chunks.
+
+        This method wraps arun_a2a_stream and transforms output to the normalized
+        dict format matching A2AConnector.astream_to_agent output, enabling direct
+        streaming without A2A server overhead.
+
+        Args:
+            query: The input query for the agent.
+            task_id: Optional task identifier for the stream.
+            context_id: Optional context identifier for the stream.
+            **kwargs: Additional arguments passed to arun_a2a_stream.
+
+        Yields:
+            SSEChunk dicts with normalized structure:
+            - status: "success" | "error"
+            - task_state: "working" | "completed" | "failed" | "canceled"
+            - content: Text content or None
+            - event_type: Always string (never enum)
+            - final: True for terminal events
+            - metadata: Normalized metadata dict
+            - artifacts: Only present when non-empty
+        """
+        if task_id is None:
+            task_id = str(uuid.uuid4())
+        if context_id is None:
+            context_id = str(uuid.uuid4())
+
+        # Extract pii_mapping from kwargs to pass to transformer (matching A2A executor behavior)
+        pii_mapping = kwargs.get("pii_mapping")
+        transformer = SSEChunkTransformer(task_id=task_id, context_id=context_id, pii_mapping=pii_mapping)
+        try:
+            stream = self.arun_a2a_stream(query, **kwargs)
+            async for chunk in transformer.transform_stream(stream):
+                yield chunk
+        except Exception as e:
+            # Lazy import to support optional guardrails dependency
+            from aip_agents.guardrails.exceptions import GuardrailViolationError
+
+            if isinstance(e, GuardrailViolationError):
+                # Re-raise guardrail violations without modification
+                raise
+            logger.error(f"Agent '{self.name}': Error in arun_sse_stream: {e}", exc_info=True)
+            yield SSEChunkTransformer._create_error_chunk(f"Error during streaming: {e}")
+
+    def _get_tool_processing_content(self, tool_names: list[str]) -> str:
+        """Generate appropriate content prefix for tool processing messages.
+
+        Args:
+            tool_names: List of tool names to process.
+
+        Returns:
+            Formatted content string with appropriate prefix.
+        """
+        unique_tool_names = list(dict.fromkeys(tool_names))
+        has_delegation_tools = any(name.startswith("delegate_to") for name in unique_tool_names)
+        content_prefix = "Processing with sub-agents:" if has_delegation_tools else "Processing with tools:"
+        return f"{content_prefix} {', '.join(unique_tool_names)}"
+
+    def _get_tool_completion_content(self, tool_names: list[str]) -> str:
+        """Generate completion message for finished tool executions.
+
+        Args:
+            tool_names: List of tool names to summarize.
+
+        Returns:
+            Content string indicating completion.
+        """
+        unique_tool_names = list(dict.fromkeys(tool_names))
+        has_delegation_tools = any(name.startswith("delegate_to") for name in unique_tool_names)
+        content_prefix = "Completed sub-agents:" if has_delegation_tools else "Completed tools:"
+        return f"{content_prefix} {', '.join(unique_tool_names)}"
+
+    def _parse_a2a_stream_message(
+        self, message: BaseMessage, state: dict[str, Any] | None = None
+    ) -> tuple[A2AEvent | None, bool]:
+        """Parse LangChain messages into semantically meaningful A2A events.
+
+        This method converts LangChain message types (AIMessage, ToolMessage) into
+        structured A2AEvent objects that preserve semantic information and eliminate
+        the need for string parsing downstream.
+
+        Args:
+            message: The LangChain message to parse (AIMessage, ToolMessage, etc.).
+            state: Optional state dictionary containing pii_mapping and other data.
+
+        Returns:
+            A tuple containing:
+            - A2AEvent | None: The parsed A2A event, or None if message should be skipped.
+            - bool: True if this is a final event that should terminate the stream.
+        """
+        if isinstance(message, AIMessage) and message.tool_calls:
+            return self._create_tool_call_event(message), False
+
+        elif isinstance(message, ToolMessage):
+            return self._create_tool_result_event(message), False
+
+        elif isinstance(message, AIMessage) and message.content:
+            return self._create_ai_message_event(message, state)
+
+        return None, False
+
+    def _link_tool_call_to_previous_status(self, event: A2AEvent) -> None:
+        """Link the tool call event to completed tool steps or the most recent status step.
+
+        Supports both parallel and sequential operation modes:
+        - "parallel": Links to ALL completed tool steps (default for backward compatibility)
+        - "sequential": Links to only the most recent completed tool step
+
+        Args:
+            event: The A2AEvent to link to previous step.
+        """
+        try:
+            thread_id = _THREAD_ID_CVAR.get()
+            if thread_id:
+                metadata = event.get("metadata", {})
+                existing_step_ids = metadata.get("previous_step_ids") or []
+
+                if existing_step_ids:
+                    return
+
+                operation_mode = _OPERATION_MODE_CVAR.get() or "parallel"
+
+                coord_completed_steps = self._coordinator_completed_tool_steps_by_thread.get(thread_id, [])
+                completed_steps = coord_completed_steps or self._completed_tool_steps_by_thread.get(thread_id, [])
+                if completed_steps:
+                    if operation_mode == "sequential":
+                        metadata["previous_step_ids"] = [completed_steps[-1]]
+                    else:
+                        metadata["previous_step_ids"] = completed_steps
+                    event["metadata"] = metadata
+                    return
+
+                last_status_id = self._last_status_step_id_by_thread.get(thread_id)
+                if last_status_id:
+                    metadata["previous_step_ids"] = [last_status_id]
+                    event["metadata"] = metadata
+        except Exception as e:
+            logger.warning("Failed linking tool call to previous step: %s", e, exc_info=True)
+
+    def _register_tool_call_parent_steps(self, event: A2AEvent, tool_calls_details: list[dict]) -> None:
+        """Register parent step IDs for each tool call ID.
+
+        Args:
+            event: The A2AEvent containing the parent step.
+            tool_calls_details: List of tool call details.
+        """
+        try:
+            thread_id = _THREAD_ID_CVAR.get()
+            if thread_id:
+                parent_step_id = event["metadata"].get("step_id")
+                if parent_step_id:
+                    parent_map = self._tool_parent_map_by_thread.setdefault(thread_id, {})
+                    for tool_call in tool_calls_details:
+                        tool_call_id = tool_call.get("id")
+                        if tool_call_id:
+                            parent_map[str(tool_call_id)] = str(parent_step_id)
+        except Exception as e:
+            logger.warning("Registering tool call parent steps failed: %s", e, exc_info=True)
+
+    def _create_tool_call_event(self, message: AIMessage) -> A2AEvent:
+        """Create an A2AEvent for tool invocation from AIMessage.
+
+        Args:
+            message: AIMessage containing tool calls.
+
+        Returns:
+            A2AEvent with TOOL_CALL event type and structured tool information.
+        """
+        tool_calls_details: list[dict[str, Any]] = []
+        manager = getattr(self, "tool_output_manager", None)
+        thread_id = _THREAD_ID_CVAR.get()
+        for tool_call in message.tool_calls:
+            args = tool_call["args"]
+            if manager and thread_id and isinstance(args, dict):
+                args = manager.rewrite_args_with_latest_reference(args, thread_id)
+            tool_calls_details.append(
+                {
+                    "id": tool_call.get("id"),
+                    "name": tool_call["name"],
+                    "args": args,
+                }
+            )
+        tool_names = [details["name"] for details in tool_calls_details]
+
+        event = self._create_a2a_event(
+            event_type=A2AStreamEventType.TOOL_CALL,
+            content=self._get_tool_processing_content(tool_names),
+            tool_info={"tool_calls": tool_calls_details, "status": "running"},
+            metadata={"status": Status.RUNNING},
+            is_final=False,
+            step_usage=self._normalize_usage_metadata(message.usage_metadata),
+        )
+
+        self._record_emitted_tool_calls(tool_calls_details)
+
+        self._link_tool_call_to_previous_status(event)
+        self._register_tool_call_parent_steps(event, tool_calls_details)
+
+        return event
+
+    def _get_sub_agent_previous_steps(self, message: ToolMessage) -> list[str] | None:
+        """Extract previous step IDs from sub-agent response metadata.
+
+        Args:
+            message: ToolMessage containing response metadata.
+
+        Returns:
+            List of previous step IDs or None if not available.
+        """
+        try:
+            if not hasattr(message, "response_metadata") or not isinstance(message.response_metadata, dict):
+                return None
+
+            sub_prev = message.response_metadata.get("previous_step_ids")
+            if isinstance(sub_prev, list) and sub_prev:
+                return [str(x) for x in sub_prev if isinstance(x, str | int)]
+            return None
+        except Exception as e:
+            logger.warning("Failed extracting sub-agent previous steps: %s", e, exc_info=True)
+            return None
+
+    def _determine_previous_step_ids(self, message: ToolMessage, sub_prev: list[str] | None) -> list[str]:
+        """Determine which previous step IDs to use for the event.
+
+        Args:
+            message: ToolMessage for the tool call.
+            sub_prev: Previous step IDs from sub-agent, if available.
+
+        Returns:
+            List of previous step IDs to use.
+        """
+        if sub_prev:
+            return sub_prev
+
+        try:
+            thread_id = _THREAD_ID_CVAR.get()
+            if thread_id:
+                parent_map = self._tool_parent_map_by_thread.get(thread_id, {})
+                parent_step = parent_map.get(str(message.tool_call_id))
+                if parent_step:
+                    return [parent_step]
+        except Exception as e:
+            logger.warning("Determining previous step IDs failed: %s", e, exc_info=True)
+
+        return []
+
+    def _record_tool_completion(self, message: ToolMessage, event: A2AEvent) -> None:
+        """Record tool completion for final event dependency tracking.
+
+        Args:
+            message: ToolMessage for the completed tool.
+            event: The A2AEvent for the tool result.
+        """
+        try:
+            thread_id = _THREAD_ID_CVAR.get()
+            if not thread_id:
+                return
+
+            completed_list = self._completed_tool_steps_by_thread.setdefault(thread_id, [])
+            coord_completed_list = self._coordinator_completed_tool_steps_by_thread.setdefault(thread_id, [])
+
+            event_sid = (event.get("metadata") or {}).get("step_id")
+            if isinstance(event_sid, str) and event_sid:
+                completed_list.append(event_sid)
+                coord_completed_list.append(event_sid)
+
+            sub_prev = self._get_sub_agent_previous_steps(message) or []
+            completed_list.extend(step_id for step_id in sub_prev if isinstance(step_id, str) and step_id)
+
+            self._completed_tool_steps_by_thread[thread_id] = list(dict.fromkeys(completed_list))
+            self._coordinator_completed_tool_steps_by_thread[thread_id] = list(dict.fromkeys(coord_completed_list))
+        except Exception as e:
+            logger.warning("Recording tool completion failed: %s", e, exc_info=True)
+
+    def _create_tool_result_event(self, message: ToolMessage) -> A2AEvent:
+        """Create an A2AEvent for tool completion from ToolMessage.
+
+        Args:
+            message: ToolMessage containing tool execution results.
+
+        Returns:
+            A2AEvent with TOOL_RESULT event type and execution details.
+        """
+        tool_info = self._extract_tool_info_from_message(message)
+        previous_ids = self._determine_previous_step_ids(
+            message,
+            self._get_sub_agent_previous_steps(message),
+        )
+
+        event = self._create_a2a_event(
+            event_type=A2AStreamEventType.TOOL_RESULT,
+            content=self._build_tool_event_content(tool_info["name"], tool_info["output"], message),
+            tool_info={
+                "id": message.tool_call_id,
+                "name": tool_info["name"],
+                "args": tool_info["args"],
+                "output": tool_info["output"],
+                "execution_time": tool_info["execution_time"],
+            },
+            metadata=self._build_tool_event_metadata(tool_info["execution_time"], previous_ids),
+            is_final=False,
+            step_usage=message.response_metadata.get(USAGE_METADATA_KEY),
+            model_usage_delta=message.response_metadata.get(MODEL_USAGE_DELTA_KEY),
+        )
+
+        self._propagate_hitl_metadata(message, event)
+        self._propagate_tool_execution_metadata(message, event)
+        self._record_tool_completion(message, event)
+        self._discard_emitted_tool_call(getattr(message, "tool_call_id", None))
+
+        return event
+
+    def _extract_tool_info_from_message(self, message: ToolMessage) -> dict[str, Any]:
+        """Extract tool details from a ToolMessage.
+
+        Args:
+            message: The ToolMessage to extract information from.
+
+        Returns:
+            Dictionary containing tool name, args, output, and execution time.
+        """
+        tool_call_info = getattr(message, "tool_calls", {})
+        tool_name = getattr(message, "name", None) or tool_call_info.get("name", "unknown")
+        return {
+            "name": tool_name,
+            "args": tool_call_info.get("args", {}),
+            "output": tool_call_info.get("output", message.content),
+            "execution_time": tool_call_info.get("time"),
+        }
+
+    def _build_tool_event_content(self, tool_name: str, tool_output: Any, message: ToolMessage) -> str:
+        """Determine event content for a tool result.
+
+        Args:
+            tool_name: Name of the tool that was executed.
+            tool_output: The output returned by the tool.
+            message: The ToolMessage containing response metadata and tool call information.
+
+        Returns:
+            String content for the tool result event.
+        """
+        response_metadata = getattr(message, "response_metadata", None) or {}
+        hitl_meta = response_metadata.get(MetadataFieldKeys.HITL) if isinstance(response_metadata, dict) else None
+
+        if hitl_meta and hitl_meta.get("required"):
+            return str(tool_output) if tool_output else self._get_tool_processing_content([tool_name])
+
+        return self._get_tool_completion_content([tool_name])
+
+    def _build_tool_event_metadata(
+        self,
+        execution_time: Any,
+        previous_ids: list[str] | None,
+    ) -> dict[str, Any]:
+        """Build metadata payload for tool result events.
+
+        Args:
+            execution_time: Time taken to execute the tool.
+            previous_ids: Optional list of previous step IDs this tool depends on.
+
+        Returns:
+            Dictionary containing status, execution time, and previous step IDs.
+        """
+        return {
+            "status": Status.FINISHED,
+            "time": execution_time,
+            "previous_step_ids": previous_ids,
+        }
+
+    def _propagate_hitl_metadata(self, message: ToolMessage, event: A2AEvent) -> None:
+        """Copy HITL metadata from ToolMessage into the event if available.
+
+        Args:
+            message: The ToolMessage containing response metadata with HITL information.
+            event: The A2AEvent to update with HITL metadata if present.
+        """
+        response_metadata = getattr(message, "response_metadata", None)
+        if not isinstance(response_metadata, dict):
+            return
+
+        hitl_meta = response_metadata.get(MetadataFieldKeys.HITL)
+        if hitl_meta is None:
+            return
+
+        try:
+            hitl_model = HitlMetadata.model_validate(hitl_meta)
+        except ValidationError as exc:
+            raise ValueError("Invalid HITL metadata payload encountered") from exc
+
+        metadata = event.get("metadata")
+        if isinstance(metadata, dict):
+            try:
+                metadata[MetadataFieldKeys.HITL] = hitl_model.as_payload()
+            except Exception as exc:  # noqa: BLE001
+                logger.warning("Failed to propagate HITL metadata to event: %s", exc)
+
+    def _propagate_tool_execution_metadata(self, message: ToolMessage, event: A2AEvent) -> None:
+        """Copy tool execution metadata into TOOL_RESULT event payload when available.
+
+        Args:
+            message: The ToolMessage containing response metadata.
+            event: The A2AEvent to update with tool execution metadata.
+        """
+        response_metadata = getattr(message, "response_metadata", None)
+        if not isinstance(response_metadata, dict):
+            return
+
+        tool_execution = response_metadata.get("tool_execution")
+        if not isinstance(tool_execution, dict) or not tool_execution:
+            return
+
+        metadata = event.get("metadata")
+        if not isinstance(metadata, dict):
+            return
+
+        metadata["tool_execution"] = dict(tool_execution)
+
+    def _should_emit_provider_reasoning(self, message: BaseMessage) -> bool:
+        """Check if an AIMessage has provider reasoning to emit.
+
+        Args:
+            message: The LangChain message to check.
+
+        Returns:
+            True if the message is an AIMessage with serialized thinking content.
+        """
+        return isinstance(message, AIMessage) and bool(message.additional_kwargs.get(THINKINGS_METADATA_KEY))
+
+    def _get_reasoning_previous_step_ids(self) -> list[str]:
+        """Get parent step IDs for a provider reasoning event."""
+        previous_step_ids = self._get_previous_step_ids() or []
+        if previous_step_ids:
+            return previous_step_ids
+
+        thread_id = _THREAD_ID_CVAR.get()
+        if not thread_id:
+            return []  # pragma: no cover
+
+        last_status_id = self._last_status_step_id_by_thread.get(thread_id)
+        return [last_status_id] if last_status_id else []
+
+    def _create_provider_reasoning_event(self, message: AIMessage) -> A2AEvent | None:
+        """Create an A2AEvent for provider reasoning/thinking content.
+
+        Extracts serialized thinking blocks from ``gllm_thinkings`` metadata
+        and emits them as a completed ``agent_thinking_step`` status update
+        using the existing ``thinking_and_activity_info`` contract.
+
+        Args:
+            message: AIMessage containing ``gllm_thinkings`` in additional_kwargs.
+
+        Returns:
+            A2AEvent for the provider reasoning step, or None if no
+            non-empty thinking text could be extracted.
+        """
+        raw_thinkings = message.additional_kwargs.get(THINKINGS_METADATA_KEY)
+        if not isinstance(raw_thinkings, list) or not raw_thinkings:
+            return None
+
+        thinking_texts: list[str] = []
+        for entry in raw_thinkings:
+            if not isinstance(entry, dict):
+                continue
+            text = entry.get("thinking", "") or ""
+            if isinstance(text, str) and text.strip():
+                thinking_texts.append(text)
+
+        combined = "\n\n".join(thinking_texts)
+        if not combined.strip():
+            return None
+
+        content_hash = hashlib.sha256(combined.encode("utf-8")).hexdigest()[:8]
+        step_id = f"provider_reasoning_{content_hash}"
+
+        thinking_and_activity_info = {
+            "id": step_id,
+            "data_type": "activity",
+            "data_value": json.dumps({"type": "reasoning", "message": combined}),
+        }
+
+        return self._create_a2a_event(
+            event_type=A2AStreamEventType.STATUS_UPDATE,
+            content="",
+            tool_info=None,
+            metadata={
+                MetadataFieldKeys.KIND: Kind.AGENT_THINKING_STEP,
+                MetadataFieldKeys.STATUS: Status.FINISHED,
+                "step_id": step_id,
+                "previous_step_ids": self._get_reasoning_previous_step_ids(),
+                MetadataFieldKeys.THINKING_AND_ACTIVITY_INFO: thinking_and_activity_info,
+            },
+            is_final=False,
+        )
+
+    def _create_ai_message_event(
+        self, message: AIMessage, state: dict[str, Any] | None = None
+    ) -> tuple[A2AEvent, bool]:
+        """Create an A2AEvent for AI-generated content from AIMessage.
+
+        Args:
+            message: AIMessage containing AI-generated content.
+            state: Optional state dictionary containing pii_mapping and other data.
+
+        Returns:
+            A tuple containing:
+            - A2AEvent: Either CONTENT_CHUNK or FINAL_RESPONSE event.
+            - bool: True if this is a final response, False for streaming content.
+        """
+        is_final_response = self._is_final_response(message)
+        metadata = self._build_metadata_for_final_response(is_final_response, state)
+        raw_content = message.content
+        content = deanonymize_final_response_content(
+            content=raw_content if isinstance(raw_content, str) else json.dumps(raw_content),
+            is_final_response=is_final_response,
+            metadata=metadata,
+        )
+        event = self._create_a2a_event(
+            event_type=A2AStreamEventType.FINAL_RESPONSE if is_final_response else A2AStreamEventType.CONTENT_CHUNK,
+            content=content,
+            tool_info=None,
+            metadata=metadata,
+            is_final=is_final_response,
+            step_usage=self._normalize_usage_metadata(message.usage_metadata),
+            model_usage_delta=(message.response_metadata or {}).get(MODEL_USAGE_DELTA_KEY),
+        )
+        return event, is_final_response
+
+    def _is_final_response(self, message: AIMessage) -> bool:
+        """Check if the message represents a final response.
+
+        Args:
+            message: AIMessage to check.
+
+        Returns:
+            True if this is a final response, False otherwise.
+        """
+        return bool(message.response_metadata) and message.response_metadata.get("finish_reason") == "stop"
+
+    def _build_metadata_for_final_response(
+        self, is_final_response: bool, state: dict[str, Any] | None = None
+    ) -> dict[str, Any]:
+        """Build metadata for final response including previous_step_ids and pii_mapping.
+
+        Args:
+            is_final_response: Whether this is a final response.
+            state: Optional state dictionary containing pii_mapping and other data.
+
+        Returns:
+            Metadata dictionary with previous_step_ids and pii_mapping if applicable.
+        """
+        metadata: dict[str, Any] = {}
+
+        if not is_final_response:
+            return metadata
+
+        try:
+            previous_step_ids = self._get_previous_step_ids()
+            if previous_step_ids:
+                metadata["previous_step_ids"] = previous_step_ids
+        except Exception as e:
+            logger.warning("Attaching previous_step_ids to final response failed: %s", e, exc_info=True)
+
+        # Add PII mapping if present in state or nested metadata
+        if state:
+            pii_mapping = state.get("pii_mapping") or state.get("metadata", {}).get("pii_mapping")
+            if pii_mapping:
+                metadata[MetadataFieldKeys.PII_MAPPING] = pii_mapping
+            a2ui_content = state.get("a2ui_content")
+            if isinstance(a2ui_content, dict):
+                metadata["a2ui_content"] = a2ui_content
+
+        return metadata
+
+    def _get_previous_step_ids(self) -> list[str] | None:
+        """Get the list of previous step IDs based on thread context and operation mode.
+
+        Returns:
+            List of step IDs or None if no thread context or steps found.
+        """
+        thread_id = _THREAD_ID_CVAR.get()
+        if not thread_id:
+            return None
+
+        operation_mode = _OPERATION_MODE_CVAR.get() or "parallel"
+
+        coord_ids = self._coordinator_completed_tool_steps_by_thread.get(thread_id, [])
+        if coord_ids:
+            return self._filter_step_ids_by_mode(coord_ids, operation_mode)
+
+        completed_ids = self._completed_tool_steps_by_thread.get(thread_id, [])
+        if completed_ids:
+            return self._filter_step_ids_by_mode(completed_ids, operation_mode)
+
+        return None
+
+    def _filter_step_ids_by_mode(self, step_ids: list[str], operation_mode: str) -> list[str]:
+        """Filter step IDs based on operation mode.
+
+        Args:
+            step_ids: List of step IDs to filter.
+            operation_mode: Either "sequential" or "parallel".
+
+        Returns:
+            Filtered list of step IDs.
+        """
+        if operation_mode == "sequential":
+            return [step_ids[-1]] if step_ids else []
+        else:
+            return step_ids
+
+    def _process_artifacts(
+        self,
+        item: dict[str, Any],
+        pending_artifacts: list,
+        seen_artifact_hashes: set,
+    ) -> None:
+        """Process artifacts from a graph stream item.
+
+        Args:
+            item: The event item from the graph stream.
+            pending_artifacts: A list of artifacts waiting to be attached to a message.
+            seen_artifact_hashes: A set of hashes of artifacts already processed.
+        """
+        if "artifacts" not in item or not item["artifacts"]:
+            return
+
+        logger.info(f"Agent '{self.name}': Artifacts: {len(item['artifacts'])}")
+        for artifact in item["artifacts"]:
+            artifact_data = str(artifact.get("data", ""))
+            artifact_name = artifact.get("name", "")
+            artifact_hash = hashlib.sha256(f"{artifact_data}:{artifact_name}".encode()).hexdigest()
+
+            if artifact_hash not in seen_artifact_hashes:
+                pending_artifacts.append(artifact)
+                seen_artifact_hashes.add(artifact_hash)
+
+    def _process_a2a_stream_item(
+        self,
+        item: dict[str, Any],
+        pending_artifacts: list,
+        seen_artifact_hashes: set,
+        processed_message_count: int,
+    ) -> tuple[list[A2AEvent], bool, int]:
+        """Process a single item from the graph stream, handling artifacts and messages.
+
+        This method processes LangGraph stream items and converts them into A2AEvent objects
+        with proper metadata merging, artifact attachment, and reference handling.
+
+        Args:
+            item: The event item from the graph stream containing messages and metadata.
+            pending_artifacts: List of artifacts waiting to be attached to messages.
+            seen_artifact_hashes: Set of hashes of artifacts already processed.
+            processed_message_count: Number of messages already processed from the stream.
+
+        Returns:
+            A tuple containing:
+            - list[A2AEvent]: List of A2A events to yield to the executor.
+            - bool: True if a final event was encountered.
+            - int: Updated count of processed messages.
+        """
+        events_to_yield = []
+        is_final_event = False
+
+        self._process_artifacts(item, pending_artifacts, seen_artifact_hashes)
+        references = item.get("references", [])
+
+        if "messages" not in item or not item["messages"]:
+            return [], False, processed_message_count
+
+        new_messages = item["messages"][processed_message_count:]
+        updated_message_count = len(item["messages"])
+        for message in new_messages:
+            message_events, message_is_final = self._process_stream_message_item(
+                message,
+                item,
+                pending_artifacts,
+                references,
+            )
+            events_to_yield.extend(message_events)
+            if message_is_final:
+                is_final_event = True
+
+        return events_to_yield, is_final_event, updated_message_count
+
+    def _process_stream_message_item(
+        self,
+        message: BaseMessage,
+        item: dict[str, Any],
+        pending_artifacts: list,
+        references: list[Any],
+    ) -> tuple[list[A2AEvent], bool]:
+        """Process a single message from the graph stream.
+
+        Returns:
+            A tuple of (events to yield, is_final).
+        """
+        event_data, is_final = self._parse_a2a_stream_message(message, item)
+        events: list[A2AEvent] = []
+        reasoning_step_id = self._emit_provider_reasoning_if_present(message, events)
+
+        if event_data and self._should_yield_a2a_event(event_data):
+            self._enhance_event_with_context(event_data, item, pending_artifacts, references, is_final)
+            if reasoning_step_id:
+                event_data.setdefault("metadata", {})["previous_step_ids"] = [reasoning_step_id]
+            events.append(event_data)
+
+        return events, is_final
+
+    def _emit_provider_reasoning_if_present(
+        self,
+        message: BaseMessage,
+        events: list[A2AEvent],
+    ) -> str | None:
+        """Emit a provider reasoning event if the message carries thinkings.
+
+        Returns:
+            The reasoning step ID if an event was emitted, else None.
+        """
+        if not self._should_emit_provider_reasoning(message):
+            return None
+        reasoning_event = self._create_provider_reasoning_event(message)
+        if not reasoning_event:
+            return None
+        events.append(reasoning_event)
+        return (reasoning_event.get("metadata") or {}).get("step_id")
+
+    def _enhance_event_with_context(
+        self,
+        event_data: A2AEvent,
+        stream_item: dict[str, Any],
+        pending_artifacts: list,
+        references: list[Any],
+        is_final: bool,
+    ) -> None:
+        """Enhance A2AEvent with context from the stream item.
+
+        This method adds metadata, artifacts, and references to the A2AEvent
+        based on the current stream item context.
+
+        Args:
+            event_data: The A2AEvent to enhance.
+            stream_item: The stream item containing context information.
+            pending_artifacts: List of artifacts to attach to the event.
+            references: List of references to attach to final events.
+            is_final: Whether this is a final event.
+        """
+        self._merge_event_metadata(event_data, stream_item)
+        self._attach_pending_artifacts(event_data, pending_artifacts)
+
+        if is_final and references:
+            self._attach_references_to_final_event(event_data, references)
+
+        if is_final and stream_item.get(TOTAL_USAGE_KEY):
+            event_data[TOTAL_USAGE_KEY] = stream_item[TOTAL_USAGE_KEY]
+        if is_final and stream_item.get(MODEL_USAGE_KEY):
+            event_data[MODEL_USAGE_KEY] = stream_item[MODEL_USAGE_KEY]
+
+    def _merge_previous_step_ids(
+        self,
+        state_prev: list[str | int | None] | None,
+        event_prev: list[str | int | None] | None,
+    ) -> list[str | int] | None:
+        """Merge previous_step_ids from state and event metadata.
+
+        Args:
+            state_prev: Previous step IDs from state metadata.
+            event_prev: Previous step IDs from event metadata.
+
+        Returns:
+            Combined list of previous step IDs, or None if no lists to merge.
+        """
+        if (state_prev is None and event_prev is None) or (
+            not isinstance(state_prev, list)
+            and state_prev is not None
+            and not isinstance(event_prev, list)
+            and event_prev is not None
+        ):
+            return None
+
+        state_list = state_prev if isinstance(state_prev, list) else []
+        event_list = event_prev if isinstance(event_prev, list) else []
+
+        combined: list[Any] = []
+
+        for step_id in event_list:
+            if step_id is not None and step_id not in combined:
+                combined.append(step_id)
+
+        for step_id in state_list:
+            if step_id is not None and step_id not in combined:
+                combined.append(step_id)
+
+        return combined
+
+    def _merge_event_metadata(self, event_data: A2AEvent, stream_item: dict[str, Any]) -> None:
+        """Merge metadata from stream item into the A2AEvent.
+
+        Args:
+            event_data: The A2AEvent to update with merged metadata.
+            stream_item: The stream item containing state metadata.
+        """
+        state_metadata = stream_item.get("metadata", {})
+        existing_metadata = event_data.get("metadata", {})
+        if isinstance(existing_metadata, dict) and isinstance(state_metadata, dict):
+            merged_metadata = {**state_metadata, **existing_metadata}
+
+            state_prev = state_metadata.get("previous_step_ids") or []
+            event_prev = existing_metadata.get("previous_step_ids") or []
+            combined_ids = self._merge_previous_step_ids(state_prev, event_prev)
+            if combined_ids is not None:
+                merged_metadata["previous_step_ids"] = combined_ids
+        else:
+            merged_metadata = state_metadata or existing_metadata
+
+        event_data["metadata"] = merged_metadata
+
+    def _attach_pending_artifacts(self, event_data: A2AEvent, pending_artifacts: list) -> None:
+        """Attach pending artifacts to the A2AEvent and clear the pending list.
+
+        Args:
+            event_data: The A2AEvent to attach artifacts to.
+            pending_artifacts: List of artifacts to attach and clear.
+        """
+        if pending_artifacts:
+            event_data["artifacts"] = pending_artifacts.copy()
+            pending_artifacts.clear()
+
+    def _attach_references_to_final_event(self, event_data: A2AEvent, references: list[Any]) -> None:
+        """Attach references to final events.
+
+        Args:
+            event_data: The final A2AEvent to attach references to.
+            references: List of references to attach.
+        """
+        if references:
+            event_data["references"] = references
+
+    def _setup_thread_context(self, config: dict[str, Any]) -> tuple[str | None, Any]:
+        """Set up thread context for step linkage during streaming.
+
+        Args:
+            config: Graph configuration
+
+        Returns:
+            Tuple of (thread_id, context_token)
+        """
+        current_thread_id: str | None = None
+        try:
+            configurable = config.get("configurable", {})
+            thread_key = self.thread_id_key or "thread_id"
+            current_thread_id = str(configurable.get(thread_key)) if configurable.get(thread_key) else None
+        except Exception:
+            current_thread_id = None
+
+        token = None
+
+        try:
+            start_step_counter_scope()
+        except Exception as exc:
+            logger.debug("Starting step counter scope failed: %s", exc)
+
+        if current_thread_id:
+            token = _THREAD_ID_CVAR.set(current_thread_id)
+            self._tool_parent_map_by_thread[current_thread_id] = {}
+            self._completed_tool_steps_by_thread[current_thread_id] = []
+            self._emitted_tool_calls_by_thread[current_thread_id] = set()
+
+        return current_thread_id, token
+
+    def _cleanup_thread_context(self, current_thread_id: str | None, token: Any) -> None:
+        """Clean up thread context and reset context variables.
+
+        Args:
+            current_thread_id: The thread ID to clean up
+            token: The context token to reset
+        """
+        try:
+            end_step_counter_scope()
+        except Exception as exc:
+            logger.debug("Ending step counter scope failed: %s", exc)
+
+        if current_thread_id:
+            self._tool_parent_map_by_thread.pop(current_thread_id, None)
+            self._completed_tool_steps_by_thread.pop(current_thread_id, None)
+            self._last_status_step_id_by_thread.pop(current_thread_id, None)
+            self._emitted_tool_calls_by_thread.pop(current_thread_id, None)
+
+        if token is not None:
+            try:
+                _THREAD_ID_CVAR.reset(token)
+            except ValueError as e:
+                logger.debug("Context variable token from different context, skipping reset: %s", e)
+            except Exception as e:
+                logger.error("Resetting _THREAD_ID_CVAR failed: %s", e, exc_info=True)
+        try:
+            _STEP_LIMIT_CONFIG_CVAR.set(None)
+        except Exception:
+            logger.debug("Failed to reset step limit config context; continuing cleanup.")
+
+    def _handle_stream_item(
+        self, item: tuple, pending_artifacts: list, seen_artifact_hashes: set, processed_message_count: int
+    ) -> tuple[list[A2AEvent], bool, int]:
+        """Handle a single stream item.
+
+        Args:
+            item: Stream item tuple (mode, data)
+            pending_artifacts: List of pending artifacts
+            seen_artifact_hashes: Set of seen artifact hashes
+            processed_message_count: Current message count
+
+        Returns:
+            Tuple of (events_to_yield, is_final, updated_message_count)
+        """
+        mode, data = item
+
+        if mode == StreamMode.CUSTOM:
+            delegation_event: A2AEvent = data
+            delegation_event = self._normalize_custom_stream_event(delegation_event)
+            if self._should_yield_a2a_event(delegation_event):
+                return [delegation_event], False, processed_message_count
+            return [], False, processed_message_count
+        elif mode == StreamMode.VALUES:
+            stream_data = data
+        else:
+            return [], False, processed_message_count
+
+        events, is_final, updated_message_count = self._process_a2a_stream_item(
+            stream_data, pending_artifacts, seen_artifact_hashes, processed_message_count
+        )
+        return events, is_final, updated_message_count
+
+    async def _arun_a2a_stream(self, query: str, **kwargs: Any) -> AsyncGenerator[A2AEvent, None]:
+        """Internal implementation of arun_a2a_stream without MCP handling.
+
+        Args:
+            query: The input query for the agent.
+            **kwargs: Additional keyword arguments.
+
+        Yields:
+            Dictionaries with "status" and "content" keys for status events.
+            Status events may include "artifacts" field when tools generate artifacts.
+            Possible statuses: "working", "completed", "failed", "canceled".
+        """
+        context = self._initialize_streaming_context(query, **kwargs)
+        await self._run_stream_before_run_hook(context)
+
+        try:
+            async for event in self._handle_streaming_process(context):
+                yield event
+
+            await self._run_stream_after_run_success_hook(context)
+
+            async for event in self._ensure_final_completion(context):
+                yield event
+
+        except Exception as e:
+            await self._run_stream_after_run_error_hook(context, e)
+            async for event in self._handle_streaming_error(context, e):
+                yield event
+        finally:
+            self._cleanup_thread_context(context.current_thread_id, context.token)
+
+    async def _run_stream_before_run_hook(self, context: "_StreamingContext") -> None:
+        """Apply run-level middleware before streaming graph execution."""
+        try:
+            before_updates = await self._middleware_manager.abefore_run(context.graph_input, context.config)
+            if before_updates:
+                self._apply_before_run_updates(context.graph_input, before_updates)
+        except SkillInstallError:
+            raise
+        except Exception:  # pragma: no cover
+            pass  # pragma: no cover
+
+    def _resolve_stream_final_text(self, context: "_StreamingContext") -> str:
+        """Resolve final text used for run-level completion middleware hook."""
+        if isinstance(context.last_final_content, str):
+            return context.last_final_content
+        if isinstance(context.final_state, dict):
+            return self._extract_output_from_final_state(context.final_state)
+        return ""
+
+    async def _run_stream_after_run_success_hook(self, context: "_StreamingContext") -> None:
+        """Run completion middleware hook for successful streaming runs."""
+        try:
+            final_text = self._resolve_stream_final_text(context)
+            await self._middleware_manager.aafter_run(
+                final_state=context.final_state or {},
+                output=final_text or None,
+                config=context.config,
+                error=None,
+            )
+        except Exception:  # pragma: no cover
+            pass  # pragma: no cover
+
+    async def _run_stream_after_run_error_hook(self, context: "_StreamingContext", error: Exception) -> None:
+        """Run completion middleware hook for failed streaming runs."""
+        try:
+            await self._middleware_manager.aafter_run(
+                final_state=context.final_state or {},
+                output=None,
+                config=context.config,
+                error=error,
+            )
+        except Exception:  # pragma: no cover
+            pass  # pragma: no cover
+
+    def _initialize_streaming_context(self, query: str, **kwargs: Any) -> "_StreamingContext":
+        """Initialize the streaming context with all necessary setup.
+
+        Args:
+            query: The user's input query to process.
+            **kwargs: Additional keyword arguments including optional metadata and configuration.
+
+        Returns:
+            Configured _StreamingContext object ready for streaming execution.
+        """
+        files = kwargs.pop("files", [])
+        if files is None:
+            files = []
+        augmented_query, input_attachments = self._prepare_file_runtime_context(query, files, kwargs)
+
+        memory_user_id: str | None = kwargs.get("memory_user_id")
+
+        # Create config first to ensure thread_id is generated
+        config = self._create_graph_config(**kwargs)
+        thread_id = self._get_thread_id_from_config(config)
+
+        graph_input = self._prepare_graph_input(augmented_query, thread_id=thread_id, **kwargs)
+        self._attach_input_attachments_to_graph_input(graph_input, input_attachments)
+
+        current_thread_id, token = self._setup_thread_context(config)
+
+        if self.enable_a2a_token_streaming and self.model:
+            self.model.disable_streaming = False
+
+        a2ui_stream_filter = (
+            self._create_a2ui_stream_filter() if self._is_a2ui_enabled(getattr(self, "config", None)) else None
+        )
+
+        return _StreamingContext(
+            original_query=query,
+            graph_input=graph_input,
+            config=config,
+            memory_user_id=memory_user_id,
+            current_thread_id=current_thread_id,
+            token=token,
+            enable_token_streaming=self.enable_a2a_token_streaming,
+            a2ui_stream_filter=a2ui_stream_filter,
+        )
+
+    @staticmethod
+    def _is_a2ui_enabled(config: Any | None) -> bool:
+        """Return True when agent config enables A2UI."""
+        if config is None:
+            return False
+        if isinstance(config, dict):
+            raw_a2ui_config = config.get("a2ui")
+        else:
+            raw_a2ui_config = getattr(config, "a2ui", None)
+        return isinstance(raw_a2ui_config, dict) and bool(raw_a2ui_config.get("enabled"))
+
+    def _create_a2ui_stream_filter(self) -> Any:
+        """Create an incremental A2UI markup filter for token streaming."""
+        from aip_agents.a2ui.stream_filter import A2UIStreamTextFilter
+
+        return A2UIStreamTextFilter()
+
+    @staticmethod
+    def _filter_a2ui_token_stream_content(context: "_StreamingContext | None", content: str) -> str | None:
+        """Return sanitized visible token text, or None when nothing should be emitted."""
+        stream_filter = None if context is None else context.a2ui_stream_filter
+        if stream_filter is None:
+            return content if content else None
+        filtered = stream_filter.feed(content)
+        return filtered if filtered else None
+
+    async def _handle_streaming_process(self, context: "_StreamingContext") -> AsyncGenerator[A2AEvent, None]:
+        """Handle the main streaming process including initial status and event processing.
+
+        Args:
+            context: The streaming context containing query, config, and thread information.
+
+        Yields:
+            Streaming events including initial status and processed streaming items.
+        """
+        if self._should_emit_initial_status_event():
+            initial_status_event = self._create_initial_status_event()
+            self._log_streaming_event_debug("initial_status", initial_status_event)
+            yield initial_status_event
+
+        async for event in self._process_streaming_items(context):
+            self._log_streaming_event_debug("process_stream_item", event)
+            yield event
+
+    def _should_emit_initial_status_event(self) -> bool:
+        """Return whether this stream should emit the default initial status."""
+        return not is_delegated_agent_execution()
+
+    def _create_initial_status_event(self) -> A2AEvent:
+        """Create and setup the initial status event."""
+        initial_status_event = self._create_a2a_event(
+            event_type=A2AStreamEventType.STATUS_UPDATE, content=DefaultStepMessages.EN.value
+        )
+
+        try:
+            thread_id = _THREAD_ID_CVAR.get()
+            if thread_id:
+                step_id = initial_status_event.get("metadata", {}).get("step_id")
+                if step_id:
+                    self._last_status_step_id_by_thread[thread_id] = str(step_id)
+        except Exception:
+            pass
+
+        return initial_status_event
+
+    async def _process_streaming_items(self, context: "_StreamingContext") -> AsyncGenerator[A2AEvent, None]:
+        """Process individual streaming items from the LangGraph execution.
+
+        Handles the core streaming logic by iterating through items produced by
+        the compiled LangGraph, processing both VALUES and CUSTOM stream modes,
+        and managing final event generation.
+
+        Args:
+            context: The streaming context containing graph input, configuration,
+                and state tracking information.
+
+        Yields:
+            dict[str, Any]: A2A events generated from the stream processing,
+                including status updates, final responses, and completion events.
+        """
+        if context.enable_token_streaming:
+            if self.event_emitter is None:
+                self.event_emitter = self._create_default_event_emitter()
+            elif not self._get_stream_handler():
+                logger.warning(
+                    "Agent '%s': No StreamEventHandler found in event_emitter. "
+                    "Reinitializing event_emitter using default emitter.",
+                    self.name,
+                )
+                self.event_emitter = self._create_default_event_emitter()
+
+            async for event in self._process_a2a_streaming_with_tokens(context):
+                yield event
+        else:
+            enhanced_input = context.graph_input
+            async for event in self._create_graph_stream_events(enhanced_input, context):
+                yield event
+
+    async def _process_a2a_streaming_with_tokens(self, context: "_StreamingContext") -> AsyncGenerator[A2AEvent, None]:
+        """Process A2A streaming with token streaming support using aiostream.
+
+        Supports both LM Invoker and LangChain models by detecting the appropriate
+        token source and merging with graph events.
+
+        Uses aiostream to merge token streaming and graph execution streams,
+        yielding events in real-time order as they arrive.
+
+        Args:
+            context: The streaming context containing graph input, configuration,
+                and state tracking information.
+
+        Yields:
+            dict[str, Any]: A2A events generated from the stream processing,
+                including status updates, final responses, and completion events.
+
+        Raises:
+            RuntimeError: If token streaming is requested but event_emitter is not available.
+        """
+        if not self.event_emitter:
+            raise RuntimeError(f"Agent '{self.name}': Event emitter required for token streaming")
+        if astream is None:
+            raise RuntimeError(
+                "aiostream is required for token streaming support. "
+                "Install the 'aiostream' dependency or disable token streaming."
+            )
+
+        try:
+            if self._has_lm_invoker():
+                token_stream, enhanced_input = self._create_token_stream(context)
+                graph_stream = self._graph_stream_events_with_token_emitter_cleanup(
+                    enhanced_input,
+                    context,
+                )
+
+                if token_stream is None:
+                    raise RuntimeError(f"Agent '{self.name}': Token stream not available for LM invoker.")
+
+                merged = astream.merge(token_stream, graph_stream)
+                async with merged.stream() as merged_stream:
+                    async for event in merged_stream:
+                        yield event
+            else:
+                _, enhanced_input = self._create_token_stream(context)
+                async for event in self._create_graph_stream_events(enhanced_input, context):
+                    yield event
+
+        except Exception as e:
+            if self.event_emitter is not None:
+                await self.event_emitter.close()
+                self.event_emitter = None
+            logger.error(f"Agent '{self.name}': Error during A2A token streaming: {e}")
+            raise
+
+    async def _graph_stream_events_with_token_emitter_cleanup(
+        self,
+        enhanced_input: dict[str, Any],
+        context: "_StreamingContext",
+    ) -> AsyncGenerator[A2AEvent, None]:
+        """Run graph streaming and close the token emitter when graph execution completes."""
+        try:  # pragma: no cover
+            async for event in self._create_graph_stream_events(enhanced_input, context):  # pragma: no cover
+                yield event  # pragma: no cover
+        finally:
+            if self.event_emitter is not None:  # pragma: no cover
+                await self.event_emitter.close()  # pragma: no cover
+                self.event_emitter = None  # pragma: no cover
+
+    async def _create_lm_invoker_token_stream(self, context: "_StreamingContext") -> AsyncGenerator[A2AEvent, None]:
+        """Generate A2A events from LM Invoker token stream.
+
+        Uses StreamEventHandler to capture tokens emitted by LM Invoker.
+
+        Yields:
+            A2A events generated from LM Invoker token stream.
+
+        Raises:
+            RuntimeError: If no StreamEventHandler is found in event_emitter.
+        """
+        stream_handler = self._get_stream_handler()
+        if stream_handler is None:
+            raise RuntimeError(f"Agent '{self.name}': StreamEventHandler is required for token streaming.")
+
+        try:
+            async for event in stream_handler.stream():
+                if event is None:
+                    break
+
+                token_event = self._convert_raw_token_to_a2a_event(event, context)
+                if token_event:
+                    yield token_event
+        except Exception as e:
+            logger.error(f"Agent '{self.name}': LM Invoker token stream error: {e}")
+
+    def _create_token_stream(
+        self,
+        context: "_StreamingContext",
+    ) -> tuple[AsyncGenerator[A2AEvent, None] | None, dict[str, Any]]:
+        """Create appropriate token stream and enhanced input for the active model backend.
+
+        Args:
+            context: Streaming context containing graph input and configuration.
+
+        Returns:
+            Tuple of (token_stream, enhanced_input) where token_stream yields A2A token
+            events and enhanced_input is the graph input dictionary (augmented with event
+            emitter when required by LM Invoker backends).
+        """
+        if self._has_lm_invoker():
+            token_stream = self._create_lm_invoker_token_stream(context)
+            enhanced_input = {**context.graph_input, "event_emitter": self.event_emitter}
+        else:
+            token_stream = None
+            enhanced_input = context.graph_input
+
+        return token_stream, enhanced_input
+
+    async def _create_graph_stream_events(
+        self, enhanced_input: dict[str, Any], context: "_StreamingContext"
+    ) -> AsyncGenerator[A2AEvent, None]:
+        """Generate A2A events from graph execution.
+
+        Args:
+            enhanced_input: The enhanced input for the graph execution.
+            context: The streaming context containing state tracking information.
+
+        Yields:
+            A2A events generated from graph execution.
+        """
+        try:
+            stream_modes = self._get_stream_modes(context)
+            graph_execution = self._graph_runtime.stream(
+                self._compiled_graph,
+                enhanced_input,
+                context.config,
+                stream_modes,
+            )
+
+            async for item in graph_execution:
+                stream_mode, stream_data = item
+
+                if stream_mode == StreamMode.MESSAGES.value:
+                    message_data = cast(tuple[Any, dict[str, Any]], stream_data)
+                    async for token_event in self._process_message_stream_item(message_data, context):
+                        yield token_event
+                    continue
+
+                async for event in self._process_graph_stream_item(item, stream_mode, stream_data, context):
+                    yield event
+        except Exception as e:
+            logger.exception("Agent '%s': Graph processing error: %s", self.name, e)
+            raise
+
+    def _get_stream_modes(self, context: "_StreamingContext") -> list[str]:
+        """Determine stream modes based on token streaming configuration.
+
+        Args:
+            context: Streaming context containing token streaming configuration.
+
+        Returns:
+            List of stream modes to use for graph execution.
+        """
+        stream_modes = [StreamMode.VALUES.value, StreamMode.CUSTOM.value]
+
+        if context.enable_token_streaming and not self._has_lm_invoker():
+            stream_modes.append(StreamMode.MESSAGES.value)
+
+        return stream_modes
+
+    async def _process_graph_stream_item(
+        self,
+        item: tuple[str, Any],
+        stream_mode: str,
+        stream_data: Any,
+        context: "_StreamingContext",
+    ) -> AsyncGenerator[A2AEvent, None]:
+        """Process a single graph stream item and yield A2A events.
+
+        Args:
+            item: The stream item tuple (mode, data).
+            stream_mode: The stream mode of this item.
+            stream_data: The data from the stream item.
+            context: Streaming context for state tracking.
+
+        Yields:
+            A2A events generated from the stream item.
+        """
+        context.final_state = copy.copy(stream_data) if stream_mode == StreamMode.VALUES.value else context.final_state
+
+        pending_artifacts = context.pending_artifacts if context.pending_artifacts is not None else []
+        seen_artifact_hashes = context.seen_artifact_hashes if context.seen_artifact_hashes is not None else set()
+        events, is_final, context.processed_message_count = self._handle_stream_item(
+            item, pending_artifacts, seen_artifact_hashes, context.processed_message_count
+        )
+
+        if is_final:
+            context.final_event_yielded = True
+
+        for event in events:
+            await self._capture_final_content_if_needed(context, event)
+            if (
+                isinstance(event, dict)
+                and event.get("event_type") in (A2AStreamEventType.FINAL_RESPONSE, "final_response")
+                and isinstance(context.final_state, dict)
+            ):
+                event_metadata = event.get("metadata")
+                if isinstance(event_metadata, dict):
+                    a2ui_content = context.final_state.get("a2ui_content")
+                    if isinstance(a2ui_content, dict):
+                        event_metadata["a2ui_content"] = a2ui_content
+            processed_event = self._update_final_response_for_streaming(context, event)
+            yield processed_event
+
+    async def _process_message_stream_item(
+        self,
+        message_data: tuple[Any, dict[str, Any]],
+        context: "_StreamingContext | None" = None,
+    ) -> AsyncGenerator[A2AEvent, None]:
+        """Process message stream items to extract token events.
+
+        The "messages" stream mode yields tuples of (AIMessageChunk, metadata).
+        This method extracts token content from AIMessageChunk and converts it
+        to A2A CONTENT_CHUNK events with TOKEN kind.
+
+        Args:
+            message_data: Tuple of (message_chunk, metadata) from messages stream
+            context: Streaming context containing optional A2UI token filter state
+
+        Yields:
+            A2A CONTENT_CHUNK events with TOKEN kind
+        """
+        try:
+            message_chunk, _ = message_data
+
+            # Filter out events with response_metadata.finish_reason attribute
+            # since it is a response from subagent
+            if hasattr(message_chunk, "response_metadata") and message_chunk.response_metadata:
+                if "finish_reason" in message_chunk.response_metadata:
+                    return
+
+            is_tool_call_event = self._is_message_stream_tool_call_chunk(message_chunk)
+            is_has_content_event = hasattr(message_chunk, "content") and message_chunk.content
+
+            if (
+                is_has_content_event
+                and not is_tool_call_event
+                and not self._is_message_stream_reasoning_chunk(message_chunk)
+            ):
+                token_content = self._extract_visible_message_stream_token_content(message_chunk.content)
+                if token_content is None:
+                    return
+                token_content = self._filter_a2ui_token_stream_content(context, token_content)
+                if token_content is None:
+                    return  # pragma: no cover
+                token_event = self._create_a2a_event(
+                    event_type=A2AStreamEventType.CONTENT_CHUNK,
+                    content=token_content,
+                    metadata={MetadataFieldKeys.KIND: Kind.TOKEN},
+                )
+                yield token_event
+
+        except Exception as e:
+            logger.error(f"Agent '{self.name}': Error processing message stream item: {e}")
+
+    @staticmethod
+    def _is_message_stream_tool_call_chunk(message_chunk: Any) -> bool:
+        """Return True when a LangGraph message-stream chunk is part of tool-call assembly."""
+        return bool(
+            (hasattr(message_chunk, "tool_calls") and message_chunk.tool_calls)
+            or (hasattr(message_chunk, "tool_call_chunks") and message_chunk.tool_call_chunks)
+        )
+
+    @staticmethod
+    def _is_message_stream_reasoning_chunk(message_chunk: Any) -> bool:
+        """Return True for provider reasoning chunks that should use activity events, not answer tokens."""
+        additional_kwargs = getattr(message_chunk, "additional_kwargs", None)
+        if not isinstance(additional_kwargs, dict):
+            return False
+
+        raw_thinkings = additional_kwargs.get(THINKINGS_METADATA_KEY)
+        if not isinstance(raw_thinkings, list) or not raw_thinkings:
+            return False
+
+        content = getattr(message_chunk, "content", None)
+        if not isinstance(content, str):
+            return False  # pragma: no cover
+
+        content_text = content.strip()
+        if not content_text:
+            return False  # pragma: no cover
+
+        thinking_texts = [
+            entry.get("thinking", "").strip()
+            for entry in raw_thinkings
+            if isinstance(entry, dict) and isinstance(entry.get("thinking"), str)
+        ]
+        combined_thinking = "\n\n".join(text for text in thinking_texts if text)
+        return content_text in {*thinking_texts, combined_thinking}
+
+    @staticmethod
+    def _extract_visible_message_stream_block_text(block: Any) -> str:
+        """Extract visible text from a single LangGraph content block."""
+        if not isinstance(block, str | dict):
+            return ""  # pragma: no cover
+
+        if isinstance(block, str):
+            return block  # pragma: no cover
+
+        block_type = str(block.get("type") or block.get("kind") or "").lower()
+        if block_type in {"reasoning", "reasoning_content", "thinking", "redacted_thinking"}:
+            return ""
+
+        if block_type not in {"text", "output_text"}:
+            return ""  # pragma: no cover
+
+        text = block.get("text")
+        if isinstance(text, str):
+            return text
+
+        text = block.get("content")  # pragma: no cover
+        return text if isinstance(text, str) else ""  # pragma: no cover
+
+    @staticmethod
+    def _extract_visible_message_stream_token_content(content: Any) -> str | None:
+        """Extract user-visible answer text from a LangGraph message-stream chunk."""
+        if isinstance(content, str):
+            return content if content else None
+
+        if not isinstance(content, list):
+            return None
+
+        visible_parts: list[str] = []
+        for block in content:
+            visible_text = BaseLangGraphAgent._extract_visible_message_stream_block_text(block)
+            if visible_text:
+                visible_parts.append(visible_text)
+
+        return "".join(visible_parts) if visible_parts else None
+
+    def _update_final_response_for_streaming(self, context: "_StreamingContext", event: A2AEvent) -> A2AEvent:
+        """Update final response events with appropriate streaming configuration.
+
+        For FINAL_RESPONSE events, this method updates the metadata while preserving
+        the canonical final response content for consumers that need transformed output.
+
+        Args:
+            context: The streaming context containing streaming configuration
+            event: The event dictionary to process
+
+        Returns:
+            The processed event dictionary with updated metadata and content
+        """
+        if event.get("event_type") == A2AStreamEventType.FINAL_RESPONSE:
+            metadata = event.get("metadata")
+            if not isinstance(metadata, dict):
+                metadata = {}
+                event["metadata"] = metadata
+            metadata[MetadataFieldKeys.TOKEN_STREAMING] = bool(context.enable_token_streaming)
+            if context.enable_token_streaming:
+                metadata[MetadataFieldKeys.CANONICAL] = True
+        return event
+
+    def _convert_raw_token_to_a2a_event(
+        self, raw_event: str, context: "_StreamingContext | None" = None
+    ) -> A2AEvent | None:
+        """Parse raw token event into A2A event.
+
+        Args:
+            raw_event: The raw event containing the raw event.
+            context: Streaming context containing optional A2UI token filter state.
+
+        Returns:
+            dict[str, Any]: A2A event generated from the stream processing,
+                including status updates, final responses, and completion events.
+        """
+        try:
+            event_data = json.loads(raw_event)
+            content = event_data.get("value", "")
+            event_type = event_data.get("type")
+            metadata = event_data.get("metadata")
+            if event_type == A2AStreamEventType.STATUS_UPDATE.value and isinstance(metadata, dict):
+                return self._create_a2a_event(
+                    event_type=A2AStreamEventType.STATUS_UPDATE,
+                    content=content if isinstance(content, str) else "",
+                    metadata=metadata,
+                )
+            if isinstance(content, str) and content:
+                filtered_content = self._filter_a2ui_token_stream_content(context, content)
+                if filtered_content is None:
+                    return None  # pragma: no cover
+                return self._create_a2a_event(
+                    event_type=A2AStreamEventType.CONTENT_CHUNK,
+                    content=filtered_content,
+                    metadata={MetadataFieldKeys.KIND: Kind.TOKEN},
+                )
+        except Exception as e:
+            logger.debug(f"Agent '{self.name}': Error parsing token event: {e}")
+        return None
+
+    async def _capture_final_content_if_needed(self, context: "_StreamingContext", event: A2AEvent) -> None:
+        """Capture final content from A2A events for streaming hooks.
+
+        Monitors A2A events for final response content and triggers early memory
+        persistence to ensure conversation content is saved even if consumers
+        stop reading the stream after receiving the final response.
+
+        Args:
+            context: The streaming context containing memory state and user
+                identification information.
+            event: The A2A event dictionary that may contain final response content.
+        """
+        try:
+            if not (isinstance(event, dict) and event.get("event_type") == A2AStreamEventType.FINAL_RESPONSE):
+                return
+
+            content = event.get("content")
+            if not isinstance(content, str) or not content:
+                return  # pragma: no cover
+
+            context.last_final_content = content
+
+            # Invoke middleware streaming hook (best-effort) for early persistence.
+            hook_context: dict[str, Any] = {
+                "agent_name": self.name,
+                "thread_id": context.current_thread_id,
+                "original_query": context.original_query,
+                "memory_user_id": context.memory_user_id,
+                "config": context.config,
+                "event": event,
+                "final_state": context.final_state,
+                "memory_persisted": context.saved_memory,
+                "a2ui_repair_handler": self._repair_a2ui_final_response,
+            }
+            await self._middleware_manager.aon_final_response(content, hook_context)
+
+            # Propagate persisted flag for dedupe.
+            if hook_context.get("memory_persisted"):
+                context.saved_memory = True
+                if isinstance(context.final_state, dict):
+                    context.final_state["memory_persisted"] = True
+
+        except Exception:
+            # Best-effort hook.
+            return
+
+    async def _repair_a2ui_final_response(
+        self,
+        *,
+        content: str,
+        context: dict[str, Any],
+        validation: dict[str, Any] | None,
+        schema_version: str,
+    ) -> str | None:
+        """Ask the configured model for one corrected A2UI final-response payload."""
+        prompt = self._build_a2ui_repair_prompt(
+            content=content,
+            original_query=context.get("original_query"),
+            validation=validation,
+            schema_version=schema_version,
+        )
+        messages = [
+            SystemMessage(
+                content=(
+                    "You repair malformed A2UI final responses. Return only corrected "
+                    "A2UI content wrapped in <a2ui-json>...</a2ui-json> tags."
+                )
+            ),
+            HumanMessage(content=prompt),
+        ]
+
+        try:
+            if self.lm_invoker is not None:
+                previous_tools = None
+                restore_tools = hasattr(self.lm_invoker, "set_tools")
+                if restore_tools and hasattr(self.lm_invoker, "tools"):
+                    previous_tools = list(self.lm_invoker.tools)
+                if restore_tools:
+                    self.lm_invoker.set_tools([])
+                try:
+                    lm_messages = convert_langchain_messages_to_gllm_messages(messages, "")
+                    lm_output = await self.lm_invoker.invoke(messages=lm_messages)
+                    return self._normalize_repair_content(convert_lm_output_to_langchain_message(lm_output).content)
+                finally:
+                    if restore_tools and previous_tools is not None:
+                        self.lm_invoker.set_tools(previous_tools)
+
+            model = getattr(self, "model", None)
+            if model is None or not hasattr(model, "ainvoke"):
+                return None  # pragma: no cover
+            ai_message = await model.ainvoke(messages)
+            return self._normalize_repair_content(getattr(ai_message, "content", ai_message))
+        except Exception as exc:
+            logger.warning("Agent '%s': A2UI repair model call failed: %s", self.name, exc, exc_info=True)
+            return None
+
+    def _build_a2ui_repair_prompt(
+        self,
+        *,
+        content: str,
+        original_query: Any,
+        validation: dict[str, Any] | None,
+        schema_version: str,
+    ) -> str:
+        """Build a bounded A2UI repair prompt from validation errors and original output."""
+        errors = []
+        if isinstance(validation, dict) and isinstance(validation.get("errors"), list):
+            errors = validation["errors"]
+
+        return (
+            "The previous final response contained invalid A2UI output.\n"
+            f"A2UI schema version: {schema_version}\n"
+            f"Original user request: {original_query or ''}\n"
+            f"Validation errors: {json.dumps(errors, ensure_ascii=False)}\n\n"
+            "Repair instructions:\n"
+            "- Return corrected A2UI only.\n"
+            "- Wrap the corrected JSON in <a2ui-json> and </a2ui-json>.\n"
+            "- Do not include explanations, markdown, diagnostics, or plain text.\n"
+            "- Preserve the user's intended UI as much as possible.\n\n"
+            "Original final response:\n"
+            f"{content}"
+        )
+
+    @staticmethod
+    def _normalize_repair_content(content: Any) -> str:
+        """Normalize repair model output to a string, handling structured content blocks."""
+        if isinstance(content, str):
+            return content
+        if isinstance(content, list):
+            parts = []
+            for block in content:
+                if isinstance(block, dict):
+                    parts.append(block.get("text", ""))
+                else:
+                    parts.append(str(block))
+            return "".join(parts)
+        return str(content)
+
+    async def _ensure_final_completion(self, context: "_StreamingContext") -> AsyncGenerator[A2AEvent, None]:
+        """Ensure final completion events are yielded if not already done.
+
+        Args:
+            context: The streaming context containing pending artifacts and
+                other state information.
+
+        Yields:
+            dict[str, Any]: The final completion event.
+        """
+        if not context.final_event_yielded:
+            pending_artifacts = context.pending_artifacts if context.pending_artifacts is not None else []
+            final_state = context.final_state or {}
+            completion_event = self._create_completion_event(pending_artifacts, final_state)
+            self._log_streaming_event_debug("final_completion", completion_event)
+            yield completion_event
+
+    async def _handle_streaming_error(
+        self,
+        context: "_StreamingContext",
+        error: Exception,
+    ) -> AsyncGenerator[A2AEvent, None]:
+        """Handle streaming errors gracefully.
+
+        Provides error handling for the A2A streaming process, ensuring errors
+        are properly logged and communicated to the client while preserving
+        any pending artifacts generated before the error occurred.
+
+        Args:
+            context: The streaming context containing pending artifacts and
+                other state information.
+            error: The exception that occurred during streaming.
+
+        Yields:
+            dict[str, Any]: An error event containing the failure status and
+                error message, optionally including any pending artifacts.
+        """
+        logger.error(f"Error in agent stream: {error}", exc_info=True)
+        error_event = self._create_a2a_event(
+            event_type=A2AStreamEventType.ERROR,
+            content=f"Error: {str(error)}",
+            metadata={"status": "failed"},
+            artifacts=context.pending_artifacts,
+            is_final=True,
+        )
+        error_event["status"] = "failed"
+        self._log_streaming_event_debug("error_event", error_event)
+        yield error_event
+
+    def _extract_references_from_state(self, final_state: dict[str, Any] | None) -> list[Chunk] | None:
+        """Extract and validate references from final state.
+
+        Args:
+            final_state: The final state of the agent.
+
+        Returns:
+            Validated references or None if not available.
+        """
+        if final_state and isinstance(final_state, dict) and final_state.get("references"):
+            try:
+                return validate_references(final_state["references"])
+            except Exception:
+                pass
+        return None
+
+    def _extract_total_usage_from_state(self, final_state: dict[str, Any] | None) -> dict[str, Any] | None:
+        """Extract total usage from final state.
+
+        Args:
+            final_state: The final state of the agent.
+
+        Returns:
+            Total usage metadata or None if not available.
+        """
+        if final_state and isinstance(final_state, dict) and final_state.get(TOTAL_USAGE_KEY):
+            return final_state[TOTAL_USAGE_KEY]
+        return None
+
+    def _extract_model_usage_from_state(self, final_state: dict[str, Any] | None) -> list[dict[str, Any]] | None:
+        """Extract model usage from final state when available."""
+        if final_state and isinstance(final_state, dict) and final_state.get(MODEL_USAGE_KEY):
+            value = final_state[MODEL_USAGE_KEY]
+            if isinstance(value, list):
+                return value
+        return None
+
+    def _build_completion_metadata(self, final_state: dict[str, Any] | None) -> dict[str, Any]:
+        """Build metadata for completion event.
+
+        Args:
+            final_state: The final state of the agent.
+
+        Returns:
+            Metadata dictionary with previous_step_ids and pii_mapping if available.
+        """
+        metadata: dict[str, Any] = {}
+
+        if final_state and isinstance(final_state, dict):
+            a2ui_content = final_state.get("a2ui_content")
+            if isinstance(a2ui_content, dict):
+                metadata["a2ui_content"] = a2ui_content
+
+        # Add previous step IDs if available
+        try:
+            thread_id = _THREAD_ID_CVAR.get()
+            if thread_id and thread_id in self._completed_tool_steps_by_thread:
+                completed_ids = self._completed_tool_steps_by_thread[thread_id]
+                if completed_ids:
+                    metadata["previous_step_ids"] = list(completed_ids)
+        except Exception as e:
+            logger.warning("Attaching previous_step_ids to completion event failed: %s", e, exc_info=True)
+
+        return metadata
+
+    def _create_completion_event(self, pending_artifacts: list, final_state: dict[str, Any]) -> A2AEvent:
+        """Helper to create the completion event with artifacts and references if available.
+
+        This method is used to create the completion event with artifacts and references if available.
+
+        Args:
+            pending_artifacts: List of artifacts waiting to be attached to a message.
+            final_state: The final state of the agent.
+
+        Returns:
+            A dictionary with "status" and "content" keys
+            Additional keys may include "artifacts" and "references" if available
+        """
+        artifacts = pending_artifacts if pending_artifacts else None
+        references = self._extract_references_from_state(final_state)
+        total_usage = self._extract_total_usage_from_state(final_state)
+        model_usage = self._extract_model_usage_from_state(final_state)
+        metadata = self._build_completion_metadata(final_state)
+
+        return self._create_a2a_event(
+            event_type=A2AStreamEventType.FINAL_RESPONSE,
+            content="Stream finished.",
+            tool_info=None,
+            metadata=metadata,
+            is_final=True,
+            artifacts=artifacts,
+            references=references,
+            step_usage=None,
+            total_usage=total_usage,
+            model_usage=model_usage,
+        )
+
+    def _extract_tool_name_prefix(self, tool_name: str) -> str:
+        """Extract a meaningful prefix from a tool name.
+
+        Args:
+            tool_name: The name of the tool.
+
+        Returns:
+            A meaningful prefix.
+        """
+        if tool_name.startswith("delegate_to_"):
+            agent_name = tool_name[12:]
+            if agent_name.endswith("Agent"):
+                agent_name = agent_name[:-5]
+            return agent_name.lower()[:4]
+
+        if "_" in tool_name:
+            parts = tool_name.split("_")
+            for part in parts:
+                if part not in ["tool", "generator", "calculator", "forecast"]:
+                    return part[:4]
+            return parts[0][:4]
+        else:
+            return tool_name[:4]
+
+    def _generate_tool_call_step_id(self, tool_info: dict[str, Any] | None, counter: int) -> str:
+        """Generate step_id for tool call events.
+
+        Args:
+            tool_info: Tool information
+            counter: Step counter
+
+        Returns:
+            Generated step_id
+        """
+        if not tool_info or not tool_info.get("tool_calls"):
+            return f"tool_start_{counter:03d}"
+
+        tool_calls = tool_info["tool_calls"]
+        if not tool_calls:
+            return f"tool_start_{counter:03d}"
+
+        prefixes = [self._extract_tool_name_prefix(tc.get("name", "")) or "unkn" for tc in tool_calls]
+        delegation_flags = self._get_delegation_info_from_tool_calls(tool_calls)
+
+        if len(tool_calls) == 1:
+            category = "agent" if delegation_flags[0] else "tool"
+            return f"{category}_{prefixes[0]}_start_{counter:03d}"
+
+        combined_name = "".join(prefixes).strip()[:6]
+        combined_name = combined_name or "multi"
+
+        if all(delegation_flags):
+            category = "agent"
+        elif any(delegation_flags):
+            category = "mixed"
+        else:
+            category = "tool"
+
+        return f"{category}_{combined_name}_parent_{counter:03d}"
+
+    def _generate_tool_result_step_id(self, tool_info: dict[str, Any] | None, counter: int) -> str:
+        """Generate step_id for tool result events.
+
+        Args:
+            tool_info: Tool information
+            counter: Step counter
+
+        Returns:
+            Generated step_id
+        """
+        if not tool_info:
+            return f"tool_done_{counter:03d}"
+
+        tool_name = tool_info.get("name", "")
+        prefix = self._extract_tool_name_prefix(tool_name) or "unkn"
+        category = "agent" if self._is_delegation_tool_from_info(tool_info) else "tool"
+        return f"{category}_{prefix}_done_{counter:03d}"
+
+    @staticmethod
+    def _is_delegation_tool_name(tool_name: str) -> bool:
+        """Check if a tool name corresponds to a delegation (sub-agent) tool.
+
+        This method maintains backward compatibility by checking the tool name pattern.
+        For new tools created by DelegationToolManager, use _is_delegation_tool() instead.
+
+        Args:
+            tool_name: The name of the tool to check.
+
+        Returns:
+            bool: True if the tool name indicates a delegation tool.
+        """
+        return isinstance(tool_name, str) and tool_name.startswith("delegate_to_")
+
+    @staticmethod
+    def _is_delegation_tool(tool_instance: Any) -> bool:
+        """Check delegation status based on metadata when available.
+
+        Args:
+            tool_instance: The tool instance to check for delegation metadata.
+
+        Returns:
+            True if the tool is marked as a delegation tool, False otherwise.
+        """
+        metadata = getattr(tool_instance, "metadata", None)
+        if not metadata or not hasattr(metadata, "get"):
+            return False
+
+        return bool(metadata.get("is_delegation_tool"))
+
+    def _get_delegation_info_from_tool_calls(self, tool_calls: list[dict[str, Any]] | None) -> list[bool]:
+        """Return delegation flags for each tool call using hybrid detection.
+
+        Args:
+            tool_calls: List of tool call dictionaries containing tool information.
+
+        Returns:
+            List of boolean flags indicating delegation status for each tool call.
+        """
+        if not tool_calls:
+            return []
+
+        delegation_flags: list[bool] = []
+        for tc in tool_calls:
+            if not isinstance(tc, dict):
+                logger.warning("Unexpected tool call payload type: %s", type(tc))
+                delegation_flags.append(False)
+                continue
+
+            delegation_flags.append(self._is_delegation_tool_from_info(tc))
+
+        return delegation_flags
+
+    def _is_delegation_tool_from_info(self, tool_info: dict[str, Any] | None) -> bool:
+        """Check delegation status from tool metadata, fallback to name pattern.
+
+        Args:
+            tool_info: Dictionary containing tool information including name and instance.
+
+        Returns:
+            True if the tool is identified as a delegation tool, False otherwise.
+        """
+        if not isinstance(tool_info, dict):
+            logger.warning("Unexpected tool info payload type: %s", type(tool_info))
+            return False
+
+        tool_instance = tool_info.get("tool_instance")
+        if tool_instance and self._is_delegation_tool(tool_instance):
+            return True
+
+        return self._is_delegation_tool_name(tool_info.get("name", ""))
+
+    def _generate_meaningful_step_id(
+        self, event_type: A2AStreamEventType, tool_info: dict[str, Any] | None = None
+    ) -> str:
+        """Generate a meaningful step_id based on event type and tool information.
+
+        Args:
+            event_type: The type of event (tool_call, tool_result, final_response, etc.)
+            tool_info: Tool information containing tool names and IDs
+
+        Returns:
+            A meaningful step_id string
+        """
+        try:
+            counter = get_next_step_number()
+
+            step_id_generators = {
+                A2AStreamEventType.TOOL_CALL: lambda: self._generate_tool_call_step_id(tool_info, counter),
+                A2AStreamEventType.TOOL_RESULT: lambda: self._generate_tool_result_step_id(tool_info, counter),
+                A2AStreamEventType.FINAL_RESPONSE: lambda: f"final_{counter:03d}",
+                A2AStreamEventType.CONTENT_CHUNK: lambda: f"content_{counter:03d}",
+            }
+
+            generator = step_id_generators.get(event_type)
+            if generator:
+                return generator()
+
+            event_value = event_type.value if hasattr(event_type, "value") else str(event_type)
+            return f"{event_value}_{counter:03d}"
+
+        except Exception:
+            return f"stp_{uuid.uuid4().hex[:8]}"
+
+    def _create_a2a_event(  # noqa: PLR0913
+        self,
+        event_type: A2AStreamEventType,
+        content: Any,
+        metadata: dict[str, Any] | None = None,
+        tool_info: dict[str, Any] | None = None,
+        thinking_and_activity_info: dict[str, Any] | None = None,
+        is_final: bool = False,
+        artifacts: list | None = None,
+        references: list | None = None,
+        step_usage: dict[str, Any] | None = None,
+        total_usage: dict[str, Any] | None = None,
+        model_usage_delta: dict[str, Any] | None = None,
+        model_usage: list[dict[str, Any]] | None = None,
+    ) -> A2AEvent:
+        """Create a structured A2AEvent dictionary.
+
+        Args:
+            event_type: The semantic type of the event.
+            content: The main text content of the event.
+            metadata: Additional metadata.
+            tool_info: Tool-specific information.
+            thinking_and_activity_info: Thinking and activity info from the model.
+            is_final: Whether this is a final event.
+            artifacts: List of artifacts to attach to the event.
+            references: List of references to attach to the event.
+            step_usage: Step-level token usage information.
+            total_usage: Total token usage information.
+            model_usage_delta: Per-call model usage delta.
+            model_usage: Final aggregated model usage.
+
+        Returns:
+            A dictionary conforming to the A2AEvent TypedDict.
+        """
+        enriched_metadata: dict[str, Any] = metadata.copy() if isinstance(metadata, dict) else {}
+        if "agent_name" not in enriched_metadata:
+            enriched_metadata["agent_name"] = self.name
+        if "step_id" not in enriched_metadata:
+            enriched_metadata["step_id"] = self._generate_meaningful_step_id(event_type, tool_info)
+        if "previous_step_ids" not in enriched_metadata:
+            enriched_metadata["previous_step_ids"] = []
+
+        # Inject cumulative time since the first STATUS_UPDATE for this thread
+        # Do not set cumulative time here; server executor enforces it for all SSE events
+
+        normalized_content = self._normalize_event_content(content)
+
+        event = {
+            "event_type": event_type,
+            "content": normalized_content,
+            "metadata": enriched_metadata,
+            "tool_info": tool_info,
+            "is_final": is_final,
+            "artifacts": artifacts,
+            "references": references,
+            STEP_USAGE_KEY: step_usage,
+            TOTAL_USAGE_KEY: total_usage,
+            MODEL_USAGE_DELTA_KEY: model_usage_delta,
+            MODEL_USAGE_KEY: model_usage,
+        }
+
+        if thinking_and_activity_info is not None:
+            event["thinking_and_activity_info"] = thinking_and_activity_info
+
+        try:
+            content_preview = normalized_content
+            logger.info(
+                "A2A emitting event: type=%s step_id=%s final=%s preview=%s",
+                getattr(event_type, "value", event_type),
+                enriched_metadata.get("step_id"),
+                is_final,
+                content_preview[:120].replace("\n", " "),
+            )
+        except Exception:
+            logger.debug("A2A emitting event (logging preview failed)", exc_info=True)
+
+        return event
+
+    def _resolve_tool_event_type(self, event_type_raw: Any) -> A2AStreamEventType | None:
+        """Normalize a raw event type to ``A2AStreamEventType``.
+
+        Args:
+            event_type_raw: Raw ``event_type`` value from a streaming chunk.
+
+        Returns:
+            The resolved ``A2AStreamEventType`` when supported, otherwise ``None``.
+        """
+        if isinstance(event_type_raw, A2AStreamEventType):
+            return event_type_raw
+        if isinstance(event_type_raw, str):
+            try:
+                return A2AStreamEventType(event_type_raw)
+            except ValueError:
+                return None
+        return None
+
+    @staticmethod
+    def _is_supported_tool_event(event_type: A2AStreamEventType) -> bool:
+        """Return True when the event type is a tool-related streaming event.
+
+        Args:
+            event_type: Candidate event type to evaluate.
+
+        Returns:
+            True when the event type should be forwarded to the client.
+        """
+        return event_type in {
+            A2AStreamEventType.TOOL_CALL,
+            A2AStreamEventType.TOOL_RESULT,
+            A2AStreamEventType.STATUS_UPDATE,
+        }
+
+    def _build_tool_activity_payload(
+        self,
+        event_type: A2AStreamEventType,
+        metadata: dict[str, Any] | None,
+        tool_info: dict[str, Any] | None,
+        activity_info: dict[str, Any] | None,
+    ) -> dict[str, Any] | None:
+        """Ensure tool events carry activity payloads per the streaming contract.
+
+        Args:
+            event_type: Stream event type emitted by the tool.
+            metadata: Optional metadata accompanying the chunk.
+            tool_info: Tool details provided by the emitting runner.
+            activity_info: Pre-built activity payload to reuse when present.
+
+        Returns:
+            Activity dictionary ready to be serialized with the tool chunk.
+        """
+        if event_type not in (A2AStreamEventType.TOOL_CALL, A2AStreamEventType.TOOL_RESULT):
+            return activity_info
+
+        if activity_info:
+            return activity_info
+
+        activity_context = self._compose_tool_activity_context(metadata, tool_info)
+        return create_tool_activity_info(activity_context)
+
+    def _compose_tool_activity_context(
+        self,
+        metadata: dict[str, Any] | None,
+        tool_info: dict[str, Any] | None,
+    ) -> dict[str, Any] | None:
+        """Create a context dictionary for downstream activity message generation.
+
+        Args:
+            metadata: Metadata payload extracted from the streaming chunk.
+            tool_info: Tool descriptor containing ids and display names.
+
+        Returns:
+            A merged context dictionary or None when no data was provided.
+        """
+        activity_context: dict[str, Any] | None = None
+        if isinstance(metadata, dict):
+            activity_context = metadata.copy()
+        if isinstance(tool_info, dict):
+            if activity_context is None:
+                activity_context = {"tool_info": tool_info}
+            else:
+                activity_context.setdefault("tool_info", tool_info)
+        return activity_context
+
+    def _create_tool_streaming_event(self, chunk: dict[str, Any], writer: StreamWriter, tool_name: str) -> None:
+        """Create and emit tool streaming events.
+
+        Only processes TOOL_CALL and TOOL_RESULT event types.
+
+        Args:
+            chunk: Streaming chunk from the tool.
+            writer: Stream writer to emit events.
+            tool_name: Name of the tool producing the chunk.
+        """
+        event_type = self._resolve_tool_event_type(chunk.get("event_type"))
+        if not event_type or not self._is_supported_tool_event(event_type):
+            return
+
+        tool_info = chunk.get("tool_info")
+        metadata = chunk.get("metadata")
+
+        if (
+            event_type == A2AStreamEventType.TOOL_RESULT
+            and isinstance(tool_info, dict)
+            and not tool_info.get("id")
+            and isinstance(tool_info.get("tool_calls"), list)
+            and tool_info.get("tool_calls")
+        ):
+            logger.info(
+                "A2A skipping streaming tool_result without id (tool=%s)",
+                tool_info.get("name"),
+            )
+            return
+
+        activity_info = self._build_tool_activity_payload(
+            event_type,
+            metadata if isinstance(metadata, dict) else None,
+            tool_info if isinstance(tool_info, dict) else None,
+            chunk.get("thinking_and_activity_info"),
+        )
+
+        a2a_event = self._create_a2a_event(
+            event_type=event_type,
+            content=chunk.get("content", f"Processing with tools: {tool_name}"),
+            metadata=metadata,
+            tool_info=tool_info,
+            thinking_and_activity_info=activity_info,
+        )
+        writer(a2a_event)
+
+    # ------------------------------------------------------------------
+    # Context compaction integration point (feature-flagged, SLICE-001)
+    # ------------------------------------------------------------------
+
+    @property
+    def _compaction_config(self) -> dict[str, Any] | None:
+        """Return compaction config if enabled, else None.
+
+        Feature flag: returns None by default until #830 SDK config lands.
+        Subclasses override to enable context compaction at the model
+        invocation boundary.
+        """
+        return None
+
+    async def _emit_compaction_event(
+        self,
+        stage: str,
+        *,
+        content: str = "",
+        metadata: dict[str, Any] | None = None,
+    ) -> None:
+        """Emit a compaction lifecycle A2A STATUS_UPDATE event if streaming.
+
+        Uses the LangGraph stream writer when available (streaming path),
+        falls back to the agent event_emitter (connector streaming path),
+        and silently no-ops during non-streaming invoke where neither is
+        available.
+
+        Args:
+            stage: Compaction lifecycle stage (started, completed, skipped, failed).
+            content: Short safe content text for the event.
+            metadata: Safe metadata dict from build_compaction_event_metadata.
+        """
+        from aip_agents.context.metadata import build_compaction_event_metadata
+
+        event_metadata: dict[str, Any] = {
+            **build_compaction_event_metadata(stage=stage),
+            **dict(metadata or {}),
+        }
+        event_metadata.setdefault("step_type", "context_compaction")
+        event_metadata.setdefault("stage", stage)
+
+        try:
+            from langgraph.config import get_stream_writer
+
+            writer = get_stream_writer()
+        except Exception:
+            writer = None
+
+        if writer is not None:
+            a2a_event = self._create_a2a_event(
+                event_type=A2AStreamEventType.STATUS_UPDATE,
+                content=content,
+                metadata=event_metadata,
+            )
+            try:
+                writer(a2a_event)
+            except Exception:
+                logger.debug("Compaction event writer failed (non-streaming path expected)", exc_info=True)
+            return
+
+        if self.event_emitter is not None:
+            try:
+                emitted = self.event_emitter.emit(
+                    Event(
+                        type=A2AStreamEventType.STATUS_UPDATE.value,
+                        value=content,
+                        metadata=event_metadata,
+                    )
+                )
+                if inspect.isawaitable(emitted):
+                    await emitted
+            except Exception:
+                logger.debug("Compaction event_emitter failed (non-streaming path expected)", exc_info=True)
+
+    async def _apply_compaction(
+        self,
+        messages: Sequence[Any],
+        *,
+        recall: Any | None = None,
+        provider_overflow: bool = False,
+        preserve_topic: str | None = None,
+    ) -> tuple[list[Any], Any | None]:
+        """Apply context compaction before model invocation.
+
+        No-op when compaction is not configured. Subclasses may override
+        to wire the full prepare_context() pipeline.
+
+        Args:
+            messages: Current message sequence (OpenAI-format dicts).
+            recall: Optional session recall snapshot.
+            provider_overflow: If True, force compaction even below
+                the normal threshold (provider-overflow retry path).
+            preserve_topic: Optional topic to preserve through compaction.
+
+        Returns:
+            Tuple of (messages, preparation_result_or_None).
+        """
+        cfg = self._compaction_config
+        if cfg is None:
+            return list(messages), None
+
+        from aip_agents.context.budget import ContextBudgetConfig
+        from aip_agents.context.metadata import (
+            build_compaction_metadata,
+            safe_compaction_content,
+        )
+        from aip_agents.context.preparation import prepare_context
+
+        budget_config = ContextBudgetConfig(
+            threshold_ratio=cfg.get("threshold_ratio", 0.85),
+            fallback_threshold_tokens=cfg.get("fallback_threshold_tokens", 170_000),
+            min_tail_tokens=cfg.get("min_tail_tokens", 6_000),
+            target_tail_ratio=cfg.get("target_tail_ratio", 0.10),
+            min_head_messages=cfg.get("min_head_messages", 2),
+        )
+
+        await self._emit_compaction_event("started", content=safe_compaction_content("started"))
+
+        effective_focus = preserve_topic or (getattr(recall, "active_task", None) if recall else None)
+        result = prepare_context(
+            messages,
+            config=budget_config,
+            focus_topic=effective_focus,
+            preserve_topic=preserve_topic,
+            provider_overflow=provider_overflow,
+        )
+
+        await self._emit_compaction_result_event(result)
+        self._last_compaction_preparation = build_compaction_metadata(
+            result.state,
+            tokens_before=result.summary.prompt_tokens_before if result.summary else None,
+            tokens_after=result.summary.prompt_tokens_after if result.summary else None,
+            summary_id_for_state=result.summary.summary_id if result.summary else None,
+        )
+        return result.messages, result
+
+    async def _emit_compaction_result_event(
+        self,
+        result: Any,
+    ) -> None:
+        """Emit compaction lifecycle event based on preparation result state."""
+        from aip_agents.context.metadata import build_compaction_event_metadata, safe_compaction_content
+
+        state = result.state
+        summary = result.summary
+        decision = result.decision
+
+        if state is None:
+            metadata = build_compaction_event_metadata(stage="failed", summary_status="none")
+            await self._emit_compaction_event("failed", content=safe_compaction_content("failed"), metadata=metadata)
+            return
+
+        if state.compaction_count > 0:
+            stage = "failed" if state.last_summary_status == "failed" else "completed"
+            stage_content = safe_compaction_content(stage)
+            tokens_before = (
+                summary.prompt_tokens_before if summary is not None and summary.prompt_tokens_before > 0 else None
+            )
+            tokens_after = (
+                summary.prompt_tokens_after if summary is not None and summary.prompt_tokens_after > 0 else None
+            )
+            metadata = build_compaction_event_metadata(
+                stage=stage,
+                trigger_reason=state.last_reason,
+                summary_status=state.last_summary_status,
+                tokens_before=tokens_before,
+                tokens_after=tokens_after,
+                offloaded_ref_count=state.offloaded_ref_count,
+                summary_id=summary.summary_id if summary is not None else None,
+            )
+            await self._emit_compaction_event(stage, content=stage_content, metadata=metadata)
+            return
+
+        metadata = build_compaction_event_metadata(
+            stage="skipped",
+            trigger_reason=decision.reason if decision else None,
+            summary_status=state.last_summary_status,
+        )
+        await self._emit_compaction_event("skipped", content=safe_compaction_content("skipped"), metadata=metadata)
+
+    @staticmethod
+    def _append_recall_guidance(instruction: str, recall_guidance: str | None) -> str:
+        """Append session recall guidance to the current instruction.
+
+        Preserves any middleware edits already applied to the instruction.
+        No-op when recall_guidance is empty or None.
+        """
+        if not recall_guidance:
+            return instruction
+        return f"{instruction}\n\n{recall_guidance}"
+
+    @staticmethod
+    def _base_message_to_dict(message: BaseMessage) -> dict[str, Any]:
+        """Convert a LangChain BaseMessage to an OpenAI-format dict."""
+        role_map = {
+            "human": "user",
+            "ai": "assistant",
+            "tool": "tool",
+            "system": "system",
+        }
+        d: dict[str, Any] = {
+            "role": role_map.get(message.type, "user"),
+            "content": message.content,
+            "type": message.type,
+        }
+        if message.additional_kwargs is not None:
+            d["additional_kwargs"] = message.additional_kwargs
+        if message.response_metadata is not None:
+            d["response_metadata"] = message.response_metadata
+        if message.id is not None:
+            d["id"] = message.id
+        if message.name is not None:
+            d["name"] = message.name
+        if getattr(message, "usage_metadata", None) is not None:
+            d["usage_metadata"] = message.usage_metadata
+        if getattr(message, "invalid_tool_calls", None) is not None:
+            d["invalid_tool_calls"] = message.invalid_tool_calls
+        if getattr(message, "artifact", None) is not None:
+            d["artifact"] = message.artifact
+        if getattr(message, "status", None) is not None:
+            d["status"] = message.status
+        if isinstance(message, ToolMessage):
+            d["tool_call_id"] = message.tool_call_id
+        if isinstance(message, AIMessage) and message.tool_calls:
+            d["tool_calls"] = [
+                {
+                    "id": tc.get("id", ""),
+                    "type": "function",
+                    "function": {
+                        "name": tc.get("name", ""),
+                        "arguments": json.dumps(tc.get("args", {}), ensure_ascii=False),
+                    },
+                }
+                for tc in message.tool_calls
+            ]
+        return d
+
+    @staticmethod
+    def _parse_tool_call_dicts(tool_calls: list[dict[str, Any]]) -> list[dict[str, Any]]:
+        """Parse OpenAI-format tool call dicts into LangChain tool call format."""
+        lc_tool_calls: list[dict[str, Any]] = []
+        for tc in tool_calls:
+            func = tc.get("function", {})
+            args_raw = func.get("arguments", "{}")
+            if isinstance(args_raw, str):
+                try:
+                    args_parsed = json.loads(args_raw)
+                except (json.JSONDecodeError, TypeError):
+                    args_parsed = {}
+            else:
+                args_parsed = args_raw
+            lc_tool_calls.append(
+                {
+                    "id": tc.get("id", ""),
+                    "name": func.get("name", ""),
+                    "args": args_parsed,
+                }
+            )
+        return lc_tool_calls
+
+    @staticmethod
+    def _message_constructor_kwargs(
+        message_cls: type[BaseMessage],
+        d: dict[str, Any],
+    ) -> dict[str, Any]:
+        """Return constructor kwargs shared by message classes."""
+        kwargs: dict[str, Any] = {"content": d.get("content", "")}
+        allowed_fields = set(message_cls.model_fields.keys())
+        for key in (
+            "additional_kwargs",
+            "response_metadata",
+            "id",
+            "name",
+            "usage_metadata",
+            "invalid_tool_calls",
+            "artifact",
+            "status",
+        ):
+            if key not in allowed_fields:
+                continue
+            if key in d:
+                kwargs[key] = d[key]
+        return kwargs
+
+    @staticmethod
+    def _build_message(message_cls: type[BaseMessage], **kwargs: Any) -> BaseMessage:
+        """Build a message, preserving provider payloads LangChain rejects at init."""
+        if isinstance(kwargs.get("content"), dict):
+            return cast(BaseMessage, message_cls.model_construct(**kwargs))
+        try:
+            return message_cls(**kwargs)
+        except ValidationError:
+            return cast(BaseMessage, message_cls.model_construct(**kwargs))
+
+    @staticmethod
+    def _dict_to_base_message(d: dict[str, Any]) -> BaseMessage:
+        """Convert an OpenAI-format dict back to a LangChain BaseMessage."""
+        role = d.get("role", "user")
+        if role == "tool":
+            return BaseLangGraphAgent._build_message(
+                ToolMessage,
+                **BaseLangGraphAgent._message_constructor_kwargs(ToolMessage, d),
+                tool_call_id=d.get("tool_call_id", ""),
+            )
+        if role == "assistant":
+            tool_calls = d.get("tool_calls")
+            if tool_calls:
+                return BaseLangGraphAgent._build_message(
+                    AIMessage,
+                    **BaseLangGraphAgent._message_constructor_kwargs(AIMessage, d),
+                    tool_calls=BaseLangGraphAgent._parse_tool_call_dicts(tool_calls),
+                )
+            return BaseLangGraphAgent._build_message(
+                AIMessage,
+                **BaseLangGraphAgent._message_constructor_kwargs(AIMessage, d),
+            )
+        if role == "system":
+            return BaseLangGraphAgent._build_message(
+                SystemMessage,
+                **BaseLangGraphAgent._message_constructor_kwargs(SystemMessage, d),
+            )
+        return BaseLangGraphAgent._build_message(
+            HumanMessage,
+            **BaseLangGraphAgent._message_constructor_kwargs(HumanMessage, d),
+        )
+
+    async def _compact_messages(
+        self,
+        message_dicts: list[dict[str, Any]],
+        system_content: str | None,
+        recall: Any,
+    ) -> list[dict[str, Any]]:
+        """Apply compaction to message dicts; return conversation-only messages.
+
+        The system prompt may be included for budget estimation, but is always
+        stripped from the result — the caller owns the system prompt separately.
+        """
+        budget_messages = (
+            [{"role": "system", "content": system_content}] + message_dicts if system_content else message_dicts
+        )
+        compacted_dicts, _ = await self._apply_compaction(budget_messages, recall=recall)
+        return [d for d in compacted_dicts if d.get("role") != "system"]
+
+    async def _compact_conversation_with_recall(
+        self,
+        message_dicts: list[dict[str, Any]],
+        system_content: str,
+    ) -> tuple[list[dict[str, Any]], str]:
+        """Apply recall guidance and compaction to conversation messages."""
+        from aip_agents.context.runtime import derive_session_recall
+        from aip_agents.context.session_recall import build_session_recall_guidance
+
+        recall = derive_session_recall(message_dicts)
+        recall_guidance = build_session_recall_guidance(recall)
+        if recall_guidance:
+            system_content = self._append_recall_guidance(system_content, recall_guidance)
+
+        compacted_conversation = await self._compact_messages(
+            message_dicts,
+            system_content,
+            recall,
+        )
+        return compacted_conversation, system_content