From a198c583a6e0c6591ecab782af89ae5a6d80225f Mon Sep 17 00:00:00 2001 From: Copilot Date: Wed, 18 Mar 2026 14:23:22 +0000 Subject: [PATCH 1/8] Fix A2AAgent context_providers not triggered (#4754) A2AAgent.run() overrode BaseAgent.run() without calling before_run/after_run on context_providers. Added _run_before_providers() to invoke before_run on each provider before the A2A call, and _run_after_providers() after the response completes, in both streaming and non-streaming paths. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .../a2a/agent_framework_a2a/_agent.py | 113 +++++++++++++++-- python/packages/a2a/tests/test_a2a_agent.py | 120 ++++++++++++++++++ 2 files changed, 219 insertions(+), 14 deletions(-) diff --git a/python/packages/a2a/agent_framework_a2a/_agent.py b/python/packages/a2a/agent_framework_a2a/_agent.py index e11aa668da..30746d3151 100644 --- a/python/packages/a2a/agent_framework_a2a/_agent.py +++ b/python/packages/a2a/agent_framework_a2a/_agent.py @@ -39,6 +39,7 @@ ContinuationToken, Message, ResponseStream, + SessionContext, normalize_messages, prepend_agent_framework_to_user_agent, ) @@ -284,22 +285,106 @@ def run( # pyright: ignore[reportIncompatibleMethodOverride] When stream=True: A ResponseStream of AgentResponseUpdate items. """ del function_invocation_kwargs, client_kwargs, kwargs - if continuation_token is not None: - a2a_stream: AsyncIterable[A2AStreamItem] = self.client.resubscribe( - TaskIdParams(id=continuation_token["task_id"]) + normalized_messages = normalize_messages(messages) if continuation_token is None else None + + if not stream: + + async def _run_non_streaming() -> AgentResponse[Any]: + active_session, session_context = await self._run_before_providers( + session=session, input_messages=normalized_messages, + ) + if continuation_token is not None: + a2a_stream: AsyncIterable[A2AStreamItem] = self.client.resubscribe( + TaskIdParams(id=continuation_token["task_id"]) + ) + else: + a2a_message = self._prepare_message_for_a2a(normalized_messages[-1]) # type: ignore[index] + a2a_stream = self.client.send_message(a2a_message) + + response_stream = ResponseStream( + self._map_a2a_stream(a2a_stream, background=background), + finalizer=AgentResponse.from_updates, + ) + result = await response_stream.get_final_response() + session_context._response = result + await self._run_after_providers(session=active_session, context=session_context) + return result + + return _run_non_streaming() + + # Streaming path + active_session_holder: dict[str, AgentSession | None] = {"session": None} + context_holder: dict[str, SessionContext | None] = {"ctx": None} + + async def _post_hook(response: AgentResponse) -> None: + session_context = context_holder["ctx"] + if session_context is None: + return + session_context._response = response + await self._run_after_providers(session=active_session_holder["session"], context=session_context) + + async def _get_stream() -> ResponseStream[AgentResponseUpdate, AgentResponse[Any]]: + active_session, session_context = await self._run_before_providers( + session=session, input_messages=normalized_messages, ) - else: - normalized_messages = normalize_messages(messages) - a2a_message = self._prepare_message_for_a2a(normalized_messages[-1]) - a2a_stream = self.client.send_message(a2a_message) - - response = ResponseStream( - self._map_a2a_stream(a2a_stream, background=background), - finalizer=AgentResponse.from_updates, + active_session_holder["session"] = active_session + context_holder["ctx"] = session_context + + if continuation_token is not None: + a2a_stream: AsyncIterable[A2AStreamItem] = self.client.resubscribe( + TaskIdParams(id=continuation_token["task_id"]) + ) + else: + a2a_message = self._prepare_message_for_a2a(normalized_messages[-1]) # type: ignore[index] + a2a_stream = self.client.send_message(a2a_message) + + return ResponseStream( + self._map_a2a_stream(a2a_stream, background=background), + finalizer=AgentResponse.from_updates, + ) + + return ( + ResponseStream + .from_awaitable(_get_stream()) + .with_result_hook(_post_hook) ) - if stream: - return response - return response.get_final_response() + + async def _run_before_providers( + self, + *, + session: AgentSession | None, + input_messages: list[Message] | None, + ) -> tuple[AgentSession | None, SessionContext]: + """Run before_run on all context providers and return the active session and context. + + Keyword Args: + session: The conversation session (None for stateless invocation). + input_messages: Messages to process. + + Returns: + A tuple of (active_session, session_context). + """ + active_session = session + if active_session is None and self.context_providers: + active_session = AgentSession() + + session_context = SessionContext( + session_id=active_session.session_id if active_session else None, + service_session_id=active_session.service_session_id if active_session else None, + input_messages=input_messages or [], + ) + + for provider in self.context_providers: + if active_session is None: + raise RuntimeError("Provider session must be available when context providers are configured.") + await provider.before_run( + agent=self, + session=active_session, + context=session_context, + state=active_session.state.setdefault(provider.source_id, {}), + ) + + return active_session, session_context async def _map_a2a_stream( self, diff --git a/python/packages/a2a/tests/test_a2a_agent.py b/python/packages/a2a/tests/test_a2a_agent.py index a426c27a7f..28da87a2d5 100644 --- a/python/packages/a2a/tests/test_a2a_agent.py +++ b/python/packages/a2a/tests/test_a2a_agent.py @@ -23,6 +23,8 @@ from agent_framework import ( AgentResponse, AgentResponseUpdate, + AgentSession, + BaseContextProvider, Content, Message, ) @@ -850,4 +852,122 @@ async def test_poll_task_completed(a2a_agent: A2AAgent, mock_a2a_client: MockA2A assert response.messages[0].text == "Poll result" +# region context_providers + + +class TrackingContextProvider(BaseContextProvider): + """Context provider that tracks before_run/after_run calls.""" + + def __init__(self, source_id: str = "tracking") -> None: + super().__init__(source_id=source_id) + self.before_run_called = False + self.after_run_called = False + self.before_run_session: AgentSession | None = None + self.after_run_session: AgentSession | None = None + self.after_run_response: AgentResponse | None = None + + async def before_run(self, *, agent, session, context, state) -> None: + self.before_run_called = True + self.before_run_session = session + + async def after_run(self, *, agent, session, context, state) -> None: + self.after_run_called = True + self.after_run_session = session + self.after_run_response = context.response + + +async def test_run_invokes_context_providers(mock_a2a_client: MockA2AClient) -> None: + """Test that run() calls before_run and after_run on context providers.""" + provider = TrackingContextProvider() + agent = A2AAgent(name="Test Agent", client=mock_a2a_client, http_client=None, context_providers=[provider]) + + mock_a2a_client.add_message_response("msg-ctx", "Hello!", "agent") + + response = await agent.run("Hi") + + assert provider.before_run_called + assert provider.after_run_called + assert provider.after_run_response is not None + assert isinstance(response, AgentResponse) + assert response.messages[0].text == "Hello!" + + +async def test_run_invokes_context_providers_with_session(mock_a2a_client: MockA2AClient) -> None: + """Test that context providers receive the provided session.""" + provider = TrackingContextProvider() + agent = A2AAgent(name="Test Agent", client=mock_a2a_client, http_client=None, context_providers=[provider]) + session = AgentSession(session_id="test-session") + + mock_a2a_client.add_message_response("msg-sess", "With session", "agent") + + await agent.run("Hi", session=session) + + assert provider.before_run_session is session + assert provider.after_run_session is session + + +async def test_run_creates_session_for_providers_when_none(mock_a2a_client: MockA2AClient) -> None: + """Test that a session is auto-created when context_providers are set but no session is passed.""" + provider = TrackingContextProvider() + agent = A2AAgent(name="Test Agent", client=mock_a2a_client, http_client=None, context_providers=[provider]) + + mock_a2a_client.add_message_response("msg-auto", "Auto session", "agent") + + await agent.run("Hi") + + assert provider.before_run_session is not None + assert provider.after_run_session is not None + + +async def test_streaming_invokes_context_providers(mock_a2a_client: MockA2AClient) -> None: + """Test that streaming run() calls before_run and after_run on context providers.""" + provider = TrackingContextProvider() + agent = A2AAgent(name="Test Agent", client=mock_a2a_client, http_client=None, context_providers=[provider]) + + mock_a2a_client.add_message_response("msg-stream-ctx", "Streamed!", "agent") + + response = await agent.run("Hi", stream=True).get_final_response() + + assert provider.before_run_called + assert provider.after_run_called + assert provider.after_run_response is not None + assert response.messages[0].text == "Streamed!" + + +async def test_run_without_providers_still_works(mock_a2a_client: MockA2AClient) -> None: + """Test that run() without context_providers still works correctly.""" + agent = A2AAgent(name="Test Agent", client=mock_a2a_client, http_client=None) + + mock_a2a_client.add_message_response("msg-no-ctx", "No providers", "agent") + + response = await agent.run("Hi") + + assert isinstance(response, AgentResponse) + assert response.messages[0].text == "No providers" + + +async def test_multiple_providers_invoked_in_order(mock_a2a_client: MockA2AClient) -> None: + """Test that multiple context providers are called in forward/reverse order.""" + call_order: list[str] = [] + + class OrderTrackingProvider(BaseContextProvider): + async def before_run(self, *, agent, session, context, state) -> None: + call_order.append(f"before:{self.source_id}") + + async def after_run(self, *, agent, session, context, state) -> None: + call_order.append(f"after:{self.source_id}") + + provider_a = OrderTrackingProvider(source_id="a") + provider_b = OrderTrackingProvider(source_id="b") + agent = A2AAgent( + name="Test Agent", client=mock_a2a_client, http_client=None, context_providers=[provider_a, provider_b] + ) + + mock_a2a_client.add_message_response("msg-order", "Ordered", "agent") + + await agent.run("Hi") + + assert call_order == ["before:a", "before:b", "after:b", "after:a"] + + # endregion From fb909b5cbc62cc2b5fc813071b6fcb84e103a248 Mon Sep 17 00:00:00 2001 From: Copilot Date: Wed, 18 Mar 2026 14:23:38 +0000 Subject: [PATCH 2/8] Apply pre-commit auto-fixes --- REPRODUCTION_REPORT.md | 44 +++++++++++++++++++ .../a2a/agent_framework_a2a/_agent.py | 12 +++-- 2 files changed, 49 insertions(+), 7 deletions(-) create mode 100644 REPRODUCTION_REPORT.md diff --git a/REPRODUCTION_REPORT.md b/REPRODUCTION_REPORT.md new file mode 100644 index 0000000000..d259938b5e --- /dev/null +++ b/REPRODUCTION_REPORT.md @@ -0,0 +1,44 @@ +# Reproduction Report: Issue #4754 + +**Issue**: [Python: [Bug]: A2AAgent inconsistent with BaseAgent: context_providers not triggered](https://github.com/microsoft/agent-framework/issues/4754) +**Repository**: microsoft/agent-framework +**Investigated**: 2026-03-18 14:19:47 UTC +**Worktree Path**: `/repos/agent-framework/.worktrees/agent/fix-4754-2` +**Gate Result**: ✅ PASSED + +--- + +## Reproduction Verdict + +| Field | Value | +|-------|-------| +| Status | `reproduced` | +| Failure observed | `yes` | +| Evidence strength | `high` | +| Confidence | `high` | +| Failing test | `python/packages/a2a/tests/test_a2a_context_providers_bug.py` | +| Gate passed | ✅ PASSED | + +## Finding + +A2AAgent.run() (lines 250-302 in _agent.py) completely overrides BaseAgent.run() without ever calling _prepare_session_and_messages() or _run_after_providers(), which are the methods that invoke before_run/after_run on context_providers. A reproduction test confirms context_providers are silently ignored — the TrackingContextProvider's before_run_called remains False after agent.run() completes. + +## Files Examined + +python/packages/a2a/agent_framework_a2a/_agent.py, python/packages/core/agent_framework/_agents.py, python/packages/core/agent_framework/_sessions.py, python/packages/a2a/tests/test_a2a_agent.py + +## Tests Run + +python/packages/a2a/tests/test_a2a_context_providers_bug.py::test_a2a_agent_context_providers_triggered + +## Full Analysis + +Now let me look at how BaseAgent's `run` method handles context providers, specifically the `_prepare_session_and_messages` and the `_run_after_providers` calls:Now I have clear evidence. The A2AAgent's `run()` method (line 250-302) completely overrides BaseAgent's `run()` without calling `_prepare_session_and_messages()` (which runs `before_run` on providers) or `_run_after_providers()` (which runs `after_run`). Let me write a reproduction test.Now I have everything I need. Let me write and run the reproduction test:Bug reproduced. Now let me write the report:Bug confirmed. A2AAgent's `run()` overrides BaseAgent without calling any context provider hooks, so `before_run`/`after_run` on providers are never invoked. + +--- + +## Next Steps + +- Reproduction confirmed — proceed to investigation and fix phase +- Review the failing test or evidence above to understand the root cause +- The issue branch is checked out in the worktree: `/repos/agent-framework/.worktrees/agent/fix-4754-2` \ No newline at end of file diff --git a/python/packages/a2a/agent_framework_a2a/_agent.py b/python/packages/a2a/agent_framework_a2a/_agent.py index 30746d3151..0fa30a9f9c 100644 --- a/python/packages/a2a/agent_framework_a2a/_agent.py +++ b/python/packages/a2a/agent_framework_a2a/_agent.py @@ -291,7 +291,8 @@ def run( # pyright: ignore[reportIncompatibleMethodOverride] async def _run_non_streaming() -> AgentResponse[Any]: active_session, session_context = await self._run_before_providers( - session=session, input_messages=normalized_messages, + session=session, + input_messages=normalized_messages, ) if continuation_token is not None: a2a_stream: AsyncIterable[A2AStreamItem] = self.client.resubscribe( @@ -325,7 +326,8 @@ async def _post_hook(response: AgentResponse) -> None: async def _get_stream() -> ResponseStream[AgentResponseUpdate, AgentResponse[Any]]: active_session, session_context = await self._run_before_providers( - session=session, input_messages=normalized_messages, + session=session, + input_messages=normalized_messages, ) active_session_holder["session"] = active_session context_holder["ctx"] = session_context @@ -343,11 +345,7 @@ async def _get_stream() -> ResponseStream[AgentResponseUpdate, AgentResponse[Any finalizer=AgentResponse.from_updates, ) - return ( - ResponseStream - .from_awaitable(_get_stream()) - .with_result_hook(_post_hook) - ) + return ResponseStream.from_awaitable(_get_stream()).with_result_hook(_post_hook) async def _run_before_providers( self, From 37c88a0f43af09a405079fc69e3f9acfcf697e87 Mon Sep 17 00:00:00 2001 From: Copilot Date: Wed, 18 Mar 2026 14:26:46 +0000 Subject: [PATCH 3/8] Fix pyright errors in A2AAgent context_providers implementation (#4754) Add type suppression comments for: - reportPrivateUsage on SessionContext._response access (cross-package) - reportUnknownMemberType/reportUnknownVariableType on ResponseStream chained calls, matching patterns used in core package Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- python/packages/a2a/agent_framework_a2a/_agent.py | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/python/packages/a2a/agent_framework_a2a/_agent.py b/python/packages/a2a/agent_framework_a2a/_agent.py index 0fa30a9f9c..1a14546e6c 100644 --- a/python/packages/a2a/agent_framework_a2a/_agent.py +++ b/python/packages/a2a/agent_framework_a2a/_agent.py @@ -307,7 +307,7 @@ async def _run_non_streaming() -> AgentResponse[Any]: finalizer=AgentResponse.from_updates, ) result = await response_stream.get_final_response() - session_context._response = result + session_context._response = result # type: ignore[assignment] # pyright: ignore[reportPrivateUsage] await self._run_after_providers(session=active_session, context=session_context) return result @@ -321,7 +321,7 @@ async def _post_hook(response: AgentResponse) -> None: session_context = context_holder["ctx"] if session_context is None: return - session_context._response = response + session_context._response = response # type: ignore[assignment] # pyright: ignore[reportPrivateUsage] await self._run_after_providers(session=active_session_holder["session"], context=session_context) async def _get_stream() -> ResponseStream[AgentResponseUpdate, AgentResponse[Any]]: @@ -345,7 +345,9 @@ async def _get_stream() -> ResponseStream[AgentResponseUpdate, AgentResponse[Any finalizer=AgentResponse.from_updates, ) - return ResponseStream.from_awaitable(_get_stream()).with_result_hook(_post_hook) + return ( + ResponseStream.from_awaitable(_get_stream()).with_result_hook(_post_hook) # type: ignore[reportUnknownMemberType] # type: ignore[reportUnknownMemberType, reportUnknownVariableType] + ) async def _run_before_providers( self, From f314a5ab0a340f12b50c28baae642b7e4515ee80 Mon Sep 17 00:00:00 2001 From: Copilot Date: Wed, 18 Mar 2026 14:29:00 +0000 Subject: [PATCH 4/8] Review fixes: remove reproduction report, add BaseHistoryProvider skip logic - Remove REPRODUCTION_REPORT.md (debugging artifact, not for commit) - Add BaseHistoryProvider load_messages=False skip in _run_before_providers to match the pattern used in _prepare_session_and_messages and WorkflowAgent Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- REPRODUCTION_REPORT.md | 44 ------------------- .../a2a/agent_framework_a2a/_agent.py | 3 ++ 2 files changed, 3 insertions(+), 44 deletions(-) delete mode 100644 REPRODUCTION_REPORT.md diff --git a/REPRODUCTION_REPORT.md b/REPRODUCTION_REPORT.md deleted file mode 100644 index d259938b5e..0000000000 --- a/REPRODUCTION_REPORT.md +++ /dev/null @@ -1,44 +0,0 @@ -# Reproduction Report: Issue #4754 - -**Issue**: [Python: [Bug]: A2AAgent inconsistent with BaseAgent: context_providers not triggered](https://github.com/microsoft/agent-framework/issues/4754) -**Repository**: microsoft/agent-framework -**Investigated**: 2026-03-18 14:19:47 UTC -**Worktree Path**: `/repos/agent-framework/.worktrees/agent/fix-4754-2` -**Gate Result**: ✅ PASSED - ---- - -## Reproduction Verdict - -| Field | Value | -|-------|-------| -| Status | `reproduced` | -| Failure observed | `yes` | -| Evidence strength | `high` | -| Confidence | `high` | -| Failing test | `python/packages/a2a/tests/test_a2a_context_providers_bug.py` | -| Gate passed | ✅ PASSED | - -## Finding - -A2AAgent.run() (lines 250-302 in _agent.py) completely overrides BaseAgent.run() without ever calling _prepare_session_and_messages() or _run_after_providers(), which are the methods that invoke before_run/after_run on context_providers. A reproduction test confirms context_providers are silently ignored — the TrackingContextProvider's before_run_called remains False after agent.run() completes. - -## Files Examined - -python/packages/a2a/agent_framework_a2a/_agent.py, python/packages/core/agent_framework/_agents.py, python/packages/core/agent_framework/_sessions.py, python/packages/a2a/tests/test_a2a_agent.py - -## Tests Run - -python/packages/a2a/tests/test_a2a_context_providers_bug.py::test_a2a_agent_context_providers_triggered - -## Full Analysis - -Now let me look at how BaseAgent's `run` method handles context providers, specifically the `_prepare_session_and_messages` and the `_run_after_providers` calls:Now I have clear evidence. The A2AAgent's `run()` method (line 250-302) completely overrides BaseAgent's `run()` without calling `_prepare_session_and_messages()` (which runs `before_run` on providers) or `_run_after_providers()` (which runs `after_run`). Let me write a reproduction test.Now I have everything I need. Let me write and run the reproduction test:Bug reproduced. Now let me write the report:Bug confirmed. A2AAgent's `run()` overrides BaseAgent without calling any context provider hooks, so `before_run`/`after_run` on providers are never invoked. - ---- - -## Next Steps - -- Reproduction confirmed — proceed to investigation and fix phase -- Review the failing test or evidence above to understand the root cause -- The issue branch is checked out in the worktree: `/repos/agent-framework/.worktrees/agent/fix-4754-2` \ No newline at end of file diff --git a/python/packages/a2a/agent_framework_a2a/_agent.py b/python/packages/a2a/agent_framework_a2a/_agent.py index 1a14546e6c..6e21bbddc0 100644 --- a/python/packages/a2a/agent_framework_a2a/_agent.py +++ b/python/packages/a2a/agent_framework_a2a/_agent.py @@ -35,6 +35,7 @@ AgentResponseUpdate, AgentSession, BaseAgent, + BaseHistoryProvider, Content, ContinuationToken, Message, @@ -375,6 +376,8 @@ async def _run_before_providers( ) for provider in self.context_providers: + if isinstance(provider, BaseHistoryProvider) and not provider.load_messages: + continue if active_session is None: raise RuntimeError("Provider session must be available when context providers are configured.") await provider.before_run( From 430069a38c571482094f83c1d4a59d6ea26c6a24 Mon Sep 17 00:00:00 2001 From: Copilot Date: Wed, 18 Mar 2026 14:42:47 +0000 Subject: [PATCH 5/8] Fix PR review feedback for #4754: type ignore syntax and history provider tests - Fix line 350: replace duplicate # type: ignore with single # pyright: ignore - Add tests for BaseHistoryProvider with load_messages=False skip logic: - before_run is NOT called when load_messages=False - before_run IS called when load_messages=True (default) - after_run is always called regardless of load_messages - Streaming variant of load_messages=False - Mixed providers: regular provider still invoked alongside skipped history provider Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .../a2a/agent_framework_a2a/_agent.py | 2 +- python/packages/a2a/tests/test_a2a_agent.py | 82 +++++++++++++++++++ 2 files changed, 83 insertions(+), 1 deletion(-) diff --git a/python/packages/a2a/agent_framework_a2a/_agent.py b/python/packages/a2a/agent_framework_a2a/_agent.py index 6e21bbddc0..e20f338e8f 100644 --- a/python/packages/a2a/agent_framework_a2a/_agent.py +++ b/python/packages/a2a/agent_framework_a2a/_agent.py @@ -347,7 +347,7 @@ async def _get_stream() -> ResponseStream[AgentResponseUpdate, AgentResponse[Any ) return ( - ResponseStream.from_awaitable(_get_stream()).with_result_hook(_post_hook) # type: ignore[reportUnknownMemberType] # type: ignore[reportUnknownMemberType, reportUnknownVariableType] + ResponseStream.from_awaitable(_get_stream()).with_result_hook(_post_hook) # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType] ) async def _run_before_providers( diff --git a/python/packages/a2a/tests/test_a2a_agent.py b/python/packages/a2a/tests/test_a2a_agent.py index 28da87a2d5..e711dbbd1e 100644 --- a/python/packages/a2a/tests/test_a2a_agent.py +++ b/python/packages/a2a/tests/test_a2a_agent.py @@ -25,6 +25,7 @@ AgentResponseUpdate, AgentSession, BaseContextProvider, + BaseHistoryProvider, Content, Message, ) @@ -970,4 +971,85 @@ async def after_run(self, *, agent, session, context, state) -> None: assert call_order == ["before:a", "before:b", "after:b", "after:a"] +class TrackingHistoryProvider(BaseHistoryProvider): + """History provider that tracks before_run/after_run calls.""" + + def __init__(self, source_id: str = "history", *, load_messages: bool = True) -> None: + super().__init__(source_id=source_id, load_messages=load_messages) + self.before_run_called = False + self.after_run_called = False + + async def before_run(self, *, agent, session, context, state) -> None: + self.before_run_called = True + + async def after_run(self, *, agent, session, context, state) -> None: + self.after_run_called = True + + async def get_messages(self, session_id, **kwargs) -> list[Message]: + return [] + + async def save_messages(self, session_id, messages, **kwargs) -> None: + pass + + +async def test_history_provider_load_messages_false_skips_before_run(mock_a2a_client: MockA2AClient) -> None: + """Test that BaseHistoryProvider with load_messages=False has before_run skipped.""" + provider = TrackingHistoryProvider(load_messages=False) + agent = A2AAgent(name="Test Agent", client=mock_a2a_client, http_client=None, context_providers=[provider]) + + mock_a2a_client.add_message_response("msg-hist", "Hello!", "agent") + + await agent.run("Hi") + + assert not provider.before_run_called + assert provider.after_run_called + + +async def test_history_provider_load_messages_true_calls_before_run(mock_a2a_client: MockA2AClient) -> None: + """Test that BaseHistoryProvider with load_messages=True (default) has before_run called.""" + provider = TrackingHistoryProvider(load_messages=True) + agent = A2AAgent(name="Test Agent", client=mock_a2a_client, http_client=None, context_providers=[provider]) + + mock_a2a_client.add_message_response("msg-hist-true", "Hello!", "agent") + + await agent.run("Hi") + + assert provider.before_run_called + assert provider.after_run_called + + +async def test_history_provider_load_messages_false_streaming(mock_a2a_client: MockA2AClient) -> None: + """Test that streaming skips before_run for BaseHistoryProvider with load_messages=False.""" + provider = TrackingHistoryProvider(load_messages=False) + agent = A2AAgent(name="Test Agent", client=mock_a2a_client, http_client=None, context_providers=[provider]) + + mock_a2a_client.add_message_response("msg-hist-stream", "Streamed!", "agent") + + await agent.run("Hi", stream=True).get_final_response() + + assert not provider.before_run_called + assert provider.after_run_called + + +async def test_mixed_providers_with_history_load_messages_false(mock_a2a_client: MockA2AClient) -> None: + """Test that a regular provider's before_run is called while history provider's is skipped.""" + context_provider = TrackingContextProvider(source_id="ctx") + history_provider = TrackingHistoryProvider(source_id="hist", load_messages=False) + agent = A2AAgent( + name="Test Agent", + client=mock_a2a_client, + http_client=None, + context_providers=[context_provider, history_provider], + ) + + mock_a2a_client.add_message_response("msg-mixed", "Mixed!", "agent") + + await agent.run("Hi") + + assert context_provider.before_run_called + assert not history_provider.before_run_called + assert context_provider.after_run_called + assert history_provider.after_run_called + + # endregion From a8f779676b21e41ba36d9051abb042d66299b933 Mon Sep 17 00:00:00 2001 From: Copilot Date: Wed, 18 Mar 2026 14:45:08 +0000 Subject: [PATCH 6/8] Address PR review feedback for #4754: fast-path and stronger test - Add fast-path in run() to skip provider hooks when no context_providers are configured, avoiding unnecessary SessionContext allocation - Add FailingHistoryProvider test: before_run raises AssertionError if called, proving load_messages=False skip logic works correctly - Previous commit already fixed pyright ignore syntax on line 350 Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .../a2a/agent_framework_a2a/_agent.py | 29 ++++++++++------- python/packages/a2a/tests/test_a2a_agent.py | 32 +++++++++++++++++++ 2 files changed, 49 insertions(+), 12 deletions(-) diff --git a/python/packages/a2a/agent_framework_a2a/_agent.py b/python/packages/a2a/agent_framework_a2a/_agent.py index e20f338e8f..58763467ef 100644 --- a/python/packages/a2a/agent_framework_a2a/_agent.py +++ b/python/packages/a2a/agent_framework_a2a/_agent.py @@ -291,10 +291,11 @@ def run( # pyright: ignore[reportIncompatibleMethodOverride] if not stream: async def _run_non_streaming() -> AgentResponse[Any]: - active_session, session_context = await self._run_before_providers( - session=session, - input_messages=normalized_messages, - ) + if self.context_providers: + active_session, session_context = await self._run_before_providers( + session=session, + input_messages=normalized_messages, + ) if continuation_token is not None: a2a_stream: AsyncIterable[A2AStreamItem] = self.client.resubscribe( TaskIdParams(id=continuation_token["task_id"]) @@ -308,8 +309,9 @@ async def _run_non_streaming() -> AgentResponse[Any]: finalizer=AgentResponse.from_updates, ) result = await response_stream.get_final_response() - session_context._response = result # type: ignore[assignment] # pyright: ignore[reportPrivateUsage] - await self._run_after_providers(session=active_session, context=session_context) + if self.context_providers: + session_context._response = result # type: ignore[assignment] # pyright: ignore[reportPrivateUsage] + await self._run_after_providers(session=active_session, context=session_context) return result return _run_non_streaming() @@ -319,6 +321,8 @@ async def _run_non_streaming() -> AgentResponse[Any]: context_holder: dict[str, SessionContext | None] = {"ctx": None} async def _post_hook(response: AgentResponse) -> None: + if not self.context_providers: + return session_context = context_holder["ctx"] if session_context is None: return @@ -326,12 +330,13 @@ async def _post_hook(response: AgentResponse) -> None: await self._run_after_providers(session=active_session_holder["session"], context=session_context) async def _get_stream() -> ResponseStream[AgentResponseUpdate, AgentResponse[Any]]: - active_session, session_context = await self._run_before_providers( - session=session, - input_messages=normalized_messages, - ) - active_session_holder["session"] = active_session - context_holder["ctx"] = session_context + if self.context_providers: + active_session, session_context = await self._run_before_providers( + session=session, + input_messages=normalized_messages, + ) + active_session_holder["session"] = active_session + context_holder["ctx"] = session_context if continuation_token is not None: a2a_stream: AsyncIterable[A2AStreamItem] = self.client.resubscribe( diff --git a/python/packages/a2a/tests/test_a2a_agent.py b/python/packages/a2a/tests/test_a2a_agent.py index e711dbbd1e..797d01a555 100644 --- a/python/packages/a2a/tests/test_a2a_agent.py +++ b/python/packages/a2a/tests/test_a2a_agent.py @@ -1005,6 +1005,38 @@ async def test_history_provider_load_messages_false_skips_before_run(mock_a2a_cl assert provider.after_run_called +async def test_history_provider_load_messages_false_raises_if_before_run_called( + mock_a2a_client: MockA2AClient, +) -> None: + """Test with a stub whose before_run raises, proving it is never invoked.""" + + class FailingHistoryProvider(BaseHistoryProvider): + def __init__(self) -> None: + super().__init__(source_id="fail-hist", load_messages=False) + self.after_run_called = False + + async def before_run(self, *, agent, session, context, state) -> None: + raise AssertionError("before_run should not be called when load_messages=False") + + async def after_run(self, *, agent, session, context, state) -> None: + self.after_run_called = True + + async def get_messages(self, session_id, **kwargs) -> list[Message]: + return [] + + async def save_messages(self, session_id, messages, **kwargs) -> None: + pass + + provider = FailingHistoryProvider() + agent = A2AAgent(name="Test Agent", client=mock_a2a_client, http_client=None, context_providers=[provider]) + + mock_a2a_client.add_message_response("msg-fail", "OK", "agent") + + # Should not raise — before_run is skipped + await agent.run("Hi") + assert provider.after_run_called + + async def test_history_provider_load_messages_true_calls_before_run(mock_a2a_client: MockA2AClient) -> None: """Test that BaseHistoryProvider with load_messages=True (default) has before_run called.""" provider = TrackingHistoryProvider(load_messages=True) From eb986bdc474657d316cd3def188a6edf07c11c28 Mon Sep 17 00:00:00 2001 From: Copilot Date: Wed, 18 Mar 2026 14:46:47 +0000 Subject: [PATCH 7/8] fix: resolve pyright 'possibly unbound' errors in A2AAgent non-streaming path (#4754) Initialize active_session and session_context before the conditional context_providers block so pyright can verify they are always bound. Add explicit None check for session_context in the after-providers guard. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- python/packages/a2a/agent_framework_a2a/_agent.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/python/packages/a2a/agent_framework_a2a/_agent.py b/python/packages/a2a/agent_framework_a2a/_agent.py index 58763467ef..da4e23dd09 100644 --- a/python/packages/a2a/agent_framework_a2a/_agent.py +++ b/python/packages/a2a/agent_framework_a2a/_agent.py @@ -291,6 +291,8 @@ def run( # pyright: ignore[reportIncompatibleMethodOverride] if not stream: async def _run_non_streaming() -> AgentResponse[Any]: + active_session: AgentSession | None = None + session_context: SessionContext | None = None if self.context_providers: active_session, session_context = await self._run_before_providers( session=session, @@ -309,7 +311,7 @@ async def _run_non_streaming() -> AgentResponse[Any]: finalizer=AgentResponse.from_updates, ) result = await response_stream.get_final_response() - if self.context_providers: + if self.context_providers and session_context is not None: session_context._response = result # type: ignore[assignment] # pyright: ignore[reportPrivateUsage] await self._run_after_providers(session=active_session, context=session_context) return result From 3165912e8e9e9c36e37df11cb9f5c21b1d85b28f Mon Sep 17 00:00:00 2001 From: Copilot Date: Wed, 18 Mar 2026 14:54:27 +0000 Subject: [PATCH 8/8] test: add continuation_token + context_providers coverage (#4754) Add two tests exercising the non-streaming continuation_token path: - test_resume_via_continuation_token_with_context_providers: validates that before_run/after_run are invoked when resuming with providers - test_resume_via_continuation_token_no_context_providers: validates no crash when resuming without providers (exercises the guard at line 314) Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- python/packages/a2a/tests/test_a2a_agent.py | 45 +++++++++++++++++++++ 1 file changed, 45 insertions(+) diff --git a/python/packages/a2a/tests/test_a2a_agent.py b/python/packages/a2a/tests/test_a2a_agent.py index 797d01a555..d2fcfba2c8 100644 --- a/python/packages/a2a/tests/test_a2a_agent.py +++ b/python/packages/a2a/tests/test_a2a_agent.py @@ -1084,4 +1084,49 @@ async def test_mixed_providers_with_history_load_messages_false(mock_a2a_client: assert history_provider.after_run_called +async def test_resume_via_continuation_token_with_context_providers(mock_a2a_client: MockA2AClient) -> None: + """Test that non-streaming run() with continuation_token correctly invokes context providers.""" + provider = TrackingContextProvider() + agent = A2AAgent(name="Test Agent", client=mock_a2a_client, http_client=None, context_providers=[provider]) + + status = TaskStatus(state=TaskState.completed, message=None) + artifact = Artifact( + artifact_id="art-ctx-resume", + name="result", + parts=[Part(root=TextPart(text="Resumed with providers"))], + ) + task = Task(id="task-ctx-resume", context_id="ctx-cr", status=status, artifacts=[artifact]) + mock_a2a_client.resubscribe_responses.append((task, None)) + + token = A2AContinuationToken(task_id="task-ctx-resume", context_id="ctx-cr") + response = await agent.run(continuation_token=token) + + assert isinstance(response, AgentResponse) + assert response.messages[0].text == "Resumed with providers" + assert provider.before_run_called + assert provider.after_run_called + assert provider.after_run_response is not None + + +async def test_resume_via_continuation_token_no_context_providers(mock_a2a_client: MockA2AClient) -> None: + """Test that run() with continuation_token and no context_providers works without crash.""" + agent = A2AAgent(name="Test Agent", client=mock_a2a_client, http_client=None) + + status = TaskStatus(state=TaskState.completed, message=None) + artifact = Artifact( + artifact_id="art-no-ctx", + name="result", + parts=[Part(root=TextPart(text="Resumed no providers"))], + ) + task = Task(id="task-no-ctx", context_id="ctx-nc", status=status, artifacts=[artifact]) + mock_a2a_client.resubscribe_responses.append((task, None)) + + token = A2AContinuationToken(task_id="task-no-ctx", context_id="ctx-nc") + response = await agent.run(continuation_token=token) + + assert isinstance(response, AgentResponse) + assert response.messages[0].text == "Resumed no providers" + assert response.continuation_token is None + + # endregion