diff --git a/python/packages/core/agent_framework/__init__.py b/python/packages/core/agent_framework/__init__.py index 241bcda3aa..d18ea7a2ed 100644 --- a/python/packages/core/agent_framework/__init__.py +++ b/python/packages/core/agent_framework/__init__.py @@ -362,7 +362,7 @@ "validate_workflow_graph", ), "._workflows._viz": ("WorkflowViz",), - "._workflows._workflow": ("Workflow", "WorkflowRunResult"), + "._workflows._workflow": ("Workflow", "WorkflowInvocationKwargs", "WorkflowRunResult"), "._workflows._workflow_builder": ("WorkflowBuilder",), "._workflows._workflow_context": ("WorkflowContext",), "._workflows._workflow_executor": ( @@ -626,6 +626,7 @@ "WorkflowEventType", "WorkflowException", "WorkflowExecutor", + "WorkflowInvocationKwargs", "WorkflowMessage", "WorkflowRunResult", "WorkflowRunState", diff --git a/python/packages/core/agent_framework/__init__.pyi b/python/packages/core/agent_framework/__init__.pyi index 5816c90600..699f4eedbd 100644 --- a/python/packages/core/agent_framework/__init__.pyi +++ b/python/packages/core/agent_framework/__init__.pyi @@ -325,7 +325,7 @@ from ._workflows._validation import ( validate_workflow_graph, ) from ._workflows._viz import WorkflowViz -from ._workflows._workflow import Workflow, WorkflowRunResult +from ._workflows._workflow import Workflow, WorkflowInvocationKwargs, WorkflowRunResult from ._workflows._workflow_builder import WorkflowBuilder from ._workflows._workflow_context import WorkflowContext from ._workflows._workflow_executor import SubWorkflowRequestMessage, SubWorkflowResponseMessage, WorkflowExecutor @@ -592,6 +592,7 @@ __all__ = [ "WorkflowExecutor", "WorkflowMessage", "WorkflowRunResult", + "WorkflowInvocationKwargs", "WorkflowRunState", "WorkflowRunnerException", "WorkflowValidationError", diff --git a/python/packages/core/agent_framework/_workflows/_agent.py b/python/packages/core/agent_framework/_workflows/_agent.py index 0df47d2b34..8936db68d2 100644 --- a/python/packages/core/agent_framework/_workflows/_agent.py +++ b/python/packages/core/agent_framework/_workflows/_agent.py @@ -45,7 +45,7 @@ from typing_extensions import TypedDict # pragma: no cover if TYPE_CHECKING: - from ._workflow import Workflow + from ._workflow import Workflow, WorkflowInvocationKwargs logger = logging.getLogger(__name__) @@ -152,28 +152,34 @@ def run( self, messages: AgentRunInputs | None = None, *, - stream: Literal[False] = ..., + stream: Literal[True], session: AgentSession | None = None, checkpoint_id: str | None = None, checkpoint_storage: CheckpointStorage | None = None, tools: ToolTypes | Callable[..., Any] | Sequence[ToolTypes | Callable[..., Any]] | None = None, - function_invocation_kwargs: Mapping[str, Any] | None = None, - client_kwargs: Mapping[str, Any] | None = None, - ) -> Awaitable[AgentResponse[Any]]: ... + function_invocation_kwargs: WorkflowInvocationKwargs + | Mapping[str, Mapping[str, Any]] + | Mapping[str, Any] + | None = None, + client_kwargs: WorkflowInvocationKwargs | Mapping[str, Mapping[str, Any]] | Mapping[str, Any] | None = None, + ) -> ResponseStream[AgentResponseUpdate, AgentResponse]: ... @overload - def run( + async def run( self, messages: AgentRunInputs | None = None, *, - stream: Literal[True], + stream: Literal[False] = ..., session: AgentSession | None = None, checkpoint_id: str | None = None, checkpoint_storage: CheckpointStorage | None = None, tools: ToolTypes | Callable[..., Any] | Sequence[ToolTypes | Callable[..., Any]] | None = None, - function_invocation_kwargs: Mapping[str, Any] | None = None, - client_kwargs: Mapping[str, Any] | None = None, - ) -> ResponseStream[AgentResponseUpdate, AgentResponse[Any]]: ... + function_invocation_kwargs: WorkflowInvocationKwargs + | Mapping[str, Mapping[str, Any]] + | Mapping[str, Any] + | None = None, + client_kwargs: WorkflowInvocationKwargs | Mapping[str, Mapping[str, Any]] | Mapping[str, Any] | None = None, + ) -> AgentResponse: ... def run( self, @@ -184,9 +190,12 @@ def run( checkpoint_id: str | None = None, checkpoint_storage: CheckpointStorage | None = None, tools: ToolTypes | Callable[..., Any] | Sequence[ToolTypes | Callable[..., Any]] | None = None, - function_invocation_kwargs: Mapping[str, Any] | None = None, - client_kwargs: Mapping[str, Any] | None = None, - ) -> ResponseStream[AgentResponseUpdate, AgentResponse[Any]] | Awaitable[AgentResponse[Any]]: + function_invocation_kwargs: WorkflowInvocationKwargs + | Mapping[str, Mapping[str, Any]] + | Mapping[str, Any] + | None = None, + client_kwargs: WorkflowInvocationKwargs | Mapping[str, Mapping[str, Any]] | Mapping[str, Any] | None = None, + ) -> ResponseStream[AgentResponseUpdate, AgentResponse] | Awaitable[AgentResponse]: """Get a response from the workflow agent. Args: @@ -254,8 +263,11 @@ async def _run_impl( checkpoint_id: str | None = None, checkpoint_storage: CheckpointStorage | None = None, tools: ToolTypes | Callable[..., Any] | Sequence[ToolTypes | Callable[..., Any]] | None = None, - function_invocation_kwargs: Mapping[str, Mapping[str, Any]] | Mapping[str, Any] | None = None, - client_kwargs: Mapping[str, Mapping[str, Any]] | Mapping[str, Any] | None = None, + function_invocation_kwargs: WorkflowInvocationKwargs + | Mapping[str, Mapping[str, Any]] + | Mapping[str, Any] + | None = None, + client_kwargs: WorkflowInvocationKwargs | Mapping[str, Mapping[str, Any]] | Mapping[str, Any] | None = None, ) -> AgentResponse: """Internal implementation of non-streaming execution. @@ -337,8 +349,11 @@ async def _run_stream_impl( checkpoint_id: str | None = None, checkpoint_storage: CheckpointStorage | None = None, tools: ToolTypes | Callable[..., Any] | Sequence[ToolTypes | Callable[..., Any]] | None = None, - function_invocation_kwargs: Mapping[str, Mapping[str, Any]] | Mapping[str, Any] | None = None, - client_kwargs: Mapping[str, Mapping[str, Any]] | Mapping[str, Any] | None = None, + function_invocation_kwargs: WorkflowInvocationKwargs + | Mapping[str, Mapping[str, Any]] + | Mapping[str, Any] + | None = None, + client_kwargs: WorkflowInvocationKwargs | Mapping[str, Mapping[str, Any]] | Mapping[str, Any] | None = None, ) -> AsyncIterable[AgentResponseUpdate]: """Internal implementation of streaming execution. @@ -419,8 +434,11 @@ async def _run_core( checkpoint_storage: CheckpointStorage | None, streaming: bool, tools: ToolTypes | Callable[..., Any] | Sequence[ToolTypes | Callable[..., Any]] | None = None, - function_invocation_kwargs: Mapping[str, Mapping[str, Any]] | Mapping[str, Any] | None = None, - client_kwargs: Mapping[str, Mapping[str, Any]] | Mapping[str, Any] | None = None, + function_invocation_kwargs: WorkflowInvocationKwargs + | Mapping[str, Mapping[str, Any]] + | Mapping[str, Any] + | None = None, + client_kwargs: WorkflowInvocationKwargs | Mapping[str, Mapping[str, Any]] | Mapping[str, Any] | None = None, ) -> AsyncIterable[WorkflowEvent]: """Core implementation that yields workflow events for both streaming and non-streaming modes. @@ -470,8 +488,7 @@ async def _run_core( # NOTE: It is possible that some pending requests are not fulfilled, # and we will let the workflow to handle this -- the agent does not # have an opinion on this. - pending_requests = await self.workflow._runner_context.get_pending_request_info_events() # pyright: ignore[reportPrivateUsage] - function_responses = self._extract_function_responses(input_messages, pending_requests) + function_responses = self._extract_function_responses(input_messages) if streaming: async for event in self.workflow.run( responses=function_responses, @@ -748,51 +765,22 @@ def _process_request_info_event( arguments=args, ) - def _extract_function_responses( - self, - input_messages: Sequence[Message], - pending_requests: Mapping[str, WorkflowEvent[Any]] | None = None, - ) -> dict[str, Any]: + def _extract_function_responses(self, input_messages: Sequence[Message]) -> dict[str, Any]: """Extract function responses from input messages. The responses are for pending requests that the workflow is waiting on, and will be passed to the workflow. The pending requests are processed to either `function_approval_request` or `function_call` content by `_process_request_info_event`. """ - pending_requests = pending_requests or {} function_responses: dict[str, Any] = {} for message in input_messages: for content in message.contents: if content.type == "function_approval_response": - request_id = content.id - if request_id is None: - raise AgentInvalidResponseException("Function approval response is missing its request ID.") + request_id: str = content.id # type: ignore[assignment] function_responses[request_id] = content elif content.type == "function_result": - request_id = content.call_id - if request_id is None: - raise AgentInvalidResponseException("Function result is missing its call ID.") - response_request_id = request_id - pending_request = pending_requests.get(response_request_id) - if pending_request is None: - matching_requests = [ - (pending_id, pending_event) - for pending_id, pending_event in pending_requests.items() - if isinstance(pending_event.data, Content) - and pending_event.data.type == "function_call" - and pending_event.data.call_id == request_id - ] - if len(matching_requests) == 1: - response_request_id, pending_request = matching_requests[0] - response_data = ( - content - if pending_request is not None - and pending_request.response_type is Content - and isinstance(pending_request.data, Content) - and pending_request.data.type == "function_call" - else content.result - ) - function_responses[response_request_id] = response_data + response_data = content.result if hasattr(content, "result") else str(content) + function_responses[content.call_id] = response_data # type: ignore else: raise AgentInvalidResponseException( "Unexpected content type while awaiting request info responses." diff --git a/python/packages/core/agent_framework/_workflows/_agent_executor.py b/python/packages/core/agent_framework/_workflows/_agent_executor.py index 512706e36f..1c3a7203f4 100644 --- a/python/packages/core/agent_framework/_workflows/_agent_executor.py +++ b/python/packages/core/agent_framework/_workflows/_agent_executor.py @@ -614,22 +614,26 @@ def _resolve_executor_kwargs(self, resolved: dict[str, Any] | None) -> dict[str, """ if not isinstance(resolved, dict): return None - # Use explicit key-presence checks so that an empty per-executor dict is - # honoured (e.g. to clear kwargs) instead of falling through to global. - if self.id in resolved: - executor_kwargs = resolved[self.id] - elif GLOBAL_KWARGS_KEY in resolved: - executor_kwargs = resolved[GLOBAL_KWARGS_KEY] - else: + 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 not isinstance(executor_kwargs, dict): + if global_kwargs is not None and not isinstance(global_kwargs, dict): logger.warning( - "Executor %s expected a dict for its kwargs, but got %s. Ignoring.", + "Executor %s expected a dict for global kwargs, but got %s. Ignoring.", self.id, - type(executor_kwargs), # type: ignore + 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 - return executor_kwargs # type: ignore + # Specific values override global values for the same function argument. + return {**(global_kwargs or {}), **(executor_kwargs or {})} diff --git a/python/packages/core/agent_framework/_workflows/_const.py b/python/packages/core/agent_framework/_workflows/_const.py index e83025bbdc..a84881196a 100644 --- a/python/packages/core/agent_framework/_workflows/_const.py +++ b/python/packages/core/agent_framework/_workflows/_const.py @@ -14,6 +14,10 @@ # to pass kwargs from workflow.run() through to agent.run() and @tool functions. WORKFLOW_RUN_KWARGS_KEY = "_workflow_run_kwargs" +# State keys used to preserve caller-provided kwargs for nested workflow routing. +RAW_FUNCTION_INVOCATION_KWARGS_KEY = "_raw_function_invocation_kwargs" +RAW_CLIENT_KWARGS_KEY = "_raw_client_kwargs" + # Sentinel key used in resolved invocation kwargs dicts to denote global kwargs # that apply to all executors (as opposed to per-executor keyed entries). GLOBAL_KWARGS_KEY = "__global__" diff --git a/python/packages/core/agent_framework/_workflows/_workflow.py b/python/packages/core/agent_framework/_workflows/_workflow.py index 65be066e06..b37cd6fd72 100644 --- a/python/packages/core/agent_framework/_workflows/_workflow.py +++ b/python/packages/core/agent_framework/_workflows/_workflow.py @@ -22,7 +22,14 @@ from ..exceptions import WorkflowException from ..observability import OtelAttr, capture_exception, create_workflow_span from ._checkpoint import CheckpointStorage -from ._const import DEFAULT_MAX_ITERATIONS, GLOBAL_KWARGS_KEY, INTERNAL_SOURCE_ID, WORKFLOW_RUN_KWARGS_KEY +from ._const import ( + DEFAULT_MAX_ITERATIONS, + GLOBAL_KWARGS_KEY, + INTERNAL_SOURCE_ID, + RAW_CLIENT_KWARGS_KEY, + RAW_FUNCTION_INVOCATION_KWARGS_KEY, + WORKFLOW_RUN_KWARGS_KEY, +) from ._edge import ( EdgeGroup, FanOutEdgeGroup, @@ -206,6 +213,20 @@ def classify(self, executor_id: str) -> Literal["output", "intermediate"] | None return None +@dataclass(frozen=True) +class WorkflowInvocationKwargs: + """Explicit global and executor-specific kwargs for a workflow run. + + Use this wrapper when shared kwargs should be combined with executor-specific + overrides. Plain mappings retain their existing global or per-executor behavior. + """ + + global_kwargs: Mapping[str, Any] = field(default_factory=lambda: dict[str, Any]()) + executor_kwargs: Mapping[str, Mapping[str, Any]] = field( + default_factory=lambda: dict[str, Mapping[str, Any]]() + ) + + class Workflow(DictConvertible): """A graph-based execution engine that orchestrates connected executors. @@ -482,8 +503,11 @@ async def _run_workflow_with_tracing( is_continuation: bool = False, streaming: bool = False, tools: ToolTypes | Callable[..., Any] | Sequence[ToolTypes | Callable[..., Any]] | None = None, - function_invocation_kwargs: Mapping[str, Mapping[str, Any]] | Mapping[str, Any] | None = None, - client_kwargs: Mapping[str, Mapping[str, Any]] | Mapping[str, Any] | None = None, + function_invocation_kwargs: WorkflowInvocationKwargs + | Mapping[str, Mapping[str, Any]] + | Mapping[str, Any] + | None = None, + client_kwargs: WorkflowInvocationKwargs | Mapping[str, Mapping[str, Any]] | Mapping[str, Any] | None = None, ) -> AsyncIterable[WorkflowEvent]: """Private method to run workflow with proper tracing. @@ -520,6 +544,7 @@ async def _run_workflow_with_tracing( OtelAttr.WORKFLOW_RUN_SPAN, attributes, ) as span: + saw_request = False emitted_in_progress_pending = False try: # Add workflow started event (telemetry + surface state to consumers) @@ -558,10 +583,18 @@ async def _run_workflow_with_tracing( combined_kwargs["function_invocation_kwargs"] = self._resolve_invocation_kwargs( function_invocation_kwargs, "function_invocation_kwargs" ) + if isinstance(function_invocation_kwargs, WorkflowInvocationKwargs) or any( + isinstance(value, Mapping) for value in function_invocation_kwargs.values() + ): + combined_kwargs[RAW_FUNCTION_INVOCATION_KWARGS_KEY] = function_invocation_kwargs if client_kwargs is not None: combined_kwargs["client_kwargs"] = self._resolve_invocation_kwargs( client_kwargs, "client_kwargs" ) + if isinstance(client_kwargs, WorkflowInvocationKwargs) or any( + isinstance(value, Mapping) for value in client_kwargs.values() + ): + combined_kwargs[RAW_CLIENT_KWARGS_KEY] = client_kwargs self._runner.state.set(WORKFLOW_RUN_KWARGS_KEY, combined_kwargs) elif not is_continuation: self._runner.state.set(WORKFLOW_RUN_KWARGS_KEY, {}) @@ -576,6 +609,9 @@ async def _run_workflow_with_tracing( # All executor executions happen within workflow span async for event in self._runner.run_until_convergence(): + # Track request events for final status determination + if event.type == "request_info": + saw_request = True yield event if event.type == "request_info" and not emitted_in_progress_pending: @@ -584,11 +620,8 @@ async def _run_workflow_with_tracing( with _framework_event_origin(): pending_status = WorkflowEvent.status(self._status) yield pending_status - # Workflow runs until idle - emit final status based on whether requests are pending. - # Continuations such as cancellation may retain an existing sibling request without - # re-emitting its request_info event during this run. - pending_requests = await self._runner.context.get_pending_request_info_events() - if pending_requests: + # Workflow runs until idle - emit final status based on whether requests are pending + if saw_request: self._status = WorkflowRunState.IDLE_WITH_PENDING_REQUESTS with _framework_event_origin(): terminal_status = WorkflowEvent.status(self._status) @@ -691,8 +724,14 @@ def run( checkpoint_id: str | None = None, checkpoint_storage: CheckpointStorage | None = None, tools: ToolTypes | Callable[..., Any] | Sequence[ToolTypes | Callable[..., Any]] | None = None, - function_invocation_kwargs: Mapping[str, Any] | None = None, - client_kwargs: Mapping[str, Any] | None = None, + function_invocation_kwargs: WorkflowInvocationKwargs + | Mapping[str, Mapping[str, Any]] + | Mapping[str, Any] + | None = None, + client_kwargs: WorkflowInvocationKwargs + | Mapping[str, Mapping[str, Any]] + | Mapping[str, Any] + | None = None, ) -> ResponseStream[WorkflowEvent, WorkflowRunResult]: ... @overload @@ -706,8 +745,8 @@ def run( checkpoint_storage: CheckpointStorage | None = None, include_status_events: bool = False, tools: ToolTypes | Callable[..., Any] | Sequence[ToolTypes | Callable[..., Any]] | None = None, - function_invocation_kwargs: Mapping[str, Any] | None = None, - client_kwargs: Mapping[str, Any] | None = None, + function_invocation_kwargs: WorkflowInvocationKwargs | Mapping[str, Any] | None = None, + client_kwargs: WorkflowInvocationKwargs | Mapping[str, Any] | None = None, ) -> Awaitable[WorkflowRunResult]: ... def run( @@ -720,8 +759,11 @@ def run( checkpoint_storage: CheckpointStorage | None = None, include_status_events: bool = False, tools: ToolTypes | Callable[..., Any] | Sequence[ToolTypes | Callable[..., Any]] | None = None, - function_invocation_kwargs: Mapping[str, Mapping[str, Any]] | Mapping[str, Any] | None = None, - client_kwargs: Mapping[str, Mapping[str, Any]] | Mapping[str, Any] | None = None, + function_invocation_kwargs: WorkflowInvocationKwargs + | Mapping[str, Mapping[str, Any]] + | Mapping[str, Any] + | None = None, + client_kwargs: WorkflowInvocationKwargs | Mapping[str, Mapping[str, Any]] | Mapping[str, Any] | None = None, ) -> ResponseStream[WorkflowEvent, WorkflowRunResult] | Awaitable[WorkflowRunResult]: """Run the workflow, optionally streaming events. @@ -746,10 +788,14 @@ def run( tools: Runtime tools available to agent executors. function_invocation_kwargs: Keyword arguments forwarded to tool invocations in subagents. Either a mapping for agent name or agent executor id to kwargs, - or a flat mapping of kwargs for all tool invocations. + a flat mapping of kwargs for all tool invocations, or a + ``WorkflowInvocationKwargs`` instance to combine global and executor-specific + kwargs. client_kwargs: Keyword arguments forwarded to chat client calls in subagents. Either a mapping for agent name or agent executor id to kwargs, - or a flat mapping of kwargs for all chat client calls. + a flat mapping of kwargs for all chat client calls, or a + ``WorkflowInvocationKwargs`` instance to combine global and executor-specific + kwargs. Returns: When stream=True: A ResponseStream[WorkflowEvent, WorkflowRunResult] for @@ -791,7 +837,7 @@ def run( checkpoint_id=checkpoint_id, checkpoint_storage=checkpoint_storage, streaming=stream, - tools=runtime_tools, + tools=tools, function_invocation_kwargs=function_invocation_kwargs, client_kwargs=client_kwargs, ), @@ -812,8 +858,11 @@ async def _run_core( checkpoint_storage: CheckpointStorage | None = None, streaming: bool = False, tools: ToolTypes | Callable[..., Any] | Sequence[ToolTypes | Callable[..., Any]] | None = None, - function_invocation_kwargs: Mapping[str, Mapping[str, Any]] | Mapping[str, Any] | None = None, - client_kwargs: Mapping[str, Mapping[str, Any]] | Mapping[str, Any] | None = None, + function_invocation_kwargs: WorkflowInvocationKwargs + | Mapping[str, Mapping[str, Any]] + | Mapping[str, Any] + | None = None, + client_kwargs: WorkflowInvocationKwargs | Mapping[str, Mapping[str, Any]] | Mapping[str, Any] | None = None, ) -> AsyncIterable[WorkflowEvent]: """Single core execution path for both streaming and non-streaming modes. @@ -1069,7 +1118,7 @@ def _get_executor_by_id(self, executor_id: str) -> Executor: def _resolve_invocation_kwargs( self, - kwargs: Mapping[str, Any], + kwargs: WorkflowInvocationKwargs | Mapping[str, Any], param_name: str, ) -> dict[str, Any]: """Resolve invocation kwargs into a normalized per-executor or global format. @@ -1077,17 +1126,24 @@ def _resolve_invocation_kwargs( Detects whether the provided kwargs dict uses per-executor targeting by checking if any top-level key matches a known executor ID in the workflow. If at least one key matches, all entries are treated as per-executor. Otherwise the dict is treated - as global kwargs that apply to every executor. + as global kwargs that apply to every executor. The ``"__global__"`` key can be used + explicitly to combine global kwargs with per-executor overrides. Args: kwargs: The raw invocation kwargs from the caller. param_name: The parameter name (for logging), e.g. ``"function_invocation_kwargs"``. Returns: - A dict with either: - - ``{"__global__": }`` for global kwargs, or - - The original dict unchanged for per-executor kwargs. + A dict containing normalized global or per-executor mappings. """ + if isinstance(kwargs, WorkflowInvocationKwargs): + resolved = {GLOBAL_KWARGS_KEY: dict(kwargs.global_kwargs)} + resolved.update({ + executor_id: dict(executor_kwargs) for executor_id, executor_kwargs in kwargs.executor_kwargs.items() + }) + logger.info("Explicit global %s provided with executor-specific overrides.", param_name) + return resolved + executor_ids = set(self.executors.keys()) matched_ids = kwargs.keys() & executor_ids if matched_ids: @@ -1211,29 +1267,10 @@ async def cancel_pending_requests( checkpoint_id: str | None = None, checkpoint_storage: CheckpointStorage | None = None, tools: ToolTypes | Callable[..., Any] | Sequence[ToolTypes | Callable[..., Any]] | None = None, - function_invocation_kwargs: Mapping[str, Mapping[str, Any]] | Mapping[str, Any] | None = None, - client_kwargs: Mapping[str, Mapping[str, Any]] | Mapping[str, Any] | None = None, + function_invocation_kwargs: WorkflowInvocationKwargs | Mapping[str, Any] | None = None, + client_kwargs: WorkflowInvocationKwargs | Mapping[str, Any] | None = None, ) -> WorkflowRunResult: - """Cancel pending external requests and release their owning executor state. - - Cancellation follows requests through nested workflows and clears any executor-owned - correlation without synthesizing a response. If cancellation drains an executor's pending - set after sibling responses were already accepted, the executor resumes through its normal - continuation path. Unknown or already-handled request IDs are ignored. - - Args: - request_ids: Request identifiers to cancel. - - Keyword Args: - checkpoint_id: Checkpoint to restore before applying cancellation. - checkpoint_storage: Runtime checkpoint storage for the cancellation continuation. - tools: Request-scoped tools available while cancellation resumes executors. - function_invocation_kwargs: Keyword arguments forwarded to resumed tool invocations. - client_kwargs: Keyword arguments forwarded to resumed chat client calls. - - Returns: - Events produced while applying cancellation and any resulting continuation. - """ + """Cancel pending external requests and continue the workflow.""" selected_ids = set(request_ids) if not all(isinstance(request_id, str) and request_id for request_id in selected_ids): raise ValueError("Pending workflow request IDs must be non-empty strings.") @@ -1250,10 +1287,7 @@ async def apply_cancellations() -> None: state=self._runner.state, runner_context=self._runner.context, ) - await executor._cancel_pending_request( # pyright: ignore[reportPrivateUsage] - request_id, - context, - ) + await executor._cancel_pending_request(request_id, context) # pyright: ignore[reportPrivateUsage] if checkpoint_storage is not None: self._runner.context.set_runtime_checkpoint_storage(checkpoint_storage) diff --git a/python/packages/core/agent_framework/_workflows/_workflow_executor.py b/python/packages/core/agent_framework/_workflows/_workflow_executor.py index f3bccd1ff1..18f581aa5e 100644 --- a/python/packages/core/agent_framework/_workflows/_workflow_executor.py +++ b/python/packages/core/agent_framework/_workflows/_workflow_executor.py @@ -4,13 +4,19 @@ import logging import sys import types +from collections.abc import Mapping from dataclasses import dataclass -from typing import TYPE_CHECKING, Any +from typing import TYPE_CHECKING, Any, cast if TYPE_CHECKING: from ._workflow import Workflow -from ._const import GLOBAL_KWARGS_KEY, WORKFLOW_RUN_KWARGS_KEY +from ._const import ( + GLOBAL_KWARGS_KEY, + RAW_CLIENT_KWARGS_KEY, + RAW_FUNCTION_INVOCATION_KWARGS_KEY, + WORKFLOW_RUN_KWARGS_KEY, +) from ._events import ( WorkflowEvent, WorkflowRunState, @@ -20,7 +26,7 @@ from ._request_info_mixin import response_handler from ._runner_context import WorkflowMessage from ._typing_utils import is_instance_of -from ._workflow import WorkflowRunResult +from ._workflow import WorkflowInvocationKwargs, WorkflowRunResult from ._workflow_context import WorkflowContext if sys.version_info >= (3, 12): @@ -375,29 +381,36 @@ async def process_workflow(self, input_data: object, ctx: WorkflowContext[Any, A # Get kwargs from parent workflow's State to propagate to subworkflow parent_kwargs: dict[str, Any] = ctx.get_state(WORKFLOW_RUN_KWARGS_KEY, {}) - # Extract invocation kwargs recognised by Workflow.run() - # The state stores resolved format (with __global__ wrapper for global kwargs). - # Unwrap __global__ before passing to the subworkflow so it gets re-resolved - # against the subworkflow's own executor IDs. - fi_kwargs: dict[str, Any] | None = None - ci_kwargs: dict[str, Any] | None = None - tools = ctx.get_runtime_tools() + # Use the caller's raw kwargs so legacy per-executor mappings are resolved + # against the child workflow's executor IDs rather than the parent's. + fi_kwargs: WorkflowInvocationKwargs | Mapping[str, Any] | None = None + ci_kwargs: WorkflowInvocationKwargs | Mapping[str, Any] | None = None for key in ("function_invocation_kwargs", "client_kwargs"): - resolved = parent_kwargs.get(key) - if isinstance(resolved, dict): - # Unwrap global sentinel; pass per-executor dicts as-is - unwrapped: dict[str, Any] = resolved.get(GLOBAL_KWARGS_KEY, resolved) # type: ignore + raw_key = ( + RAW_FUNCTION_INVOCATION_KWARGS_KEY if key == "function_invocation_kwargs" else RAW_CLIENT_KWARGS_KEY + ) + raw_value = parent_kwargs.get(raw_key) + if raw_value is not None: + resolved = cast(WorkflowInvocationKwargs | Mapping[str, Any], raw_value) + else: + normalized: Any = parent_kwargs.get(key) + if isinstance(normalized, dict): + normalized_dict = cast(dict[str, Any], normalized) + if len(normalized_dict) == 1 and GLOBAL_KWARGS_KEY in normalized_dict: + normalized = normalized_dict[GLOBAL_KWARGS_KEY] + resolved = cast(WorkflowInvocationKwargs | Mapping[str, Any] | None, normalized) + if resolved is not None: if key == "function_invocation_kwargs": - fi_kwargs = unwrapped # type: ignore + fi_kwargs = resolved else: - ci_kwargs = unwrapped # type: ignore + ci_kwargs = resolved # Run the sub-workflow and collect all events, passing parent kwargs result = await self.workflow.run( input_data, - tools=tools, - function_invocation_kwargs=fi_kwargs, # type: ignore - client_kwargs=ci_kwargs, # type: ignore + tools=ctx.get_runtime_tools(), + function_invocation_kwargs=fi_kwargs, + client_kwargs=ci_kwargs, ) logger.debug(f"WorkflowExecutor {self.id} sub-workflow {self.workflow.id} completed with {len(result)} events") @@ -450,15 +463,6 @@ async def handle_propagated_request_response( ctx=ctx, ) - @override - async def _cancel_pending_request(self, request_id: str, ctx: WorkflowContext[Any, Any]) -> None: - """Propagate cancellation into the wrapped workflow.""" - result = await self.workflow.cancel_pending_requests( - [request_id], - tools=ctx.get_runtime_tools(), - ) - await self._process_workflow_result(result, ctx) - @override async def on_checkpoint_save(self) -> dict[str, Any]: """Get the current state of the WorkflowExecutor for checkpointing purposes.""" @@ -617,5 +621,5 @@ async def _handle_response( # Forward the response to the sub-workflow, which resumes and validates it against its own # pending requests, then process whatever the sub-workflow produces. - result = await self.workflow.run(responses={request_id: response}, tools=ctx.get_runtime_tools()) + result = await self.workflow.run(responses={request_id: response}) await self._process_workflow_result(result, ctx) diff --git a/python/packages/core/tests/core/test_serializable_mixin.py b/python/packages/core/tests/core/test_serializable_mixin.py index d5b14cb001..59c4238dfe 100644 --- a/python/packages/core/tests/core/test_serializable_mixin.py +++ b/python/packages/core/tests/core/test_serializable_mixin.py @@ -586,6 +586,7 @@ def __init__(self, raw_representation: Any): def test_pickle_restores_slot_fields(self): """Pickle state should include fields declared in slots.""" + class TestClass(SerializationMixin): __slots__ = ("value",) @@ -614,6 +615,7 @@ def __init__(self): def test_pickle_omission_is_separate_from_shallow_copy_policy(self): """Fields shallow-copied by default remain persistent unless explicitly omitted.""" + class TestClass(SerializationMixin): _PICKLE_OMIT_FIELDS = set() diff --git a/python/packages/core/tests/workflow/test_agent_executor.py b/python/packages/core/tests/workflow/test_agent_executor.py index 2cc2ed2ce6..9d7e4978a0 100644 --- a/python/packages/core/tests/workflow/test_agent_executor.py +++ b/python/packages/core/tests/workflow/test_agent_executor.py @@ -1,7 +1,6 @@ # Copyright (c) Microsoft. All rights reserved. import pickle - from collections.abc import AsyncIterable, Awaitable from typing import Any, Literal, overload @@ -338,9 +337,6 @@ class _NonCopyableRaw: def __deepcopy__(self, memo: dict) -> Any: raise TypeError("Cannot deepcopy this object") - def __reduce__(self) -> Any: - raise TypeError("Cannot pickle this object") - class _AgentWithRawRepr(BaseAgent): """Agent that returns responses with a non-copyable raw_representation.""" @@ -392,20 +388,6 @@ async def test_agent_executor_workflow_with_non_copyable_raw_representation() -> assert agent_responses[0].raw_representation is raw -def test_serialization_mixin_omits_non_pickleable_raw_representation() -> None: - """Pickling framework objects should not include runtime-only raw representations.""" - raw = _NonCopyableRaw() - response = AgentResponse( - messages=[Message("assistant", [Content.from_text(text="reply", raw_representation=raw)])], - raw_representation=raw, - ) - - restored = pickle.loads(pickle.dumps(response)) - - assert restored.raw_representation is None - assert restored.messages[0].contents[0].raw_representation is None - - # --------------------------------------------------------------------------- # Context mode tests # --------------------------------------------------------------------------- @@ -641,15 +623,15 @@ async def test_resolve_executor_kwargs_returns_none_for_none_input() -> None: assert result is None -async def test_resolve_executor_kwargs_prefers_executor_id_over_global() -> None: - """_resolve_executor_kwargs prefers executor-specific entry over __global__.""" +async def test_resolve_executor_kwargs_merges_executor_id_over_global() -> None: + """_resolve_executor_kwargs merges executor-specific entries over __global__.""" agent = _CountingAgent(id="a", name="A") executor = AgentExecutor(agent, id="exec_a") # Dict has both a per-executor entry and a global entry resolved = {"exec_a": {"specific": True}, GLOBAL_KWARGS_KEY: {"global": True}} result = executor._resolve_executor_kwargs(resolved) # pyright: ignore[reportPrivateUsage] - assert result == {"specific": True} + assert result == {"global": True, "specific": True} async def test_prepare_agent_run_args_extracts_function_invocation_kwargs() -> None: @@ -708,16 +690,15 @@ async def test_prepare_agent_run_args_per_executor_no_match() -> None: assert fi_kwargs is None -async def test_resolve_executor_kwargs_empty_per_executor_does_not_fallback_to_global() -> None: - """An explicit empty per-executor dict should not fall through to global kwargs.""" +async def test_resolve_executor_kwargs_empty_per_executor_keeps_global_kwargs() -> None: + """An explicit empty per-executor dict keeps the global kwargs.""" agent = _CountingAgent(id="a", name="A") executor = AgentExecutor(agent, id="exec_a") - # Per-executor entry for exec_a is empty, but global has values. - # The empty dict should be honoured (no fallback to global). + # Per-executor entry for exec_a is empty, so only global values apply. resolved = {"exec_a": {}, GLOBAL_KWARGS_KEY: {"global_key": "global_val"}} # type: ignore[var-annotated] result = executor._resolve_executor_kwargs(resolved) # pyright: ignore[reportPrivateUsage] - assert result == {} + assert result == {"global_key": "global_val"} # region Tool approval emission diff --git a/python/packages/core/tests/workflow/test_workflow_kwargs.py b/python/packages/core/tests/workflow/test_workflow_kwargs.py index c21a25a5d7..d91d224c54 100644 --- a/python/packages/core/tests/workflow/test_workflow_kwargs.py +++ b/python/packages/core/tests/workflow/test_workflow_kwargs.py @@ -17,6 +17,7 @@ FunctionTool, Message, ResponseStream, + WorkflowInvocationKwargs, WorkflowRunState, ) from agent_framework._workflows._const import WORKFLOW_RUN_KWARGS_KEY @@ -1070,6 +1071,35 @@ async def test_nested_subworkflow_kwargs_propagation() -> None: ) +async def test_mixed_kwargs_route_through_subworkflow() -> None: + """Mixed kwargs preserve global values and child executor-specific routing.""" + from agent_framework._workflows._workflow_executor import WorkflowExecutor + + inner_agent1 = _KwargsCapturingAgent(name="inner_agent1") + inner_agent2 = _KwargsCapturingAgent(name="inner_agent2") + inner_workflow = SequentialBuilder(participants=[inner_agent1, inner_agent2]).build() + subworkflow_executor = WorkflowExecutor(workflow=inner_workflow, id="subworkflow") + outer_workflow = SequentialBuilder(participants=[subworkflow_executor]).build() + + fi_kwargs = WorkflowInvocationKwargs( + global_kwargs={"shared": "value", "overridden": "global"}, + executor_kwargs={"inner_agent2": {"overridden": "inner_agent2"}}, + ) + + async for event in outer_workflow.run("test", stream=True, function_invocation_kwargs=fi_kwargs): + if event.type == "status" and event.state == WorkflowRunState.IDLE: + break + + assert inner_agent1.captured_kwargs[0].get("function_invocation_kwargs") == { + "shared": "value", + "overridden": "global", + } + assert inner_agent2.captured_kwargs[0].get("function_invocation_kwargs") == { + "shared": "value", + "overridden": "inner_agent2", + } + + # endregion @@ -1151,6 +1181,35 @@ async def test_per_executor_function_invocation_kwargs_routes_to_correct_agent() assert agent2.captured_kwargs[0].get("function_invocation_kwargs") == {"tool_param": "value_for_agent2"} +async def test_global_and_per_executor_function_invocation_kwargs_are_merged() -> None: + """Global function kwargs are merged with executor-specific overrides.""" + agent1 = _KwargsCapturingAgent(name="agent1") + agent2 = _KwargsCapturingAgent(name="agent2") + workflow = SequentialBuilder(participants=[agent1, agent2]).build() + + fi_kwargs = WorkflowInvocationKwargs( + global_kwargs={"shared": "value", "overridden": "global"}, + executor_kwargs={ + "agent1": {"overridden": "agent1"}, + "agent2": {"agent_only": True}, + }, + ) + + async for event in workflow.run("test", stream=True, function_invocation_kwargs=fi_kwargs): + if event.type == "status" and event.state == WorkflowRunState.IDLE: + break + + assert agent1.captured_kwargs[0].get("function_invocation_kwargs") == { + "shared": "value", + "overridden": "agent1", + } + assert agent2.captured_kwargs[0].get("function_invocation_kwargs") == { + "shared": "value", + "overridden": "global", + "agent_only": True, + } + + async def test_per_executor_kwargs_unmatched_agent_gets_none() -> None: """An agent not targeted in per-executor kwargs should receive None for that kwarg.""" agent1 = _KwargsCapturingAgent(name="agent1") @@ -1221,6 +1280,35 @@ async def test_per_executor_client_kwargs_routes_correctly() -> None: assert agent2.captured_kwargs[0].get("client_kwargs") == {"temperature": 0.9} +async def test_global_and_per_executor_client_kwargs_are_merged() -> None: + """Global client kwargs are merged with executor-specific overrides.""" + agent1 = _KwargsCapturingAgent(name="agent1") + agent2 = _KwargsCapturingAgent(name="agent2") + workflow = SequentialBuilder(participants=[agent1, agent2]).build() + + ci_kwargs = WorkflowInvocationKwargs( + global_kwargs={"shared": "value", "overridden": "global"}, + executor_kwargs={ + "agent1": {"overridden": "agent1"}, + "agent2": {"agent_only": True}, + }, + ) + + async for event in workflow.run("test", stream=True, client_kwargs=ci_kwargs): + if event.type == "status" and event.state == WorkflowRunState.IDLE: + break + + assert agent1.captured_kwargs[0].get("client_kwargs") == { + "shared": "value", + "overridden": "agent1", + } + assert agent2.captured_kwargs[0].get("client_kwargs") == { + "shared": "value", + "overridden": "global", + "agent_only": True, + } + + async def test_resolve_invocation_kwargs_logs_per_executor(caplog: "LogCaptureFixture") -> None: """Workflow._resolve_invocation_kwargs logs info when per-executor format is detected.""" import logging