agent/agent-framework
2
0
Fork 0
mirror of https://github.com/microsoft/agent-framework.git synced 2026-06-16 21:04:09 +08:00
Code Issues Packages Projects Releases Wiki Activity

Python: Fix A2AAgent to invoke context providers before and after run (#4757)

Browse Source 
* Fix A2AAgent to invoke context providers before and after run

A2AAgent.run() bypassed the context provider lifecycle (before_run/after_run)
that BaseAgent defines as a contract for all agents. This caused A2AAgent to
violate the semantic definition of BaseAgent, resulting in inconsistency with
other agent implementations.

The fix follows the same pattern used by WorkflowAgent:
- Create SessionContext and run before_run on all context providers before
  processing the A2A stream
- Collect response updates and run after_run on all context providers after
  the stream is fully consumed
- Auto-create a session when context providers are configured but no session
  is explicitly passed

Fixes #4754

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

* Apply pre-commit auto-fixes

* Remove reproduction report from repository

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

* Address PR review feedback for #4754

- Validate messages when no continuation_token: raise ValueError if
  normalized_messages is empty, preventing IndexError on messages[-1]
- Import BaseContextProvider/SessionContext from public agent_framework
  package instead of internal agent_framework._sessions module
- Add test for ValueError on run(None) without continuation_token

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

* Improve test coverage for empty-messages guard in A2AAgent.run (#4754)

- Parameterize test to cover both messages=None and messages=[] inputs
- Add test verifying run(None, continuation_token=...) does not raise

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

---------

Co-authored-by: Copilot <copilot@github.com>
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
This commit is contained in:
Eduard van Valkenburg
2026-03-19 11:45:42 +01:00
committed by GitHub
Unverified
parent bf8d9672e1
commit 4b21f38650
2 changed files with 243 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/packages/a2a/agent_framework_a2a/_agent.py
+54 -3
View File
@@ -35,10 +35,12 @@ from agent_framework import (
AgentResponseUpdate,
AgentSession,
BaseAgent,
BaseHistoryProvider,
Content,
ContinuationToken,
Message,
ResponseStream,
SessionContext,
normalize_messages,
prepend_agent_framework_to_user_agent,
)
@@ -284,17 +286,36 @@ class A2AAgent(AgentTelemetryLayer, BaseAgent):
When stream=True: A ResponseStream of AgentResponseUpdate items.
"""
del function_invocation_kwargs, client_kwargs, kwargs
normalized_messages = normalize_messages(messages)
if continuation_token is not None:
a2a_stream: AsyncIterable[A2AStreamItem] = self.client.resubscribe(
TaskIdParams(id=continuation_token["task_id"])
)
else:
normalized_messages = normalize_messages(messages)
if not normalized_messages:
raise ValueError("At least one message is required when starting a new task (no continuation_token).")
a2a_message = self._prepare_message_for_a2a(normalized_messages[-1])
a2a_stream = self.client.send_message(a2a_message)
provider_session = session
if provider_session is None and self.context_providers:
provider_session = AgentSession()
session_context = SessionContext(
session_id=provider_session.session_id if provider_session else None,
service_session_id=provider_session.service_session_id if provider_session else None,
input_messages=normalized_messages or [],
options={},
)
response = ResponseStream(
self._map_a2a_stream(a2a_stream, background=background),
self._map_a2a_stream(
a2a_stream,
background=background,
session=provider_session,
session_context=session_context,
),
finalizer=AgentResponse.from_updates,
)
if stream:
@@ -306,6 +327,8 @@ class A2AAgent(AgentTelemetryLayer, BaseAgent):
a2a_stream: AsyncIterable[A2AStreamItem],
*,
background: bool = False,
session: AgentSession | None = None,
session_context: SessionContext | None = None,
) -> AsyncIterable[AgentResponseUpdate]:
"""Map raw A2A protocol items to AgentResponseUpdates.
@@ -316,24 +339,52 @@ class A2AAgent(AgentTelemetryLayer, BaseAgent):
background: When False, in-progress task updates are silently
consumed (the stream keeps iterating until a terminal state).
When True, they are yielded with a continuation token.
session: The agent session for context providers.
session_context: The session context for context providers.
"""
if session_context is None:
session_context = SessionContext(input_messages=[], options={})
# Run before_run providers (forward order)
for provider in self.context_providers:
if isinstance(provider, BaseHistoryProvider) and not provider.load_messages:
continue
if session is None:
raise RuntimeError("Provider session must be available when context providers are configured.")
await provider.before_run(
agent=self, # type: ignore[arg-type]
session=session,
context=session_context,
state=session.state.setdefault(provider.source_id, {}),
)
all_updates: list[AgentResponseUpdate] = []
async for item in a2a_stream:
if isinstance(item, A2AMessage):
# Process A2A Message
contents = self._parse_contents_from_a2a(item.parts)
yield AgentResponseUpdate(
update = AgentResponseUpdate(
contents=contents,
role="assistant" if item.role == A2ARole.agent else "user",
response_id=str(getattr(item, "message_id", uuid.uuid4())),
raw_representation=item,
)
all_updates.append(update)
yield update
elif isinstance(item, tuple) and len(item) == 2 and isinstance(item[0], Task):
task, _update_event = item
for update in self._updates_from_task(task, background=background):
all_updates.append(update)
yield update
else:
raise NotImplementedError("Only Message and Task responses are supported")
# Set the response on the context for after_run providers
if all_updates:
session_context._response = AgentResponse.from_updates(all_updates) # type: ignore[assignment]
await self._run_after_providers(session=session, context=session_context)
# ------------------------------------------------------------------
# Task helpers
# ------------------------------------------------------------------