Python: Reduce ResponseStream hot-path latency

python/CHANGELOG.md

@@ -36,6 +36,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - **docs**, **samples**: Cross-reference security patterns, clarify file-search retrieval scope, improve local-shell filtering guidance, and correct README instructions ([#8726](https://github.com/microsoft/agent-framework/pull/8726), [#8591](https://github.com/microsoft/agent-framework/pull/8591), [#8766](https://github.com/microsoft/agent-framework/pull/8766), [#8822](https://github.com/microsoft/agent-framework/pull/8822))
 
 ### Fixed
+- **agent-framework-core**: Reduce response-stream hot-path overhead and avoid loading tool-approval middleware for agent runs without middleware ([#8971](https://github.com/microsoft/agent-framework/pull/8971))
 - **agent-framework-core**: Preserve functional-workflow replay identity, scope invocation arguments and activity IDs to workflow runs, deduplicate fan-out targets, propagate nested cancellation arguments, and support orchestration checkpoint deserialization ([#8887](https://github.com/microsoft/agent-framework/pull/8887), [#8596](https://github.com/microsoft/agent-framework/pull/8596), [#8549](https://github.com/microsoft/agent-framework/pull/8549), [#8739](https://github.com/microsoft/agent-framework/pull/8739), [#8602](https://github.com/microsoft/agent-framework/pull/8602), [#8723](https://github.com/microsoft/agent-framework/pull/8723))
 - **agent-framework-core**: Preserve repeated history turns, compaction through middleware rewrites, checkpoint dictionary subclasses, response metadata, annotations, native vector keys, integer precision, and boolean set values ([#8800](https://github.com/microsoft/agent-framework/pull/8800), [#8671](https://github.com/microsoft/agent-framework/pull/8671), [#8519](https://github.com/microsoft/agent-framework/pull/8519), [#8464](https://github.com/microsoft/agent-framework/pull/8464), [#8701](https://github.com/microsoft/agent-framework/pull/8701), [#8535](https://github.com/microsoft/agent-framework/pull/8535), [#8637](https://github.com/microsoft/agent-framework/pull/8637), [#8681](https://github.com/microsoft/agent-framework/pull/8681))
 - **agent-framework-core**: Improve MCP skill archive and resource validation, stop failed MCP lifecycle owners, distinguish missing tool arguments from parse failures, settle fatal unknown calls, and omit absent MCP prompt-argument descriptions ([#8937](https://github.com/microsoft/agent-framework/pull/8937), [#8690](https://github.com/microsoft/agent-framework/pull/8690), [#8755](https://github.com/microsoft/agent-framework/pull/8755), [#8609](https://github.com/microsoft/agent-framework/pull/8609), [#8640](https://github.com/microsoft/agent-framework/pull/8640), [#8733](https://github.com/microsoft/agent-framework/pull/8733))

python/packages/core/agent_framework/_agents.py

@@ -105,9 +105,12 @@
 
 def _tool_approval_source_ids(middleware: Sequence[MiddlewareTypes] | None) -> frozenset[str]:
     """Return session-state keys owned by ToolApprovalMiddleware instances."""
+    if not middleware:
+        return frozenset()
+
     from ._harness._tool_approval import ToolApprovalMiddleware
 
-    return frozenset(item.source_id for item in middleware or () if isinstance(item, ToolApprovalMiddleware))
+    return frozenset(item.source_id for item in middleware if isinstance(item, ToolApprovalMiddleware))
 
 
 def _merge_delegated_session_state(

python/packages/core/agent_framework/_types.py

@@ -3775,17 +3775,20 @@ async def _apply_transforms(
         return value, transformed
 
     async def _record_update(self, update: UpdateT, *, run_after_gates: bool) -> UpdateT:
-        await self._run_gates(
-            cast(Sequence[Callable[[Any], object]], self._update_gates_before),
-            update,
-            target="update",
-        )
+        if self._update_gates_before:
+            await self._run_gates(
+                cast(Sequence[Callable[[Any], object]], self._update_gates_before),
+                update,
+                target="update",
+            )
         self._updates.append(update)
-        update, _ = await self._apply_transforms(
-            cast(Sequence[Callable[[Any], Any | Awaitable[Any | None] | None]], self._transform_hooks),
-            update,
-        )
-        if run_after_gates:
+        for transform in self._transform_hooks:
+            transformed_update = transform(update)
+            if isawaitable(transformed_update):
+                transformed_update = await transformed_update
+            if transformed_update is not None:
+                update = cast(UpdateT, transformed_update)
+        if run_after_gates and self._update_gates_after:
             await self._run_gates(
                 cast(Sequence[Callable[[Any], object]], self._update_gates_after),
                 update,
@@ -3801,17 +3804,42 @@ async def _pull_next_update(self, *, run_after_gates: bool) -> UpdateT:
                     run_after_gates=run_after_gates,
                 )
 
-            with contextlib.ExitStack() as stack:
-                for factory in self._pull_context_manager_factories:
-                    stack.enter_context(factory())
-                # Resolve the underlying stream inside the pull contexts so that any
-                # spans/contexts created during stream resolution (e.g. inner chat
-                # completion spans created on the first pull of a wrapped agent stream)
-                # inherit the active context (e.g. an outer agent invoke span).
+            context_factories = self._pull_context_manager_factories
+            update: UpdateT
+            if not context_factories:
                 if self._iterator is None:
                     stream = await self._get_stream()
                     self._iterator = stream.__aiter__()
-                update: UpdateT = await self._iterator.__anext__()
+                update = await self._iterator.__anext__()
+            elif len(context_factories) == 1:
+                with context_factories[0]():
+                    if len(context_factories) == 1:
+                        if self._iterator is None:
+                            stream = await self._get_stream()
+                            self._iterator = stream.__aiter__()
+                        update = await self._iterator.__anext__()
+                    else:
+                        with contextlib.ExitStack() as stack:
+                            context_index = 1
+                            while context_index < len(context_factories):
+                                stack.enter_context(context_factories[context_index]())
+                                context_index += 1
+                            if self._iterator is None:
+                                stream = await self._get_stream()
+                                self._iterator = stream.__aiter__()
+                            update = await self._iterator.__anext__()
+            else:
+                with contextlib.ExitStack() as stack:
+                    for factory in context_factories:
+                        stack.enter_context(factory())
+                    # Resolve the underlying stream inside the pull contexts so that any
+                    # spans/contexts created during stream resolution (e.g. inner chat
+                    # completion spans created on the first pull of a wrapped agent stream)
+                    # inherit the active context (e.g. an outer agent invoke span).
+                    if self._iterator is None:
+                        stream = await self._get_stream()
+                        self._iterator = stream.__aiter__()
+                    update = await self._iterator.__anext__()
             if self._flat_map_update is not None:
                 mapped_updates = self._flat_map_update(update)
                 if isawaitable(mapped_updates):
@@ -3881,41 +3909,48 @@ async def _prepare_final_result(self) -> None:
         else:
             result = list(self._updates)
 
-        await self._run_gates(
-            cast(Sequence[Callable[[Any], object]], self._result_gates_before),
-            result,
-            target="result",
-        )
-        result, self._result_was_transformed = await self._apply_transforms(
-            cast(Sequence[Callable[[Any], Any | Awaitable[Any | None] | None]], self._result_hooks),
-            result,
-        )
+        if self._result_gates_before:
+            await self._run_gates(
+                cast(Sequence[Callable[[Any], object]], self._result_gates_before),
+                result,
+                target="result",
+            )
+        if self._result_hooks:
+            result, self._result_was_transformed = await self._apply_transforms(
+                cast(Sequence[Callable[[Any], Any | Awaitable[Any | None] | None]], self._result_hooks),
+                result,
+            )
+        else:
+            self._result_was_transformed = False
         self._final_result = result
         self._result_prepared = True
 
     async def _complete_final_result(self) -> None:
         if self._finalized:
             return
         await self._prepare_final_result()
-        await self._run_gates(
-            cast(Sequence[Callable[[Any], object]], self._result_gates_after),
-            self._final_result,
-            target="result",
-        )
-        terminal_result, transformed = await self._apply_transforms(
-            cast(
-                Sequence[Callable[[Any], Any | Awaitable[Any | None] | None]],
-                self._terminal_result_transforms,
-            ),
-            self._final_result,
-        )
-        self._final_result = terminal_result
-        self._result_was_transformed = self._result_was_transformed or transformed
-        await self._run_gates(
-            cast(Sequence[Callable[[Any], object]], self._terminal_result_gates),
-            self._final_result,
-            target="result",
-        )
+        if self._result_gates_after:
+            await self._run_gates(
+                cast(Sequence[Callable[[Any], object]], self._result_gates_after),
+                self._final_result,
+                target="result",
+            )
+        if self._terminal_result_transforms:
+            terminal_result, transformed = await self._apply_transforms(
+                cast(
+                    Sequence[Callable[[Any], Any | Awaitable[Any | None] | None]],
+                    self._terminal_result_transforms,
+                ),
+                self._final_result,
+            )
+            self._final_result = terminal_result
+            self._result_was_transformed = self._result_was_transformed or transformed
+        if self._terminal_result_gates:
+            await self._run_gates(
+                cast(Sequence[Callable[[Any], object]], self._terminal_result_gates),
+                self._final_result,
+                target="result",
+            )
         self._finalized = True
 
     async def _finish_consumption(self) -> None:

python/packages/core/tests/core/test_types.py

@@ -2,6 +2,7 @@
 
 import asyncio
 import base64
+import contextlib
 import json
 import warnings
 from collections.abc import AsyncIterable, Awaitable, Callable, Sequence
@@ -4276,6 +4277,38 @@ async def test_updates_property_returns_collected(self) -> None:
         assert stream.updates[0].text == "update_0"
         assert stream.updates[1].text == "update_1"
 
+    async def test_pull_context_added_by_factory_applies_to_current_pull(self) -> None:
+        """A pull context registered by the first factory wraps the same iterator pull."""
+        events: list[str] = []
+
+        async def updates() -> AsyncIterable[ChatResponseUpdate]:
+            events.append("pull")
+            yield ChatResponseUpdate(contents=[Content.from_text("update")], role="assistant")
+
+        @contextlib.contextmanager
+        def added_context() -> Any:
+            events.append("added_enter")
+            try:
+                yield
+            finally:
+                events.append("added_exit")
+
+        @contextlib.contextmanager
+        def initial_context() -> Any:
+            events.append("initial_enter")
+            stream.with_pull_context_manager(added_context)
+            try:
+                yield
+            finally:
+                events.append("initial_exit")
+
+        stream = ResponseStream(updates(), finalizer=_combine_updates).with_pull_context_manager(initial_context)
+
+        await anext(stream)
+        await stream.close()
+
+        assert events == ["initial_enter", "added_enter", "pull", "added_exit", "initial_exit"]
+
     async def test_auto_finalize_on_iteration_completion(self) -> None:
         """Stream auto-finalizes when async iteration completes."""
         stream = ResponseStream(_generate_updates(2), finalizer=_combine_updates)
@@ -4447,6 +4480,31 @@ async def async_hook(update: ChatResponseUpdate) -> ChatResponseUpdate:
 
         assert collected == ["async_update_0", "async_update_1"]
 
+    async def test_transform_added_while_async_transform_waits_applies_to_current_update(self) -> None:
+        """A transform registered during an awaited transform applies before the current update is released."""
+        transform_started = asyncio.Event()
+        release_transform = asyncio.Event()
+
+        async def waiting_transform(update: ChatResponseUpdate) -> ChatResponseUpdate:
+            transform_started.set()
+            await release_transform.wait()
+            return update
+
+        def uppercase_transform(update: ChatResponseUpdate) -> ChatResponseUpdate:
+            return ChatResponseUpdate(
+                contents=[Content.from_text((update.text or "").upper())],
+                role=cast(Any, update.role),
+            )
+
+        stream = ResponseStream(_generate_updates(1), finalizer=_combine_updates).with_transform_hook(waiting_transform)
+        pending_update = asyncio.create_task(anext(stream))
+
+        await transform_started.wait()
+        stream.with_transform_hook(uppercase_transform)
+        release_transform.set()
+
+        assert (await pending_update).text == "UPDATE_0"
+
 
 class TestResponseStreamCleanupHooks:
     """Tests for cleanup hooks (after stream consumption, before finalizer)."""