diff --git a/python/CHANGELOG.md b/python/CHANGELOG.md index 33c1b85262a..f2525a94073 100644 --- a/python/CHANGELOG.md +++ b/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)) diff --git a/python/packages/core/agent_framework/_agents.py b/python/packages/core/agent_framework/_agents.py index db264535aa1..c37156a779b 100644 --- a/python/packages/core/agent_framework/_agents.py +++ b/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( diff --git a/python/packages/core/agent_framework/_types.py b/python/packages/core/agent_framework/_types.py index e148e1d2e4e..ba1703c0caf 100644 --- a/python/packages/core/agent_framework/_types.py +++ b/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,15 +3909,19 @@ 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 @@ -3897,25 +3929,28 @@ 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: diff --git a/python/packages/core/tests/core/test_types.py b/python/packages/core/tests/core/test_types.py index e607e1b9e34..9360e1a7609 100644 --- a/python/packages/core/tests/core/test_types.py +++ b/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)."""