Update Claude agent connector layering

Author: eavanvalkenburg

python/packages/a2a/agent_framework_a2a/_agent.py

@@ -29,7 +29,7 @@
     AgentResponse,
     AgentResponseUpdate,
     AgentThread,
-    BareAgent,
+    BaseAgent,
     ChatMessage,
     Content,
     ResponseStream,
@@ -58,12 +58,12 @@ def _get_uri_data(uri: str) -> str:
     return match.group("base64_data")
 
 
-class A2AAgent(AgentTelemetryLayer, BareAgent):
+class A2AAgent(AgentTelemetryLayer, BaseAgent):
     """Agent2Agent (A2A) protocol implementation.
 
     Wraps an A2A Client to connect the Agent Framework with external A2A-compliant agents
     via HTTP/JSON-RPC. Converts framework ChatMessages to A2A Messages on send, and converts
-    A2A responses (Messages/Tasks) back to framework types. Inherits BareAgent capabilities
+    A2A responses (Messages/Tasks) back to framework types. Inherits BaseAgent capabilities
     while managing the underlying A2A protocol communication.
 
     Can be initialized with a URL, AgentCard, or existing A2A Client instance.
@@ -99,7 +99,7 @@ def __init__(
             timeout: Request timeout configuration. Can be a float (applied to all timeout components),
                 httpx.Timeout object (for full control), or None (uses 10.0s connect, 60.0s read,
                 10.0s write, 5.0s pool - optimized for A2A operations).
-            kwargs: any additional properties, passed to BareAgent.
+            kwargs: any additional properties, passed to BaseAgent.
         """
         super().__init__(id=id, name=name, description=description, **kwargs)
         self._http_client: httpx.AsyncClient | None = http_client

python/packages/ag-ui/tests/__init__.py

@@ -0,0 +1 @@
+# Copyright (c) Microsoft. All rights reserved.

python/packages/claude/agent_framework_claude/_agent.py

@@ -2,12 +2,12 @@
 
 import contextlib
 import sys
-from collections.abc import AsyncIterable, Callable, MutableMapping, Sequence
+from collections.abc import AsyncIterable, Awaitable, Callable, MutableMapping, Sequence
 from pathlib import Path
-from typing import TYPE_CHECKING, Any, ClassVar, Generic
+from typing import TYPE_CHECKING, Any, ClassVar, Generic, Literal, overload
 
 from agent_framework import (
-    AgentMiddlewareTypes,
+    AgentMiddlewareLayer,
     AgentResponse,
     AgentResponseUpdate,
     AgentThread,
@@ -16,13 +16,16 @@
     Content,
     ContextProvider,
     FunctionTool,
+    ResponseStream,
     Role,
     ToolProtocol,
     get_logger,
+    merge_chat_options,
     normalize_messages,
+    normalize_tools,
 )
-from agent_framework._types import normalize_tools
 from agent_framework.exceptions import ServiceException, ServiceInitializationError
+from agent_framework.observability import AgentTelemetryLayer
 from claude_agent_sdk import (
     ClaudeAgentOptions as SDKOptions,
 )
@@ -145,7 +148,7 @@ class ClaudeAgentOptions(TypedDict, total=False):
 )
 
 
-class ClaudeAgent(BaseAgent, Generic[TOptions]):
+class RawClaudeAgent(BaseAgent, Generic[TOptions]):
     """Claude Agent using Claude Code CLI.
 
     Wraps the Claude Agent SDK to provide agentic capabilities including
@@ -175,7 +178,7 @@ class ClaudeAgent(BaseAgent, Generic[TOptions]):
         .. code-block:: python
 
             async with ClaudeAgent() as agent:
-                async for update in agent.run_stream("Write a poem"):
+                async for update in agent.run("Write a poem", stream=True):
                     print(update.text, end="", flush=True)
 
         With session management:
@@ -214,7 +217,6 @@ def __init__(
         name: str | None = None,
         description: str | None = None,
         context_provider: ContextProvider | None = None,
-        middleware: Sequence[AgentMiddlewareTypes] | None = None,
         tools: ToolProtocol
         | Callable[..., Any]
         | MutableMapping[str, Any]
@@ -224,8 +226,9 @@ def __init__(
         default_options: TOptions | MutableMapping[str, Any] | None = None,
         env_file_path: str | None = None,
         env_file_encoding: str | None = None,
+        **kwargs: Any,
     ) -> None:
-        """Initialize a ClaudeAgent instance.
+        """Initialize a Claude agent instance.
 
         Args:
             instructions: System prompt for the agent.
@@ -237,20 +240,20 @@ def __init__(
             name: Name of the agent.
             description: Description of the agent.
             context_provider: Context provider for the agent.
-            middleware: List of middleware.
             tools: Tools for the agent. Can be:
                 - Strings for built-in tools (e.g., "Read", "Write", "Bash", "Glob")
                 - Functions or ToolProtocol instances for custom tools
             default_options: Default ClaudeAgentOptions including system_prompt, model, etc.
             env_file_path: Path to .env file.
             env_file_encoding: Encoding of .env file.
+            kwargs: Additional keyword arguments passed to BaseAgent.
         """
         super().__init__(
             id=id,
             name=name,
             description=description,
             context_provider=context_provider,
-            middleware=middleware,
+            **kwargs,
         )
 
         self._client = client
