From 18775471d7113c90f226dd3067d613209f21ca1f Mon Sep 17 00:00:00 2001 From: Jeremy Schoemaker Date: Fri, 11 Sep 2026 13:28:47 -0500 Subject: [PATCH] =?UTF-8?q?fix(orchestrations):=20give=20the=20GroupChat?= =?UTF-8?q?=20orchestrator=20agent=20the=20workflow=20run=20kwargs=20?= =?UTF-8?q?=F0=9F=94=8C?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Fixes #8304. workflow.run(function_invocation_kwargs=..., client_kwargs=...) is stored in workflow state under WORKFLOW_RUN_KWARGS_KEY and AgentExecutor forwards it to every participant agent. The GroupChat orchestrator agent runs outside AgentExecutor, in AgentBasedGroupChatOrchestrator._invoke_agent, which called self._agent.run() with only messages, session and options. So a host that puts request-scoped values there (a user id for tool ACL, an id for metrics) had them reach the participants and silently not reach the orchestrator, even though the orchestrator is an agent with its own tools and middleware reading AgentContext.function_invocation_kwargs. _invoke_agent now takes the WorkflowContext both of its callers already hold, reads the same state key AgentExecutor reads, and forwards both kwargs. The resolution itself is not reimplemented. AgentExecutor._prepare_agent_run_args and _resolve_executor_kwargs depended on nothing but self.id, so they move to _agent_utils as prepare_agent_run_args(executor_id, ...) and resolve_executor_kwargs(executor_id, ...), and both AgentExecutor methods stay as one-line delegates. The orchestrator is itself an Executor, so passing self.id gives it the same semantics participants get: __global__ kwargs apply, per-executor entries are keyed by the orchestrator's own id, and specific values override global ones. Not changed: the Magentic manager's _complete, named in the issue as the same pattern at a separate call site. Its callers are plan/replan/ create_progress_ledger/prepare_final_answer on MagenticManagerBase, which is not an Executor and holds no WorkflowContext, so reaching the run kwargs there means widening a public extension point rather than reading state that is already in hand. That is a design call for the team, not a drive-by. 2 tests added in packages/orchestrations/tests/test_group_chat.py. The first fails against unpatched sources with AssertionError: assert None == {'user_id': 'user-123'} the orchestrator having been invoked with only {'options': {'response_format': AgentOrchestrationOutput}}. The second pins the no-kwargs case so the resolution keeps returning None rather than an empty dict, which also holds before the change. packages/core/tests and packages/orchestrations/tests: 5381 tests, 0 failures, 0 errors, 142 skipped. poe syntax and poe test-typing clean for both packages; poe pyright reports the same 155 pre-existing core errors before and after, none in the touched files. --- .../_workflows/_agent_executor.py | 37 ++-------- .../_workflows/_agent_utils.py | 70 ++++++++++++++++++ .../_group_chat.py | 21 ++++-- .../orchestrations/tests/test_group_chat.py | 72 +++++++++++++++++++ 4 files changed, 162 insertions(+), 38 deletions(-) diff --git a/python/packages/core/agent_framework/_workflows/_agent_executor.py b/python/packages/core/agent_framework/_workflows/_agent_executor.py index 1c3a7203f49..cd69504c71d 100644 --- a/python/packages/core/agent_framework/_workflows/_agent_executor.py +++ b/python/packages/core/agent_framework/_workflows/_agent_executor.py @@ -14,8 +14,8 @@ from .._agents import SupportsAgentRun from .._sessions import AgentSession from .._types import AgentResponse, AgentResponseUpdate, Message, ResponseStream -from ._agent_utils import resolve_agent_id -from ._const import GLOBAL_KWARGS_KEY, INTERNAL_SOURCE_ID, WORKFLOW_RUN_KWARGS_KEY +from ._agent_utils import prepare_agent_run_args, resolve_agent_id, resolve_executor_kwargs +from ._const import INTERNAL_SOURCE_ID, WORKFLOW_RUN_KWARGS_KEY from ._executor import Executor, handler from ._message_utils import normalize_messages_input from ._request_info_mixin import response_handler @@ -594,12 +594,7 @@ def _prepare_agent_run_args( Returns: A 2-tuple of (function_invocation_kwargs, client_kwargs). """ - fi_resolved = raw_run_kwargs.get("function_invocation_kwargs") - ci_resolved = raw_run_kwargs.get("client_kwargs") - function_invocation_kwargs = self._resolve_executor_kwargs(fi_resolved) - client_kwargs = self._resolve_executor_kwargs(ci_resolved) - - return function_invocation_kwargs, client_kwargs + return prepare_agent_run_args(self.id, raw_run_kwargs) def _resolve_executor_kwargs(self, resolved: dict[str, Any] | None) -> dict[str, Any] | None: """Extract this executor's kwargs from a resolved invocation kwargs dict. @@ -612,28 +607,4 @@ def _resolve_executor_kwargs(self, resolved: dict[str, Any] | None) -> dict[str, Returns: The kwargs for this executor, or ``None`` if not applicable. """ - if not isinstance(resolved, dict): - return None - global_kwargs: Any = resolved.get(GLOBAL_KWARGS_KEY) - executor_kwargs: Any = resolved.get(self.id) - if global_kwargs is None and executor_kwargs is None: - return None - - if global_kwargs is not None and not isinstance(global_kwargs, dict): - logger.warning( - "Executor %s expected a dict for global kwargs, but got %s. Ignoring.", - self.id, - cast(type[Any], type(global_kwargs)), - ) - return None - - if executor_kwargs is not None and not isinstance(executor_kwargs, dict): - logger.warning( - "Executor %s expected a dict for its kwargs, but got %s. Ignoring.", - self.id, - cast(type[Any], type(executor_kwargs)), - ) - return None - - # Specific values override global values for the same function argument. - return {**(global_kwargs or {}), **(executor_kwargs or {})} + return resolve_executor_kwargs(self.id, resolved) diff --git a/python/packages/core/agent_framework/_workflows/_agent_utils.py b/python/packages/core/agent_framework/_workflows/_agent_utils.py index f70d524ceeb..08034123328 100644 --- a/python/packages/core/agent_framework/_workflows/_agent_utils.py +++ b/python/packages/core/agent_framework/_workflows/_agent_utils.py @@ -1,6 +1,12 @@ # Copyright (c) Microsoft. All rights reserved. +import logging +from typing import Any, cast + from .._agents import SupportsAgentRun +from ._const import GLOBAL_KWARGS_KEY + +logger = logging.getLogger(__name__) def resolve_agent_id(agent: SupportsAgentRun) -> str: @@ -15,3 +21,67 @@ def resolve_agent_id(agent: SupportsAgentRun) -> str: The resolved unique identifier for the agent. """ return agent.name if agent.name else agent.id + + +def resolve_executor_kwargs(executor_id: str, resolved: dict[str, Any] | None) -> dict[str, Any] | None: + """Extract one executor's kwargs from a resolved invocation kwargs dict. + + Args: + executor_id: The id of the executor whose kwargs are wanted. + resolved: The resolved dict produced by ``Workflow._resolve_invocation_kwargs``, + containing either a ``__global__`` key (global kwargs) or executor-ID keys + (per-executor kwargs). May also be ``None``. + + Returns: + The kwargs for that executor, or ``None`` if not applicable. + """ + if not isinstance(resolved, dict): + return None + global_kwargs: Any = resolved.get(GLOBAL_KWARGS_KEY) + executor_kwargs: Any = resolved.get(executor_id) + if global_kwargs is None and executor_kwargs is None: + return None + + if global_kwargs is not None and not isinstance(global_kwargs, dict): + logger.warning( + "Executor %s expected a dict for global kwargs, but got %s. Ignoring.", + executor_id, + cast(type[Any], type(global_kwargs)), + ) + return None + + if executor_kwargs is not None and not isinstance(executor_kwargs, dict): + logger.warning( + "Executor %s expected a dict for its kwargs, but got %s. Ignoring.", + executor_id, + cast(type[Any], type(executor_kwargs)), + ) + return None + + # Specific values override global values for the same function argument. + return {**(global_kwargs or {}), **(executor_kwargs or {})} + + +def prepare_agent_run_args( + executor_id: str, + raw_run_kwargs: dict[str, Any], +) -> tuple[dict[str, Any] | None, dict[str, Any] | None]: + """Prepare function_invocation_kwargs and client_kwargs for agent.run(). + + Extracts ``function_invocation_kwargs`` and ``client_kwargs`` from the workflow state + dict, resolving per-executor entries using ``executor_id``. The ``__global__`` sentinel + key (set by ``Workflow._resolve_invocation_kwargs``) denotes global kwargs that apply to + all executors. Per-executor dicts use executor IDs as keys; only the entry for + ``executor_id`` is extracted. + + Args: + executor_id: The id of the executor about to invoke the agent. + raw_run_kwargs: The workflow state dict stored under ``WORKFLOW_RUN_KWARGS_KEY``. + + Returns: + A 2-tuple of (function_invocation_kwargs, client_kwargs). + """ + function_invocation_kwargs = resolve_executor_kwargs(executor_id, raw_run_kwargs.get("function_invocation_kwargs")) + client_kwargs = resolve_executor_kwargs(executor_id, raw_run_kwargs.get("client_kwargs")) + + return function_invocation_kwargs, client_kwargs diff --git a/python/packages/orchestrations/agent_framework_orchestrations/_group_chat.py b/python/packages/orchestrations/agent_framework_orchestrations/_group_chat.py index bbb61edae28..879e09236cf 100644 --- a/python/packages/orchestrations/agent_framework_orchestrations/_group_chat.py +++ b/python/packages/orchestrations/agent_framework_orchestrations/_group_chat.py @@ -32,8 +32,9 @@ from agent_framework import Agent, AgentResponse, AgentResponseUpdate, AgentSession, Message, SupportsAgentRun from agent_framework._telemetry import mark_feature_used from agent_framework._workflows._agent_executor import AgentExecutor, AgentExecutorRequest, AgentExecutorResponse -from agent_framework._workflows._agent_utils import resolve_agent_id +from agent_framework._workflows._agent_utils import prepare_agent_run_args, resolve_agent_id from agent_framework._workflows._checkpoint import CheckpointStorage +from agent_framework._workflows._const import WORKFLOW_RUN_KWARGS_KEY from agent_framework._workflows._executor import Executor from agent_framework._workflows._workflow import Workflow from agent_framework._workflows._workflow_context import WorkflowContext @@ -353,7 +354,7 @@ async def _handle_messages( ): return - agent_orchestration_output = await self._invoke_agent() + agent_orchestration_output = await self._invoke_agent(cast(WorkflowContext[Any, Any], ctx)) if await self._check_agent_terminate_and_yield( agent_orchestration_output, cast(WorkflowContext[Never, AgentResponse | AgentResponseUpdate], ctx), @@ -393,7 +394,7 @@ async def _handle_response( ): return - agent_orchestration_output = await self._invoke_agent() + agent_orchestration_output = await self._invoke_agent(cast(WorkflowContext[Any, Any], ctx)) if await self._check_agent_terminate_and_yield( agent_orchestration_output, cast(WorkflowContext[Never, AgentResponse | AgentResponseUpdate], ctx), @@ -486,8 +487,16 @@ def _parse_agent_output(cls, agent_response: Any) -> AgentOrchestrationOutput: raise ValueError("Failed to parse agent orchestration output.") from last_error - async def _invoke_agent(self) -> AgentOrchestrationOutput: - """Invoke the orchestrator agent to determine the next speaker and termination.""" + async def _invoke_agent(self, ctx: WorkflowContext[Any, Any]) -> AgentOrchestrationOutput: + """Invoke the orchestrator agent to determine the next speaker and termination. + + Args: + ctx: The workflow context, read for the run kwargs stored by ``Workflow.run``. + The orchestrator agent runs outside ``AgentExecutor``, so it has to resolve + those itself or it is the only agent in the group chat that does not see them. + """ + raw_run_kwargs: dict[str, Any] = ctx.get_state(WORKFLOW_RUN_KWARGS_KEY, {}) + function_invocation_kwargs, client_kwargs = prepare_agent_run_args(self.id, raw_run_kwargs) async def _invoke_agent_helper(conversation: list[Message]) -> AgentOrchestrationOutput: # Run the agent in non-streaming mode for simplicity @@ -495,6 +504,8 @@ async def _invoke_agent_helper(conversation: list[Message]) -> AgentOrchestratio messages=conversation, session=self._session, options={"response_format": AgentOrchestrationOutput}, + function_invocation_kwargs=function_invocation_kwargs, + client_kwargs=client_kwargs, ) # Parse and validate the structured output agent_orchestration_output = self._parse_agent_output(agent_response) diff --git a/python/packages/orchestrations/tests/test_group_chat.py b/python/packages/orchestrations/tests/test_group_chat.py index 935111ccaad..92d600ba9ce 100644 --- a/python/packages/orchestrations/tests/test_group_chat.py +++ b/python/packages/orchestrations/tests/test_group_chat.py @@ -275,6 +275,78 @@ async def test_group_chat_as_agent_accepts_conversation() -> None: assert response.messages, "Expected agent conversation output" +class KwargsRecordingManagerAgent(StubManagerAgent): + """Manager agent that records the run kwargs it was invoked with.""" + + def __init__(self) -> None: + super().__init__() + self.seen_kwargs: list[dict[str, Any]] = [] + + async def run( # type: ignore[override] # ty: ignore[invalid-method-override] + self, + messages: str | Content | Message | Sequence[str | Content | Message] | None = None, + *, + session: AgentSession | None = None, + **kwargs: Any, + ) -> AgentResponse[Any]: + self.seen_kwargs.append(dict(kwargs)) + return await super().run(messages, session=session, **kwargs) + + +async def test_agent_manager_receives_workflow_run_kwargs() -> None: + """The orchestrator agent gets the same run kwargs the participants already get. + + ``workflow.run(function_invocation_kwargs=...)`` is stored in workflow state and + ``AgentExecutor`` forwards it to every participant agent. The orchestrator agent runs + outside ``AgentExecutor``, so it has to resolve the same state itself or hosts that put + request-scoped values there (a user id for tool ACL, say) silently get them on the + participants and not on the orchestrator. + """ + manager = KwargsRecordingManagerAgent() + worker = StubAgent("agent", "worker response") + + workflow = GroupChatBuilder( + participants=[worker], + orchestrator_agent=manager, + ).build() + + async for _ in workflow.run( + "coordinate task", + stream=True, + function_invocation_kwargs={"user_id": "user-123"}, + client_kwargs={"trace_id": "trace-abc"}, + ): + pass + + assert manager.seen_kwargs, "Expected the orchestrator agent to be invoked" + for call in manager.seen_kwargs: + assert call.get("function_invocation_kwargs") == {"user_id": "user-123"} + assert call.get("client_kwargs") == {"trace_id": "trace-abc"} + + +async def test_agent_manager_receives_no_run_kwargs_when_none_supplied() -> None: + """With nothing declared on the run, the orchestrator is invoked with both kwargs as None. + + This pins the shape rather than just the happy path: the resolution has to return None, + not an empty dict, so a client that distinguishes the two is not handed a stray {}. + """ + manager = KwargsRecordingManagerAgent() + worker = StubAgent("agent", "worker response") + + workflow = GroupChatBuilder( + participants=[worker], + orchestrator_agent=manager, + ).build() + + async for _ in workflow.run("coordinate task", stream=True): + pass + + assert manager.seen_kwargs, "Expected the orchestrator agent to be invoked" + for call in manager.seen_kwargs: + assert call.get("function_invocation_kwargs") is None + assert call.get("client_kwargs") is None + + async def test_agent_manager_handles_concatenated_json_output() -> None: manager = ConcatenatedJsonManagerAgent() worker = StubAgent("agent", "worker response")