@@ -295,6 +298,11 @@ def __init__(
         self._started = False
         self._current_session_id: str | None = None
 
+    @property
+    def default_options(self) -> dict[str, Any]:
+        """Expose default options for telemetry and middleware layers."""
+        return dict(self._default_options)
+
     def _normalize_tools(
         self,
         tools: ToolProtocol
@@ -328,7 +336,7 @@ def _normalize_tools(
                 normalized = normalize_tools(tool)
                 self._custom_tools.extend(normalized)
 
-    async def __aenter__(self) -> "ClaudeAgent[TOptions]":
+    async def __aenter__(self) -> "RawClaudeAgent[TOptions]":
         """Start the agent when entering async context."""
         await self.start()
         return self
@@ -549,55 +557,80 @@ def _format_prompt(self, messages: list[ChatMessage] | None) -> str:
             return ""
         return "\n".join([msg.text or "" for msg in messages])
 
-    async def run(
+    @overload
+    def run(
         self,
         messages: str | ChatMessage | Sequence[str | ChatMessage] | None = None,
         *,
+        stream: Literal[False] = ...,
         thread: AgentThread | None = None,
-        options: TOptions | MutableMapping[str, Any] | None = None,
         **kwargs: Any,
-    ) -> AgentResponse[Any]:
-        """Run the agent with the given messages.
+    ) -> Awaitable[AgentResponse[Any]]: ...
 
-        Args:
-            messages: The messages to process.
+    @overload
+    def run(
+        self,
+        messages: str | ChatMessage | Sequence[str | ChatMessage] | None = None,
+        *,
+        stream: Literal[True],
+        thread: AgentThread | None = None,
+        **kwargs: Any,
+    ) -> ResponseStream[AgentResponseUpdate, AgentResponse[Any]]: ...
 
-        Keyword Args:
-            thread: The conversation thread. If thread has service_thread_id set,
-                the agent will resume that session.
-            options: Runtime options (model, permission_mode can be changed per-request).
-            kwargs: Additional keyword arguments.
+    def run(
+        self,
+        messages: str | ChatMessage | Sequence[str | ChatMessage] | None = None,
+        *,
+        stream: bool = False,
+        thread: AgentThread | None = None,
+        **kwargs: Any,
+    ) -> Awaitable[AgentResponse[Any]] | ResponseStream[AgentResponseUpdate, AgentResponse[Any]]:
+        """Run the agent with the given messages."""
+        options = kwargs.pop("options", None)
+        if stream:
+
+            def _finalize(updates: Sequence[AgentResponseUpdate]) -> AgentResponse[Any]:
+                response = AgentResponse.from_agent_run_response_updates(updates)
+                session_id = _get_session_id_from_updates(updates)
+                if session_id and thread is not None:
+                    thread.service_thread_id = session_id
+                return response
+
+            return ResponseStream(
+                self._stream_updates(messages=messages, thread=thread, options=options, **kwargs),
+                finalizer=_finalize,
+            )
 
-        Returns:
-            AgentResponse with the agent's response.
-        """
+        return self._run_impl(messages=messages, thread=thread, options=options, **kwargs)
+
+    async def _run_impl(
+        self,
+        messages: str | ChatMessage | Sequence[str | ChatMessage] | None = None,
+        *,
+        thread: AgentThread | None = None,
+        options: TOptions | MutableMapping[str, Any] | None = None,
+        **kwargs: Any,
+    ) -> AgentResponse[Any]:
+        """Non-streaming implementation of run."""
         thread = thread or self.get_new_thread()
-        return await AgentResponse.from_agent_response_generator(
-            self.run_stream(messages, thread=thread, options=options, **kwargs)
-        )
+        updates: list[AgentResponseUpdate] = []
+        async for update in self._stream_updates(messages=messages, thread=thread, options=options, **kwargs):
+            updates.append(update)
+        response = AgentResponse.from_agent_run_response_updates(updates)
+        session_id = _get_session_id_from_updates(updates)
+        if session_id:
+            thread.service_thread_id = session_id
+        return response
 
-    async def run_stream(
+    async def _stream_updates(
         self,
         messages: str | ChatMessage | Sequence[str | ChatMessage] | None = None,
         *,
         thread: AgentThread | None = None,
         options: TOptions | MutableMapping[str, Any] | None = None,
         **kwargs: Any,
     ) -> AsyncIterable[AgentResponseUpdate]:
-        """Stream the agent's response.
-
-        Args:
-            messages: The messages to process.
-
-        Keyword Args:
-            thread: The conversation thread. If thread has service_thread_id set,
-                the agent will resume that session.
-            options: Runtime options (model, permission_mode can be changed per-request).
-            kwargs: Additional keyword arguments.
-
-        Yields:
-            AgentResponseUpdate objects containing chunks of the response.
-        """
+        """Stream the agent's response updates."""
         thread = thread or self.get_new_thread()
 
         # Ensure we're connected to the right session
@@ -606,12 +639,18 @@ async def run_stream(
         if not self._client:
             raise ServiceException("Claude SDK client not initialized.")
 
+        merged_options = merge_chat_options(
+            {"instructions": self._default_options.get("system_prompt")},
+            merge_chat_options(self._default_options, dict(options) if options else None),
+        )
+        runtime_options = dict(merged_options)
+        runtime_options.pop("system_prompt", None)
+        runtime_options.pop("instructions", None)
+
         prompt = self._format_prompt(normalize_messages(messages))
 
         # Apply runtime options (model, permission_mode)
-        await self._apply_runtime_options(dict(options) if options else None)
-
-        session_id: str | None = None
+        await self._apply_runtime_options(runtime_options if runtime_options else None)
 
         await self._client.query(prompt)
         async for message in self._client.receive_response():
@@ -638,8 +677,29 @@ async def run_stream(
                                 raw_representation=message,
                             )
             elif isinstance(message, ResultMessage):
-                session_id = message.session_id
-
-        # Update thread with session ID
-        if session_id:
-            thread.service_thread_id = session_id
+                if message.session_id:
+                    yield AgentResponseUpdate(
+                        role=Role.ASSISTANT,
+                        contents=[Content.from_text(text="", raw_representation=message)],
+                        raw_representation=message,
+                    )
+
+
+class ClaudeAgent(  # type: ignore[misc]
+    AgentTelemetryLayer,
+    AgentMiddlewareLayer,
+    RawClaudeAgent[TOptions],
+    Generic[TOptions],
+):
+    """Claude agent with middleware and telemetry layers applied."""
+
+    pass
+
+
+def _get_session_id_from_updates(updates: Sequence[AgentResponseUpdate]) -> str | None:
+    """Extract session_id from ResultMessage entries in updates."""
+    for update in updates:
+        raw = update.raw_representation
+        if isinstance(raw, ResultMessage):
+            return raw.session_id
+    return None

python/packages/claude/tests/test_claude_agent.py

@@ -312,7 +312,7 @@ async def test_run_with_thread(self) -> None:
 
 
 class TestClaudeAgentRunStream:
-    """Tests for ClaudeAgent run_stream method."""
+    """Tests for ClaudeAgent streaming run method."""
 
     @staticmethod
     async def _create_async_generator(items: list[Any]) -> Any:
@@ -332,7 +332,7 @@ def _create_mock_client(self, messages: list[Any]) -> MagicMock:
         return mock_client
 
     async def test_run_stream_yields_updates(self) -> None:
-        """Test run_stream yields AgentResponseUpdate objects."""
+        """Test run(stream=True) yields AgentResponseUpdate objects."""
         from claude_agent_sdk import AssistantMessage, ResultMessage, TextBlock
         from claude_agent_sdk.types import StreamEvent
 
@@ -371,10 +371,10 @@ async def test_run_stream_yields_updates(self) -> None:
         with patch("agent_framework_claude._agent.ClaudeSDKClient", return_value=mock_client):
             agent = ClaudeAgent()
             updates: list[AgentResponseUpdate] = []
-            async for update in agent.run_stream("Hello"):
+            async for update in agent.run("Hello", stream=True):
                 updates.append(update)
             # StreamEvent yields text deltas
-            assert len(updates) == 2
+            assert len(updates) == 3
             assert updates[0].role == Role.ASSISTANT
             assert updates[0].text == "Streaming "
             assert updates[1].text == "response"

python/packages/copilotstudio/agent_framework_copilotstudio/_agent.py

@@ -8,7 +8,7 @@
     AgentResponse,
     AgentResponseUpdate,
     AgentThread,
-    BareAgent,
+    BaseAgent,
     ChatMessage,
     Content,
     ContextProvider,
@@ -69,7 +69,7 @@ class CopilotStudioSettings(AFBaseSettings):
     tenantid: str | None = None
 
 
-class CopilotStudioAgent(BareAgent):
+class CopilotStudioAgent(BaseAgent):
     """A Copilot Studio Agent."""
 
     def __init__(