diff --git a/docs/decisions/0046-python-declarative-request-scope-recovery.md b/docs/decisions/0046-python-declarative-request-scope-recovery.md new file mode 100644 index 00000000000..6199a27a6ec --- /dev/null +++ b/docs/decisions/0046-python-declarative-request-scope-recovery.md @@ -0,0 +1,189 @@ +--- +status: proposed +contact: baywet +date: 2026-10-15 +deciders: eavanvalkenburg, westey-m +--- + +# Preserve declarative request origins with opt-in legacy recovery + +## Context and Problem Statement + +A declarative workflow can retain requests from several fresh runs. Resuming a +question, tool approval, or agent external-input request must restore its +originating workflow scope rather than inherit the most recent run's MCP session. +Older checkpoints do not retain that origin, so the framework cannot infer which +session should be resumed. + +## Decision Drivers + +- All declarative pause types need the same origin and response-batch rules. +- Existing checkpoints need an explicit recovery path without guessing an origin. +- Repeated restoration must select the same migrated scope without sharing a newer run's session. + +## Considered Options + +- Use the current workflow scope for requests without an origin. This needs no + host changes and retains apparent continuity, but can attach an older request + to an unrelated fresh run. +- Reject all unknown origins and require restarting the workflow. This avoids + guessing, but prevents hosts from recovering otherwise usable pending requests. +- Automatically assign isolated scopes. This permits recovery without host + changes, but silently loses the previous MCP session's server-side state. +- Require explicit host opt-in to assign stable, isolated scopes. This makes the + continuity trade-off visible and permits recovery, but requires host changes + and separate resumes for legacy requests. + +## Decision Outcome + +Choose explicit host opt-in through +`WorkflowFactory(migrate_legacy_request_scopes=True)`, defaulting to `False`: + +- Questions/external input, function approvals, MCP approvals, and agent external + input retain their originating scope in the request. Response handlers restore + it before continuing. Agent requests also retain their own loop configuration. +- Factory-built workflows validate all submitted request origins before consuming + responses. Mixed-scope batches are rejected without changing request metadata. +- Checkpoint restores validate ordinary queued response origins before restoring + shared or executor state, including embedded child checkpoints. Submitted and + ordinary queued responses share batch validation before pending requests are consumed. + Wrapped child replies validate selected routes at delivery, after executor restoration. +- Nested restoration preflights a detached copy of the complete embedded tree + with each child's own policy before applying live state or queues. Migration + metadata is confined to that candidate; a rejected child leaves earlier + siblings, parent executors, and the supplied checkpoint unchanged. +- Without opt-in, a missing origin produces an actionable recovery error. +- With opt-in, a missing origin becomes + `legacy-request:`. The distinct namespace never selects + the active run's scope; preserving the request ID makes migration stable across + workers and repeated restoration. Legacy requests must be resumed separately. +- Existing valid origins are unchanged. Invalid stored origins are not migrated. + Migration does not change header verification or authorize an unverifiable + MCP header context. +- Newly emitted requests and subsequent checkpoints retain the migrated scope. + The old MCP session and its server-side history cannot be reconstructed. + Hosts requiring that continuity must restart rather than opt in. +- Legacy agent loop configuration remains available from its old checkpoint + slot; new requests own their configuration instead of replacing shared state. + +Checkpoint contents and framework request IDs remain trusted host-owned data. +Migration does not reconstruct state already overwritten in a legacy checkpoint. +This is a checkpoint-compatibility change with an explicit recovery path, not +transparent restoration of an old session. + +## Nested entry checkpoint replay + +Selecting a random child scope only when the child first executes is insufficient: +the parent's entry checkpoint predates that selection. Rebuilding the graph and +replaying that checkpoint would select a different child scope. Pre-seeding every +descendant's mutable state before the parent snapshot would instead couple parent +checkpoint creation to child initialization. + +Select a parent scope before its entry checkpoint and derive each child scope from +that scope and the binding ID. Agent turns select a stable scope for their session; +direct fresh runs select new scopes. Declarative response handlers restore the +originating parent scope before further nested invocation. Child scope selection +occurs only after the child accepts a fresh run, not before checkpoint validation +or on response-only continuations. + +Declarative executors manually composed with Core `WorkflowBuilder` read the +active Core scope instead of reusing a stale declarative cache. Factory +initialization still selects its existing declarative scopes before the entry +checkpoint; standalone state helpers retain their previous initialization behavior. +Restored request scopes take precedence over inherited metadata from later runs. + +Older parent checkpoints without this scope cannot reconstruct a child's original +scope if that child had not started. Such parents receive a new isolated scope +with a warning; subsequent checkpoints retain it. Existing request-origin +validation and explicit legacy-request migration rules are unchanged. + +## Nested response boundaries + +Restoring only a child's scope leaves its ancestors on a later fresh run's scope. +If runs A and B pause in the same child and B is answered before A, a downstream +binding could therefore select B's MCP session while processing A's continuation. + +Retain the originating parent run context for each forwarded child request in +framework-owned parent shared state. Unlike transient wrapper fields, this state +is checkpointed with its owning workflow, including at every embedded boundary. +After a child successfully resumes or cancels a request, restore the saved parent +context before forwarding outputs or capturing the origins of follow-up requests. +Release completed origins, preserve still-pending origins across new inputs, and +reject mixed-parent-scope ordinary response batches before consuming any requests, +including batches spanning sibling bindings. Cancellation batches obey the same +parent-scope rule because cancellation can also produce downstream outputs. +Checkpoint preflight applies the same validation to queued propagated replies +before replacing live state. Wrapped replies use the delivery boundary below. + +Older checkpoints without boundary origins cannot reconstruct those ancestors' +original scopes. A successful child continuation therefore receives a new isolated +parent scope with a warning rather than reusing the latest parent's session. +This does not bypass the child's declarative request-origin checks or explicit +legacy-request migration policy. + +## Ordinary Core requests and live response preflight + +Ordinary `WorkflowContext.request_info()` requests can also pause before a child +binding. Capturing only forwarded child requests leaves these continuations on the +latest fresh run's scope. Capture their framework-owned run metadata in the same +checkpointed state, and restore it before Core invokes the response handler, +including for queued checkpoint replies. Core cancellation handlers restore the +same metadata. Propagated child responses use this path once; wrapped child +messages retain their boundary-specific capture and restoration. + +Wrapped origins are owned independently by each child binding, not by the +coordinator's external-request ID. Relaying a child request with that same ID +therefore retains two independently consumable origins across checkpoints. +Capturing or releasing an origin stages only that mapping entry in shared state, +with committed values unchanged until the superstep commits. Discard rolls back +these entry updates. Full-map validation and detached copies remain checkpoint +and response-batch operations, not per-request insertion or removal work. + +Preflight includes propagated responses already queued for delivery. Repeat +validation before each live message batch is drained. Cancellation validates ordinary queued replies together with +the complete selected cancellation set before consuming requests or invoking +cancellation handlers. Legacy runner contexts use a non-persisting snapshot when +optional queued-message inspection is unavailable. + +Deep-copying all live request payloads to preflight migrations imposes an unrelated +serialization requirement on non-checkpointed, in-process workflows. Instead, +response policies validate without mutation and return deferred updates to the +scope metadata they own. Apply these updates only after every policy in the tree +has validated. This preserves cross-child rejection atomicity without copying user +payloads or changing their identity. Stored checkpoint restoration still preflights +ordinary responses in its detached candidate tree before replacing live state. + +## State-dependent wrapped-reply routing + +Edge callbacks can read executor fields changed by preceding handlers or restored +from a checkpoint. Evaluating them during queue or detached-checkpoint preflight +and caching their decisions changes normal routing behavior. Copying executors +does not rebind arbitrary callback closures, and speculative handler or restore-hook +execution can have external side effects. An explicit stable/snapshot-aware routing +API would impose a new contract on existing graphs. + +Instead, finish preceding messages from a source before selecting a wrapped +reply's targets. Evaluate each edge callback once, validate all selected targets +together against their own child policies, then deliver that reply. Retain the +accepted parent scope in checkpoint-backed shared state while the continuation +has in-flight messages. Later supersteps and mid-run checkpoint replay cannot +introduce a different scope after earlier origins have been released. Clear this +guard at quiescence before capturing the final checkpoint, so a later response-only +run may resume another origin. An accepted new run with no in-flight work also +resets the guard; restoring an in-flight checkpoint preserves it. Validate stored +guard metadata before replacing live state. +Normal message batches +retain concurrent edge delivery; wrapped replies form delivery barriers, with +their selected targets still executing concurrently. + +This deliberately narrows atomicity for wrapped replies: rejection retains the +rejected and subsequent undelivered messages, but cannot undo earlier handlers, +child continuations, external effects, or already restored checkpoint state. +Recover a partially executed failed run from a known-good checkpoint rather than +retrying its remaining queue. Ordinary response/cancellation batches and detached +structural and ordinary-response checkpoint validation retain their preflight +guarantees. No callback is assumed pure, and no executor is cloned or replayed +speculatively. + +The proposed deciders are Python code owners; maintainer and +engineering-management approval is still required. diff --git a/python/packages/core/AGENTS.md b/python/packages/core/AGENTS.md index aaba1c8bc80..dcbd2a459cf 100644 --- a/python/packages/core/AGENTS.md +++ b/python/packages/core/AGENTS.md @@ -338,6 +338,38 @@ The vector store API is experimental under the shared `VECTOR_STORES` feature ID by that continuation, can atomically restore a supplied checkpoint before cancellation, and returns the resulting `WorkflowRunResult`. Cancelling a computer request also cancels the other pending requests in that agent's batch, retains already resolved sibling evidence in terminal output, and resets its session before new input. +- **Nested workflow scopes** - Fresh runs select a checkpoint-backed scope before the entry snapshot. + `WorkflowExecutor` derives a child scope from the parent scope and binding ID after the child run is accepted. + Agent turns retain their session scope; direct fresh runs select a new scope. Entry replay therefore preserves + child scopes even when child execution had not started at the snapshot. Response-only continuations restore + their originating scope before further nested invocation. Every `WorkflowContext.request_info()` request + retains its run context in checkpointed shared state, including ordinary Core requests before a child binding. + Core restores and releases that context before invoking response or cancellation handlers, including queued + checkpoint replies. Wrapped child replies use the same helpers at their message boundary; propagated replies + are restored only once. Wrapped origins belong to their binding independently of external requests, even when + a coordinator relays the same request ID. Origin insertion and removal stage individual mapping entries; + committed snapshots remain unchanged until commit and discard rolls back pending changes. + Ordinary response and cancellation batches reject mixed parent scopes before consuming requests. + Wrapped child replies select routes immediately before delivery, after preceding handlers on compatible + edge runners have completed, without waiting for unrelated outgoing handlers, and validate all selected targets + together using each child's own policy. Predicates run once per delivery, not during checkpoint or queue + preflight. Fan-in aggregates are not individual wrapped replies. + Selected replies cannot introduce a different parent scope while a continuation has in-flight messages. + The accepted scope is checkpoint-backed across supersteps and clears at quiescence before the final + checkpoint, so a later response-only run can select another origin. A rejected reply + and subsequent messages remain queued, but earlier handlers and child continuations are not rolled back. + Recover a partially executed failed run from a known-good checkpoint rather than retrying its remaining queue. +- **Nested checkpoint restoration** - Stored and embedded restores preflight a detached checkpoint tree before + replacing live state or queues. `WorkflowExecutor` delegates its async private validation hook to its child's runner, + so ordinary queued responses use each child's own policy and migration metadata stays in the candidate snapshot. + Wrapped-reply policies run only at delivery after executor restoration; rejection does not undo that restoration. +- **Response policy staging** - Live response validation does not copy request payloads. Policies validate + without mutation and may return deferred metadata updates; all policies in the workflow tree must succeed + before those updates are applied. This preserves in-process non-copyable payloads and atomic legacy-scope + migration without imposing checkpoint serialization requirements on ordinary requests. +- **Runner context compatibility** - Queued-message inspection is an optional capability, not a required + `RunnerContext` member. Contexts without it use their non-persisting `build_checkpoint` snapshot so response + validation still includes queued replies without draining messages or consuming pending requests. - **`WorkflowBuilder`** - Fluent API for building workflows, including explicit `output_from` / `intermediate_output_from` selection for caller-facing emissions. `output_from` is an allow-list for **Workflow Output**; unselected executor payloads are hidden unless diff --git a/python/packages/core/agent_framework/_workflows/_agent.py b/python/packages/core/agent_framework/_workflows/_agent.py index e71a62a6eb3..a120fab3b0d 100644 --- a/python/packages/core/agent_framework/_workflows/_agent.py +++ b/python/packages/core/agent_framework/_workflows/_agent.py @@ -388,6 +388,7 @@ async def _run_impl( output_events: list[WorkflowEvent[Any]] = [] async for event in self._run_core( session_messages, + session.session_id if session is not None else None, checkpoint_id, checkpoint_storage, streaming=False, @@ -475,6 +476,7 @@ async def _run_stream_impl( role_gate = _StreamingRoleGate() async for event in self._run_core( session_messages, + session.session_id if session is not None else None, checkpoint_id, checkpoint_storage, streaming=True, @@ -496,6 +498,7 @@ async def _run_stream_impl( async def _run_core( self, input_messages: Sequence[Message], + agent_session_id: str | None, checkpoint_id: str | None, checkpoint_storage: CheckpointStorage | None, streaming: bool, @@ -510,6 +513,7 @@ async def _run_core( Args: input_messages: Normalized input messages to process. + agent_session_id: Framework-local identity of the supplied agent session. checkpoint_id: ID of checkpoint to restore from. checkpoint_storage: Runtime checkpoint storage. streaming: Whether to use streaming workflow methods. @@ -527,19 +531,21 @@ async def _run_core( logger.debug(f"Restoring workflow from checkpoint {checkpoint_id}") # Restore the workflow from checkpoint if streaming: - async for _ in self.workflow.run( + restore_stream = self.workflow.run( stream=True, checkpoint_id=checkpoint_id, checkpoint_storage=checkpoint_storage, tools=tools, - ): + ) + async for _ in restore_stream: pass else: - _ = await self.workflow.run( + restore_result = self.workflow.run( checkpoint_id=checkpoint_id, checkpoint_storage=checkpoint_storage, tools=tools, ) + _ = await restore_result if not input_messages: logger.info("No input messages provided; the workflow has been restored to the checkpoint state.") return @@ -557,43 +563,49 @@ async def _run_core( 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) if streaming: - async for event in self.workflow.run( + response_stream = self.workflow.run( responses=function_responses, stream=True, checkpoint_storage=checkpoint_storage, tools=tools, function_invocation_kwargs=function_invocation_kwargs, client_kwargs=client_kwargs, - ): + ) + async for event in response_stream: yield event else: - for event in await self.workflow.run( + response_result = self.workflow.run( responses=function_responses, checkpoint_storage=checkpoint_storage, tools=tools, function_invocation_kwargs=function_invocation_kwargs, client_kwargs=client_kwargs, - ): + ) + for event in await response_result: yield event elif final_state == WorkflowRunState.IDLE: if streaming: - async for event in self.workflow.run( + response_stream = self.workflow.run( message=input_messages, stream=True, checkpoint_storage=checkpoint_storage, tools=tools, function_invocation_kwargs=function_invocation_kwargs, client_kwargs=client_kwargs, - ): + ) + self.workflow._set_agent_run_context(agent_session_id) # pyright: ignore[reportPrivateUsage] + async for event in response_stream: yield event else: - for event in await self.workflow.run( + response_result = self.workflow.run( message=input_messages, checkpoint_storage=checkpoint_storage, tools=tools, function_invocation_kwargs=function_invocation_kwargs, client_kwargs=client_kwargs, - ): + ) + self.workflow._set_agent_run_context(agent_session_id) # pyright: ignore[reportPrivateUsage] + for event in await response_result: yield event else: raise AgentException(f"The underlying workflow is in an invalid state to restart: {final_state}.") diff --git a/python/packages/core/agent_framework/_workflows/_const.py b/python/packages/core/agent_framework/_workflows/_const.py index db99b45426a..17b97f79360 100644 --- a/python/packages/core/agent_framework/_workflows/_const.py +++ b/python/packages/core/agent_framework/_workflows/_const.py @@ -1,5 +1,7 @@ # Copyright (c) Microsoft. All rights reserved. +from __future__ import annotations + # Default maximum iterations for workflow execution. DEFAULT_MAX_ITERATIONS = 100 @@ -17,6 +19,19 @@ # to pass kwargs from workflow.run() through to agent.run() and @tool functions. WORKFLOW_RUN_KWARGS_KEY = "_workflow_run_kwargs" +# State key for the AgentSession identity supplied through WorkflowAgent. +WORKFLOW_AGENT_SESSION_ID_KEY = "_workflow_agent_session_id" + +# State key identifying WorkflowAgent turns, including turns without an AgentSession. +WORKFLOW_AGENT_RUN_KEY = "_workflow_agent_run" + +# Checkpoint-backed workflow scope and the binding-specific scope inherited by a child. +WORKFLOW_SESSION_ID_KEY = "_workflow_session_id" +WORKFLOW_PARENT_SESSION_ID_KEY = "_workflow_parent_session_id" +WORKFLOW_AGENT_SESSION_NAMESPACE_KEY = "_workflow_agent_session_namespace" +WORKFLOW_REQUEST_PARENT_CONTEXTS_KEY = "_workflow_request_parent_contexts" +WORKFLOW_RESPONSE_SCOPE_KEY = "_workflow_response_scope" + # Key used to store collision-free, executor-aware workflow invocation kwargs. # WORKFLOW_RUN_KWARGS_KEY remains in the legacy dict format for compatibility with # older first-party packages and checkpoints. diff --git a/python/packages/core/agent_framework/_workflows/_edge_runner.py b/python/packages/core/agent_framework/_workflows/_edge_runner.py index 89d2ecf3a6b..8824ff10bdf 100644 --- a/python/packages/core/agent_framework/_workflows/_edge_runner.py +++ b/python/packages/core/agent_framework/_workflows/_edge_runner.py @@ -1,5 +1,7 @@ # Copyright (c) Microsoft. All rights reserved. +from __future__ import annotations + import asyncio import json import logging @@ -65,7 +67,6 @@ def __init__(self, edge_group: EdgeGroup, executors: dict[str, Executor]) -> Non self._edge_group = edge_group self._executors = executors - @abstractmethod async def send_message( self, message: WorkflowMessage, @@ -83,8 +84,37 @@ async def send_message( bool: True if the message was processed successfully, False if the target executor cannot handle the message. """ + return await self._send_selected_message(message, state, ctx, await self._select_targets(message)) + + async def _send_selected_message( + self, + message: WorkflowMessage, + state: State, + ctx: RunnerContext, + delivery: tuple[bool, list[str]], + ) -> bool: + processed, target_ids = delivery + source_id = self._edge_group.edges[0].source_id + if len(target_ids) == 1: + await self._execute_on_target(target_ids[0], [source_id], message, state, ctx) + return processed + await gather_cancelling_siblings_on_error( + *(self._execute_on_target(target_id, [source_id], message, state, ctx) for target_id in target_ids) + ) + return processed + + @abstractmethod + async def _select_targets(self, message: WorkflowMessage) -> tuple[bool, list[str]]: + """Determine delivery targets without executing or buffering the message.""" raise NotImplementedError + def _has_compatible_target(self, message: WorkflowMessage) -> bool: + """Check target and handler compatibility without evaluating routing predicates.""" + return any( + (not message.target_id or message.target_id == edge.target_id) and self._can_handle(edge.target_id, message) + for edge in self._edge_group.edges + ) + @property def state_key(self) -> str: """Return a topology-derived key identifying this runner across workflow instances. @@ -159,13 +189,8 @@ def __init__(self, edge_group: SingleEdgeGroup | InternalEdgeGroup, executors: d super().__init__(edge_group, executors) self._edge = edge_group.edges[0] - async def send_message( - self, - message: WorkflowMessage, - state: State, - ctx: RunnerContext, - ) -> bool: - """Send a message through the single edge.""" + async def _select_targets(self, message: WorkflowMessage) -> tuple[bool, list[str]]: + """Select the single edge's delivery target without executing it.""" should_execute = False target_id: str | None = None source_id: str | None = None @@ -183,7 +208,7 @@ async def send_message( OtelAttr.EDGE_GROUP_DELIVERED: False, OtelAttr.EDGE_GROUP_DELIVERY_STATUS: EdgeGroupDeliveryStatus.DROPPED_TARGET_MISMATCH.value, }) - return False + return False, [] if self._can_handle(self._edge.target_id, message): route_result = await self._edge.should_route(message.data) @@ -202,13 +227,13 @@ async def send_message( OtelAttr.EDGE_GROUP_DELIVERY_STATUS: EdgeGroupDeliveryStatus.DROPPED_CONDITION_FALSE.value, }) # Return True here because message was processed, just condition failed - return True + return True, [] else: span.set_attributes({ OtelAttr.EDGE_GROUP_DELIVERED: False, OtelAttr.EDGE_GROUP_DELIVERY_STATUS: EdgeGroupDeliveryStatus.DROPPED_TYPE_MISMATCH.value, }) - return False + return False, [] except Exception as e: span.set_attributes({ OtelAttr.EDGE_GROUP_DELIVERED: False, @@ -216,12 +241,10 @@ async def send_message( }) raise e - # Execute outside the span if should_execute and target_id and source_id: - await self._execute_on_target(target_id, [source_id], message, state, ctx) - return True + return True, [target_id] - return False + return False, [] class FanOutEdgeRunner(EdgeRunner): @@ -236,13 +259,8 @@ def __init__(self, edge_group: FanOutEdgeGroup, executors: dict[str, Executor]) Callable[[Any, list[str]], list[str]] | None, getattr(edge_group, "selection_func", None) ) - async def send_message( - self, - message: WorkflowMessage, - state: State, - ctx: RunnerContext, - ) -> bool: - """Send a message through all edges in the fan-out edge group.""" + async def _select_targets(self, message: WorkflowMessage) -> tuple[bool, list[str]]: + """Select fan-out delivery targets without executing them.""" deliverable_edges: list[Edge] = [] single_target_edge: Edge | None = None # Process routing logic within span @@ -269,14 +287,14 @@ async def send_message( OtelAttr.EDGE_GROUP_DELIVERED: False, OtelAttr.EDGE_GROUP_DELIVERY_STATUS: EdgeGroupDeliveryStatus.DROPPED_TARGET_MISMATCH.value, }) - return False + return False, [] if not any(self._can_handle(target_id, message) for target_id in self._target_ids): span.set_attributes({ OtelAttr.EDGE_GROUP_DELIVERED: False, OtelAttr.EDGE_GROUP_DELIVERY_STATUS: EdgeGroupDeliveryStatus.DROPPED_TYPE_MISMATCH.value, }) - return False + return False, [] selection_results = ( self._selection_func(message.data, self._target_ids) if self._selection_func else self._target_ids @@ -313,21 +331,21 @@ async def send_message( OtelAttr.EDGE_GROUP_DELIVERY_STATUS: EdgeGroupDeliveryStatus.DROPPED_CONDITION_FALSE.value, # ruff:ignore[line-too-long] }) # For targeted messages with condition failure, return True (message was processed) - return True + return True, [] else: span.set_attributes({ OtelAttr.EDGE_GROUP_DELIVERED: False, OtelAttr.EDGE_GROUP_DELIVERY_STATUS: EdgeGroupDeliveryStatus.DROPPED_TYPE_MISMATCH.value, # ruff:ignore[line-too-long] }) # For targeted messages that can't be handled, return False - return False + return False, [] else: span.set_attributes({ OtelAttr.EDGE_GROUP_DELIVERED: False, OtelAttr.EDGE_GROUP_DELIVERY_STATUS: EdgeGroupDeliveryStatus.DROPPED_TARGET_MISMATCH.value, }) # For targeted messages not in selection, return False - return False + return False, [] else: # If no target ID, send the message to the selected targets for target_id in selection_results: @@ -355,27 +373,14 @@ async def send_message( }) raise e - # Execute outside the span if single_target_edge: - await self._execute_on_target( - single_target_edge.target_id, - [single_target_edge.source_id], - message, - state, - ctx, - ) - return True + return True, [single_target_edge.target_id] if deliverable_edges: - - async def send_to_edge(edge: Edge) -> None: - await self._execute_on_target(edge.target_id, [edge.source_id], message, state, ctx) - - await gather_cancelling_siblings_on_error(*(send_to_edge(edge) for edge in deliverable_edges)) - return True + return True, [edge.target_id for edge in deliverable_edges] # If we get here, it's a broadcast message with no deliverable edges - return False + return False, [] def _validate_selection_result(self, selection_results: list[str]) -> bool: """Validate the selection results to ensure all IDs are valid target executor IDs.""" @@ -394,6 +399,22 @@ def __init__(self, edge_group: FanInEdgeGroup, executors: dict[str, Executor]) - # Key is the source executor ID, value is a list of messages self._buffer: dict[str, list[WorkflowMessage]] = defaultdict(list) + async def _select_targets(self, message: WorkflowMessage) -> tuple[bool, list[str]]: + # Fan-in delivers an aggregate rather than the wrapped reply itself. + return False, [] + + def _has_compatible_target(self, message: WorkflowMessage) -> bool: + return False + + async def _send_selected_message( + self, + message: WorkflowMessage, + state: State, + ctx: RunnerContext, + delivery: tuple[bool, list[str]], + ) -> bool: + return await self.send_message(message, state, ctx) + def snapshot_state(self) -> dict[str, Any] | None: """Capture the buffered messages that are still waiting for the remaining sources. diff --git a/python/packages/core/agent_framework/_workflows/_executor.py b/python/packages/core/agent_framework/_workflows/_executor.py index 7c1a6ffd524..6971d6c3c4b 100644 --- a/python/packages/core/agent_framework/_workflows/_executor.py +++ b/python/packages/core/agent_framework/_workflows/_executor.py @@ -1,5 +1,7 @@ # Copyright (c) Microsoft. All rights reserved. +from __future__ import annotations + import asyncio import contextlib import copy @@ -7,7 +9,7 @@ import inspect import logging import types -from collections.abc import Awaitable, Callable +from collections.abc import Awaitable, Callable, Sequence from typing import Any, TypeVar, overload from ..observability import create_processing_span @@ -17,6 +19,7 @@ _framework_event_origin, # type: ignore[reportPrivateUsage] ) from ._model_utils import DictConvertible +from ._request_context import restore_request_context from ._request_info_mixin import RequestInfoMixin from ._runner_context import MessageType, RunnerContext, WorkflowMessage from ._state import State @@ -31,6 +34,10 @@ logger = logging.getLogger(__name__) +# Metadata migrations are applied only after every validator in the tree succeeds. +_ResponseBatchValidator = Callable[[Sequence[WorkflowEvent]], Callable[[], None] | None] +_ResponseBatchValidation = tuple[_ResponseBatchValidator, Sequence[WorkflowEvent]] + # region Executor class Executor(RequestInfoMixin, DictConvertible): @@ -309,6 +316,8 @@ async def execute( invoke_event = WorkflowEvent.executor_invoked(self.id, copy.deepcopy(message)) await context.add_event(invoke_event) try: + if context.request_id is not None: + restore_request_context(state, context.request_id) await handler(message, context) except Exception as exc: # Surface structured executor failure before propagating @@ -533,6 +542,17 @@ async def _cancel_pending_request(self, request_id: str, ctx: WorkflowContext[An """Release executor-owned state for a cancelled workflow request.""" del request_id, ctx + async def _get_response_validation_batches( + self, requests: Sequence[WorkflowEvent] + ) -> list[_ResponseBatchValidation]: + """Collect delegated response policies without changing pending requests.""" + del requests + return [] + + async def _validate_checkpoint_restore(self, state: dict[str, Any], requests: Sequence[WorkflowEvent] = ()) -> None: + """Validate a checkpoint snapshot without changing live executor state.""" + del state, requests + async def on_checkpoint_save(self) -> dict[str, Any]: """Hook called when the workflow is being saved to a checkpoint. diff --git a/python/packages/core/agent_framework/_workflows/_request_context.py b/python/packages/core/agent_framework/_workflows/_request_context.py new file mode 100644 index 00000000000..5cc2af6852c --- /dev/null +++ b/python/packages/core/agent_framework/_workflows/_request_context.py @@ -0,0 +1,134 @@ +# Copyright (c) Microsoft. All rights reserved. + +from __future__ import annotations + +import logging +import uuid +from collections.abc import Collection, Mapping, Sequence +from typing import Any, cast + +from ._const import ( + WORKFLOW_AGENT_RUN_KEY, + WORKFLOW_AGENT_SESSION_ID_KEY, + WORKFLOW_PARENT_SESSION_ID_KEY, + WORKFLOW_REQUEST_PARENT_CONTEXTS_KEY, + WORKFLOW_RESPONSE_SCOPE_KEY, + WORKFLOW_SESSION_ID_KEY, +) +from ._events import WorkflowEvent +from ._state import State + +logger = logging.getLogger(__name__) + +_RUN_CONTEXT_KEYS = ( + WORKFLOW_SESSION_ID_KEY, + WORKFLOW_PARENT_SESSION_ID_KEY, + WORKFLOW_AGENT_RUN_KEY, + WORKFLOW_AGENT_SESSION_ID_KEY, +) +_MISSING_ORIGIN = object() + + +def get_response_scopes(state: State | Mapping[str, Any]) -> set[str]: + """Read the accepted scope retained while a response continuation has in-flight work.""" + scope = state.get(WORKFLOW_RESPONSE_SCOPE_KEY) + if scope is None: + return set() + if not isinstance(scope, str) or not scope: + raise ValueError("Invalid parent workflow response scope.") + return {scope} + + +def request_context_state_key(boundary_id: str | None = None) -> str: + """Keep wrapped child origins independent of external requests with the same ID.""" + if boundary_id is None: + return WORKFLOW_REQUEST_PARENT_CONTEXTS_KEY + return f"{WORKFLOW_REQUEST_PARENT_CONTEXTS_KEY}:{boundary_id}" + + +def _validate_origin(value: Any) -> dict[str, Any]: + if not isinstance(value, dict): + raise ValueError("Invalid parent workflow request context.") + context = cast(dict[str, Any], value) + if set(context) != set(_RUN_CONTEXT_KEYS): + raise ValueError("Invalid parent workflow request context.") + scope = context[WORKFLOW_SESSION_ID_KEY] + inherited_scope = context[WORKFLOW_PARENT_SESSION_ID_KEY] + agent_session = context[WORKFLOW_AGENT_SESSION_ID_KEY] + if ( + not isinstance(scope, str) + or not scope + or (inherited_scope is not None and (not isinstance(inherited_scope, str) or not inherited_scope)) + or not isinstance(context[WORKFLOW_AGENT_RUN_KEY], bool) + or (agent_session is not None and not isinstance(agent_session, str)) + ): + raise ValueError("Invalid parent workflow request context.") + return context + + +def validate_request_contexts( + contexts: Any, + requests: Sequence[WorkflowEvent] = (), + *, + boundary_contexts: Sequence[tuple[Any, Sequence[WorkflowEvent]]] = (), + previous_scopes: Collection[str] = (), +) -> set[str]: + """Validate selected origins together with earlier deliveries in the same continuation.""" + scopes = set(previous_scopes) + for pending_contexts, pending_requests in ((contexts, requests), *boundary_contexts): + if not isinstance(pending_contexts, dict): + raise ValueError("Invalid parent workflow request contexts.") + pending_contexts = cast(dict[str, Any], pending_contexts) + for request_id, value in pending_contexts.items(): + if not isinstance(request_id, str): + raise ValueError("Invalid parent workflow request context.") + _validate_origin(value) + scopes.update( + pending_contexts[request.request_id][WORKFLOW_SESSION_ID_KEY] + for request in pending_requests + if request.request_id in pending_contexts + ) + if len(scopes) > 1: + raise ValueError("Responses span multiple parent workflow scopes; submit each scope in a separate run.") + return scopes + + +def capture_request_context(state: State, request_id: str, *, boundary_id: str | None = None) -> None: + """Associate a request with framework-owned run metadata, without copying its payload.""" + state_key = request_context_state_key(boundary_id) + existing = state._get_mapping_item(state_key, request_id, _MISSING_ORIGIN) # pyright: ignore[reportPrivateUsage] + if existing is not _MISSING_ORIGIN: + _validate_origin(existing) + return + origin = {key: state.get(key) for key in _RUN_CONTEXT_KEYS} + if origin[WORKFLOW_SESSION_ID_KEY] is None: + logger.warning("Workflow request %s has no active run scope; assigning a new isolated scope.", request_id) + origin[WORKFLOW_SESSION_ID_KEY] = uuid.uuid4().hex + state.set(WORKFLOW_SESSION_ID_KEY, origin[WORKFLOW_SESSION_ID_KEY]) + origin[WORKFLOW_AGENT_RUN_KEY] = origin[WORKFLOW_AGENT_RUN_KEY] is True + _validate_origin(origin) + state._set_mapping_item(state_key, request_id, origin) # pyright: ignore[reportPrivateUsage] + + +def restore_request_context(state: State, request_id: str, *, boundary_id: str | None = None) -> None: + """Restore and release a request's run metadata before its continuation executes.""" + state_key = request_context_state_key(boundary_id) + if boundary_id is not None and not state.has(state_key): + # Checkpoints predating separate boundary ownership used the external-request map. + state_key = WORKFLOW_REQUEST_PARENT_CONTEXTS_KEY + origin = state._get_mapping_item(state_key, request_id, _MISSING_ORIGIN) # pyright: ignore[reportPrivateUsage] + if origin is not _MISSING_ORIGIN: + _validate_origin(origin) + state._pop_mapping_item(state_key, request_id) # pyright: ignore[reportPrivateUsage] + else: + logger.warning( + "Workflow request %s has no checkpointed run origin; assigning a new isolated scope.", request_id + ) + origin = { + WORKFLOW_SESSION_ID_KEY: uuid.uuid4().hex, + WORKFLOW_PARENT_SESSION_ID_KEY: None, + WORKFLOW_AGENT_RUN_KEY: False, + WORKFLOW_AGENT_SESSION_ID_KEY: None, + } + for key in _RUN_CONTEXT_KEYS: + state.set(key, origin[key]) diff --git a/python/packages/core/agent_framework/_workflows/_runner.py b/python/packages/core/agent_framework/_workflows/_runner.py index f41183cb1c2..86614d1db9b 100644 --- a/python/packages/core/agent_framework/_workflows/_runner.py +++ b/python/packages/core/agent_framework/_workflows/_runner.py @@ -1,10 +1,13 @@ # Copyright (c) Microsoft. All rights reserved. +from __future__ import annotations + import asyncio +import copy import logging import warnings from collections import defaultdict -from collections.abc import AsyncGenerator, Sequence +from collections.abc import AsyncGenerator, Awaitable, Callable, Mapping, Sequence from typing import TYPE_CHECKING, Any, cast from ..exceptions import ( @@ -12,14 +15,17 @@ WorkflowConvergenceException, ) from ._checkpoint import CheckpointID, CheckpointStorage, WorkflowCheckpoint -from ._const import EDGE_STATE_KEY, EXECUTOR_STATE_KEY +from ._const import EDGE_STATE_KEY, EXECUTOR_STATE_KEY, WORKFLOW_RESPONSE_SCOPE_KEY from ._edge import EdgeGroup from ._edge_runner import EdgeRunner, create_edge_runner, gather_cancelling_siblings_on_error from ._events import WorkflowEvent from ._executor import Executor +from ._request_context import get_response_scopes, request_context_state_key, validate_request_contexts from ._runner_context import ( RunnerContext, WorkflowMessage, + _get_queued_response_info_events, # pyright: ignore[reportPrivateUsage] + _SupportsQueuedMessages, # pyright: ignore[reportPrivateUsage] ) from ._state import State @@ -83,6 +89,11 @@ def __init__( # Checkpointing related attributes self._previous_checkpoint_id: CheckpointID | None = None + self._response_batch_validator: Callable[[Sequence[WorkflowEvent]], None] | None = None + self._queued_response_validator: Callable[[set[str]], Awaitable[None]] | None = None + self._wrapped_response_validator: ( + Callable[[Mapping[str, Sequence[WorkflowEvent]], set[str]], Awaitable[None]] | None + ) = None @property def context(self) -> RunnerContext: @@ -105,7 +116,7 @@ def reset_iteration_count(self) -> None: """ self._iteration = 0 - async def run_until_convergence(self) -> AsyncGenerator[WorkflowEvent, None]: + async def run_until_convergence(self) -> AsyncGenerator[WorkflowEvent]: """Run the workflow until no more messages are sent. The runner's checkpoint responsibility is intentionally narrow: it creates a @@ -163,6 +174,9 @@ async def run_until_convergence(self) -> AsyncGenerator[WorkflowEvent, None]: logger.info(f"Completed superstep {self._iteration}") + # Idle checkpoints must allow a later response run to select a different origin. + if not await self._ctx.has_messages() and self._state.has(WORKFLOW_RESPONSE_SCOPE_KEY): + self._state.delete(WORKFLOW_RESPONSE_SCOPE_KEY) # Commit pending state changes at superstep boundary self._state.commit() committed = True @@ -204,38 +218,106 @@ async def _run_iteration(self) -> None: a source by specifying the target, the messages will be delivered to the specified targets in the order they were sent. This is because all messages go through the same edge runner instance which preserves message order. + - Wrapped replies wait for preceding messages only on compatible edge runners, then select and + validate all targets before delivery. Unrelated edge runners remain concurrent. """ + from ._workflow_executor import SubWorkflowResponseMessage + + scopes = get_response_scopes(self._state) + next_message: dict[str, int] = {} async def _deliver_messages(source_executor_id: str, source_messages: list[WorkflowMessage]) -> None: """Outer loop to concurrently deliver messages from all sources to their targets.""" - - async def _deliver_message_inner(edge_runner: EdgeRunner, message: WorkflowMessage) -> bool: - """Inner loop to deliver a single message through an edge runner.""" - return await edge_runner.send_message(message, self._state, self._ctx) - - async def _deliver_messages_for_edge_runner(edge_runner: EdgeRunner) -> None: - # Preserve message order per edge runner (and therefore per routed target path) - # while still allowing parallelism across different edge runners. - for message in source_messages: - await _deliver_message_inner(edge_runner, message) - # Route all messages through normal workflow edges associated_edge_runners = self._edge_runner_map.get(source_executor_id, []) if not associated_edge_runners: # This is expected for terminal nodes (e.g., EndWorkflow, last action in workflow) logger.debug(f"No outgoing edges found for executor {source_executor_id}; dropping messages.") + next_message[source_executor_id] = len(source_messages) return - tasks = [_deliver_messages_for_edge_runner(edge_runner) for edge_runner in associated_edge_runners] - await gather_cancelling_siblings_on_error(*tasks) + selections: dict[int, dict[EdgeRunner, asyncio.Future[tuple[bool, list[str]]]]] = {} + validated: dict[int, asyncio.Event] = {} + for index, message in enumerate(source_messages): + if isinstance(message.data, SubWorkflowResponseMessage): + selections[index] = { + runner: asyncio.get_running_loop().create_future() + for runner in associated_edge_runners + if runner._has_compatible_target(message) # pyright: ignore[reportPrivateUsage] + } + validated[index] = asyncio.Event() + + async def _deliver_messages_for_edge_runner(edge_runner: EdgeRunner) -> None: + for index, message in enumerate(source_messages): + if index not in selections or edge_runner not in selections[index]: + await edge_runner.send_message(message, self._state, self._ctx) + continue + # Select routes after this edge's earlier handlers, not unrelated outgoing handlers. + delivery = await edge_runner._select_targets(message) # pyright: ignore[reportPrivateUsage] + selections[index][edge_runner].set_result(delivery) + await validated[index].wait() + await edge_runner._send_selected_message( # pyright: ignore[reportPrivateUsage] + message, self._state, self._ctx, delivery + ) + + async def _validate_reply(index: int) -> None: + targets: set[str] = set() + for selection in selections[index].values(): + _, target_ids = await selection + targets.update(target_ids) + message = source_messages[index] + if self._wrapped_response_validator is not None: + await self._wrapped_response_validator( + {target: [message.data.source_event] for target in targets}, scopes + ) + validated[index].set() + # Independent paths may validate out of order; retain the earliest unvalidated reply on failure. + while next_message[source_executor_id] < len(source_messages): + pending_index = next_message[source_executor_id] + if pending_index in validated and not validated[pending_index].is_set(): + break + next_message[source_executor_id] += 1 + + next_message[source_executor_id] = next(iter(selections), len(source_messages)) + await gather_cancelling_siblings_on_error( + *(_validate_reply(index) for index in selections), + *(_deliver_messages_for_edge_runner(runner) for runner in associated_edge_runners), + ) + if self._queued_response_validator is not None: + await self._queued_response_validator(scopes) message_batches = await self._ctx.drain_messages() - await gather_cancelling_siblings_on_error( - *( - _deliver_messages(source_executor_id, source_messages) - for source_executor_id, source_messages in message_batches.items() + try: + await gather_cancelling_siblings_on_error( + *( + _deliver_messages(source_executor_id, source_messages) + for source_executor_id, source_messages in message_batches.items() + ) ) + except BaseException: + remaining = [ + message + for source_id, messages in message_batches.items() + for message in messages[next_message.get(source_id, 0) :] + ] + if remaining: + # Siblings have stopped; retain undelivered replies ahead of newly emitted messages. + queued = await self._ctx.drain_messages() + for message in remaining: + await self._ctx.send_message(message) + for messages in queued.values(): + for message in messages: + await self._ctx.send_message(message) + raise + + async def _get_queued_messages(self) -> Mapping[str, Sequence[WorkflowMessage]]: + if isinstance(self._ctx, _SupportsQueuedMessages): + return await self._ctx.get_queued_messages() + # Legacy contexts expose their queues through non-persisting checkpoint snapshots. + checkpoint = await self._ctx.build_checkpoint( + self._workflow_name, self._graph_signature_hash, self._state, None, self._iteration ) + return checkpoint.messages async def _prepare_checkpoint_state(self) -> None: """Persist executor snapshots into committed shared state. @@ -319,12 +401,8 @@ async def restore_from_checkpoint( logger.error(f"Checkpoint {checkpoint_id} not found") raise WorkflowCheckpointException(f"Checkpoint {checkpoint_id} not found") - # Validate the loaded checkpoint against the workflow - if self._graph_signature_hash != checkpoint.graph_signature_hash: - raise WorkflowCheckpointException( - "Workflow graph has changed since the checkpoint was created. " - "Please rebuild the original workflow before resuming." - ) + checkpoint = copy.deepcopy(checkpoint) + await self._validate_checkpoint_tree(checkpoint) # Restore state. Clear first so import_state (which merges) does # not leak stale keys from a prior run on this Workflow instance. @@ -382,13 +460,10 @@ async def restore_checkpoint(self, checkpoint: WorkflowCheckpoint) -> None: WorkflowCheckpointException: If the checkpoint's graph signature does not match this runner's workflow, or if restoration otherwise fails. """ - if self._graph_signature_hash != checkpoint.graph_signature_hash: - raise WorkflowCheckpointException( - "Workflow graph has changed since the checkpoint was created. " - "Please rebuild the original workflow before resuming." - ) - try: + checkpoint = copy.deepcopy(checkpoint) + await self._validate_checkpoint_tree(checkpoint) + # Clear first so import_state (which merges) does not leak stale keys from a # prior run on this Workflow instance. self._state.clear() @@ -397,10 +472,55 @@ async def restore_checkpoint(self, checkpoint: WorkflowCheckpoint) -> None: self._restore_edge_runner_states() await self._ctx.apply_checkpoint(checkpoint) self._mark_resumed(checkpoint) + except WorkflowCheckpointException: + raise except Exception as e: logger.error(f"Failed to restore from checkpoint {checkpoint.checkpoint_id}: {e}") raise WorkflowCheckpointException(f"Failed to restore from checkpoint {checkpoint.checkpoint_id}") from e + async def _validate_checkpoint_tree( + self, checkpoint: WorkflowCheckpoint, delegated_requests: Sequence[WorkflowEvent] = () + ) -> None: + from ._const import WORKFLOW_REQUEST_PARENT_CONTEXTS_KEY + + if self._graph_signature_hash != checkpoint.graph_signature_hash: + raise WorkflowCheckpointException( + "Workflow graph has changed since the checkpoint was created. " + "Please rebuild the original workflow before resuming." + ) + requests = [ + *_get_queued_response_info_events(checkpoint.messages), + *( + checkpoint.pending_request_info_events[request.request_id] + for request in delegated_requests + if request.request_id in checkpoint.pending_request_info_events + ), + ] + if requests and self._response_batch_validator is not None: + self._response_batch_validator(requests) + parent_contexts = checkpoint.state.get(WORKFLOW_REQUEST_PARENT_CONTEXTS_KEY, {}) + validate_request_contexts( + parent_contexts, + requests, + previous_scopes=get_response_scopes(checkpoint.state), + boundary_contexts=[ + ( + checkpoint.state.get(request_context_state_key(executor_id), parent_contexts), + (), + ) + for executor_id in self._executors + if request_context_state_key(executor_id) in checkpoint.state + ], + ) + if EXECUTOR_STATE_KEY in checkpoint.state: + delegated_by_executor: defaultdict[str, list[WorkflowEvent]] = defaultdict(list) + for request in requests: + delegated_by_executor[request.source_executor_id].append(request) + for _, executor, state in self._get_executor_states(checkpoint.state[EXECUTOR_STATE_KEY]): + await executor._validate_checkpoint_restore( # pyright: ignore[reportPrivateUsage] + state, delegated_by_executor[executor.id] + ) + async def _save_executor_states(self) -> None: """Populate executor state by calling checkpoint hooks on executors.""" for exec_id, executor in self._executors.items(): @@ -420,14 +540,20 @@ async def _save_executor_states(self) -> None: async def _restore_executor_states(self) -> None: """Restore executor state by calling restore hooks on executors.""" - has_executor_states = self._state.has(EXECUTOR_STATE_KEY) - if not has_executor_states: + if not self._state.has(EXECUTOR_STATE_KEY): return - executor_states = self._state.get(EXECUTOR_STATE_KEY) + for executor_id, executor, state in self._get_executor_states(self._state.get(EXECUTOR_STATE_KEY)): + try: + await executor.on_checkpoint_restore(state) + except Exception as ex: # pragma: no cover - defensive + raise WorkflowCheckpointException(f"Executor {executor_id} on_checkpoint_restore failed") from ex + + def _get_executor_states(self, executor_states: Any) -> list[tuple[str, Executor, dict[str, Any]]]: if not isinstance(executor_states, dict): raise WorkflowCheckpointException("Executor states in shared state is not a dictionary. Unable to restore.") + snapshots: list[tuple[str, Executor, dict[str, Any]]] = [] for executor_id, state in executor_states.items(): # pyright: ignore[reportUnknownVariableType] if not isinstance(executor_id, str): raise WorkflowCheckpointException("Executor ID in executor states is not a string. Unable to restore.") @@ -439,12 +565,8 @@ async def _restore_executor_states(self) -> None: executor = self._executors.get(executor_id) if not executor: raise WorkflowCheckpointException(f"Executor {executor_id} not found during state restoration.") - - # Try the updated behavior only if backward compatibility did not restore - try: - await executor.on_checkpoint_restore(state) # pyright: ignore[reportUnknownArgumentType] - except Exception as ex: # pragma: no cover - defensive - raise WorkflowCheckpointException(f"Executor {executor_id} on_checkpoint_restore failed") from ex + snapshots.append((executor_id, executor, cast(dict[str, Any], state))) + return snapshots def _save_edge_runner_states(self) -> None: """Store edge runner delivery state, such as partially filled fan-in buffers, in state. diff --git a/python/packages/core/agent_framework/_workflows/_runner_context.py b/python/packages/core/agent_framework/_workflows/_runner_context.py index edbbfea50c9..71bc526d77c 100644 --- a/python/packages/core/agent_framework/_workflows/_runner_context.py +++ b/python/packages/core/agent_framework/_workflows/_runner_context.py @@ -5,7 +5,7 @@ import asyncio import logging import warnings -from collections.abc import Callable +from collections.abc import Callable, Mapping, Sequence from copy import copy from dataclasses import dataclass from enum import Enum @@ -96,6 +96,24 @@ def from_dict(data: dict[str, Any]) -> WorkflowMessage: ) +def _get_queued_response_info_events(messages: Mapping[str, Sequence[WorkflowMessage]]) -> list[WorkflowEvent[Any]]: + return [ + message.original_request_info_event + for batch in messages.values() + for message in batch + if message.type == MessageType.RESPONSE and message.original_request_info_event is not None + ] + + +@runtime_checkable +class _SupportsQueuedMessages(Protocol): + """Optional fast path for inspecting queued messages without building a snapshot.""" + + async def get_queued_messages(self) -> Mapping[str, Sequence[WorkflowMessage]]: + """Inspect queued messages without consuming them or copying their payloads.""" + ... + + @runtime_checkable class RunnerContext(Protocol): """Protocol for the execution context used by the runner. @@ -611,6 +629,14 @@ async def get_pending_request_info_events(self) -> dict[str, WorkflowEvent[Any]] """ return dict(self._pending_request_info_events) + async def get_queued_response_info_events(self) -> list[WorkflowEvent[Any]]: + """Get original requests for queued responses without consuming messages.""" + return _get_queued_response_info_events(self._messages) + + async def get_queued_messages(self) -> dict[str, list[WorkflowMessage]]: + """Inspect queued messages without consuming them or copying their payloads.""" + return {source_id: list(messages) for source_id, messages in self._messages.items()} + async def cancel_request_info_events(self, request_ids: set[str]) -> dict[str, WorkflowEvent[Any]]: """Remove and return pending request_info events selected for cancellation.""" return { diff --git a/python/packages/core/agent_framework/_workflows/_state.py b/python/packages/core/agent_framework/_workflows/_state.py index 136b70de14e..288812f5073 100644 --- a/python/packages/core/agent_framework/_workflows/_state.py +++ b/python/packages/core/agent_framework/_workflows/_state.py @@ -1,8 +1,10 @@ # Copyright (c) Microsoft. All rights reserved. +from __future__ import annotations + import copy from collections.abc import Callable -from typing import Any +from typing import Any, cast class State: @@ -28,6 +30,7 @@ def __init__(self) -> None: """Initialize the state.""" self._committed: dict[str, Any] = {} self._pending: dict[str, Any] = {} + self._pending_mapping_updates: dict[str, dict[str, Any]] = {} def set(self, key: str, value: Any) -> None: """Set a value in the pending state buffer. @@ -44,6 +47,7 @@ def set(self, key: str, value: Any) -> None: in a superstep see the same committed state at the start. """ self._pending[key] = copy.deepcopy(value) + self._pending_mapping_updates.pop(key, None) def get(self, key: str, default: Any = None) -> Any: """Get a value from state, checking pending first then committed. @@ -56,14 +60,51 @@ def get(self, key: str, default: Any = None) -> Any: A deep copy of the value if found, otherwise the default value. Mutate the returned value and pass it to :meth:`set` to update workflow state. """ - if key in self._pending: - value = self._pending[key] - if value is _DeleteSentinel: - return default - return copy.deepcopy(value) - if key in self._committed: - return copy.deepcopy(self._committed[key]) - return default + value = self._get_value(key) + return default if value is _DeleteSentinel else copy.deepcopy(value) + + def _get_value(self, key: str) -> Any: + value = self._pending.get(key, self._committed.get(key, _DeleteSentinel)) + updates = self._pending_mapping_updates.get(key) + if updates: + merged: dict[str, Any] = dict(value) + for item_key, item_value in updates.items(): + if item_value is _DeleteSentinel: + merged.pop(item_key, None) + else: + merged[item_key] = item_value + return merged + return value + + def _get_mapping(self, key: str) -> dict[str, Any]: + value = self._pending.get(key, self._committed.get(key, _DeleteSentinel)) + if value is _DeleteSentinel: + return {} + if not isinstance(value, dict): + raise ValueError(f"Workflow state '{key}' must be a mapping.") + return cast(dict[str, Any], value) + + def _get_mapping_item(self, key: str, item_key: str, default: Any = None) -> Any: + """Read one framework-owned mapping entry without copying its siblings.""" + mapping = self._get_mapping(key) + value = self._pending_mapping_updates.get(key, {}).get(item_key, mapping.get(item_key, _DeleteSentinel)) + return default if value is _DeleteSentinel else copy.deepcopy(value) + + def _set_mapping_item(self, key: str, item_key: str, value: Any) -> None: + """Stage a detached mapping entry while preserving commit/discard semantics.""" + self._get_mapping(key) + copied_value = copy.deepcopy(value) + if not self.has(key): + self._pending[key] = {} + self._pending_mapping_updates.setdefault(key, {})[item_key] = copied_value + + def _pop_mapping_item(self, key: str, item_key: str) -> Any: + """Stage removal of one entry without copying or mutating committed siblings.""" + value = self._get_mapping_item(key, item_key, _DeleteSentinel) + if value is not _DeleteSentinel: + self._pending_mapping_updates.setdefault(key, {})[item_key] = _DeleteSentinel + return value + return None def has(self, key: str) -> bool: """Check if a key exists in pending or committed state.""" @@ -77,7 +118,7 @@ def _validate(self, key: str, validator: Callable[[Any], None]) -> None: The internal validator must not mutate or retain the value. Missing and pending-deleted keys are not validated, matching ``get`` visibility. """ - value = self._pending.get(key, self._committed.get(key, _DeleteSentinel)) + value = self._get_value(key) if value is not _DeleteSentinel: validator(value) @@ -97,11 +138,13 @@ def delete(self, key: str) -> None: elif key in self._pending: # Only exists in pending, safe to just remove del self._pending[key] + self._pending_mapping_updates.pop(key, None) def clear(self) -> None: """Clear both committed and pending state.""" self._committed.clear() self._pending.clear() + self._pending_mapping_updates.clear() def commit(self) -> None: """Commit pending changes to the committed state. @@ -113,11 +156,20 @@ def commit(self) -> None: self._committed.pop(key, None) else: self._committed[key] = value + for key, updates in self._pending_mapping_updates.items(): + mapping = self._committed[key] + for item_key, value in updates.items(): + if value is _DeleteSentinel: + mapping.pop(item_key, None) + else: + mapping[item_key] = value self._pending.clear() + self._pending_mapping_updates.clear() def discard(self) -> None: """Discard all pending changes without committing.""" self._pending.clear() + self._pending_mapping_updates.clear() def export_state(self) -> dict[str, Any]: """Export a deepcopy of the committed state. diff --git a/python/packages/core/agent_framework/_workflows/_workflow.py b/python/packages/core/agent_framework/_workflows/_workflow.py index c60617dabd5..95be9656356 100644 --- a/python/packages/core/agent_framework/_workflows/_workflow.py +++ b/python/packages/core/agent_framework/_workflows/_workflow.py @@ -35,7 +35,14 @@ RAW_FUNCTION_INVOCATION_KWARGS_KEY, RESOLVED_WORKFLOW_RUN_KWARGS_KEY, ROUTED_WORKFLOW_RUN_KWARGS_KEY, + WORKFLOW_AGENT_RUN_KEY, + WORKFLOW_AGENT_SESSION_ID_KEY, + WORKFLOW_AGENT_SESSION_NAMESPACE_KEY, + WORKFLOW_PARENT_SESSION_ID_KEY, + WORKFLOW_REQUEST_PARENT_CONTEXTS_KEY, + WORKFLOW_RESPONSE_SCOPE_KEY, WORKFLOW_RUN_KWARGS_KEY, + WORKFLOW_SESSION_ID_KEY, ) from ._edge import ( EdgeGroup, @@ -48,10 +55,24 @@ WorkflowRunState, _framework_event_origin, # type: ignore ) -from ._executor import Executor +from ._executor import ( + Executor, + _ResponseBatchValidation, # pyright: ignore[reportPrivateUsage] + _ResponseBatchValidator, # pyright: ignore[reportPrivateUsage] +) from ._model_utils import DictConvertible +from ._request_context import ( + get_response_scopes, + request_context_state_key, + restore_request_context, + validate_request_contexts, +) from ._runner import RunnerImpl -from ._runner_context import RunnerContext, WorkflowMessage +from ._runner_context import ( + RunnerContext, + WorkflowMessage, + _get_queued_response_info_events, # pyright: ignore[reportPrivateUsage] +) from ._state import State from ._typing_utils import is_instance_of, try_coerce_to_type from ._validation import ValidationTypeEnum, WorkflowValidationError @@ -411,6 +432,15 @@ def __init__( # ever iterating, the weakref dereferences to ``None`` once Python collects it, # so a subsequent ``run()`` is allowed. self._active_run: weakref.ref[ResponseStream[WorkflowEvent, WorkflowRunResult]] | None = None + self._response_batch_validator: _ResponseBatchValidator | None = None + self._runner._response_batch_validator = self._validate_response_batch # pyright: ignore[reportPrivateUsage] + self._runner._queued_response_validator = ( # pyright: ignore[reportPrivateUsage] + lambda scopes: self._validate_response_requests((), scopes=scopes) + ) + self._runner._wrapped_response_validator = ( # pyright: ignore[reportPrivateUsage] + lambda wrapped, scopes: self._validate_response_requests((), wrapped_requests=wrapped, scopes=scopes) + ) + self._initial_state_initializer: Callable[[State], None] | None = None # Run-scoped pause checkpoint bookkeeping (owned by Workflow, not callers). # Captured at the start of each ``_run_core`` so ``resolve_pause_checkpoint_id`` @@ -429,6 +459,44 @@ def status(self) -> WorkflowRunState: """ return self._status + def _set_agent_run_context(self, session_id: str | None) -> None: + self._runner.state.set(WORKFLOW_AGENT_RUN_KEY, True) + self._runner.state.set(WORKFLOW_AGENT_SESSION_ID_KEY, session_id) + + def _initialize_workflow_session_id(self) -> None: + state = self._runner.state + scope = state.get(WORKFLOW_PARENT_SESSION_ID_KEY) + if scope is None: + if state.get(WORKFLOW_AGENT_RUN_KEY) is True: + namespace = state.get(WORKFLOW_AGENT_SESSION_NAMESPACE_KEY) + if namespace is None: + namespace = uuid.uuid4().hex + state.set(WORKFLOW_AGENT_SESSION_NAMESPACE_KEY, namespace) + if not isinstance(namespace, str): + raise ValueError("Invalid workflow agent session namespace.") + scope = uuid.uuid5(uuid.UUID(namespace), json.dumps([state.get(WORKFLOW_AGENT_SESSION_ID_KEY)])).hex + else: + scope = uuid.uuid4().hex + if not isinstance(scope, str) or not scope: + raise ValueError("Invalid workflow session scope.") + state.set(WORKFLOW_SESSION_ID_KEY, scope) + + def _set_parent_run_context(self, parent_state: State, binding_id: str) -> None: + parent_scope = parent_state.get(WORKFLOW_SESSION_ID_KEY) + if parent_scope is None: + logger.warning("Parent workflow has no checkpointed session scope; assigning a new isolated scope.") + parent_scope = uuid.uuid4().hex + parent_state.set(WORKFLOW_SESSION_ID_KEY, parent_scope) + if not isinstance(parent_scope, str) or not parent_scope: + raise ValueError("Invalid parent workflow session scope.") + state = self._runner.state + state.set( + WORKFLOW_PARENT_SESSION_ID_KEY, + uuid.uuid5(uuid.NAMESPACE_URL, json.dumps([parent_scope, binding_id])).hex, + ) + state.set(WORKFLOW_AGENT_RUN_KEY, parent_state.get(WORKFLOW_AGENT_RUN_KEY) is True) + state.set(WORKFLOW_AGENT_SESSION_ID_KEY, parent_state.get(WORKFLOW_AGENT_SESSION_ID_KEY)) + def to_dict(self) -> dict[str, Any]: """Serialize the workflow definition into a JSON-ready dictionary.""" data: dict[str, Any] = { @@ -761,6 +829,10 @@ async def _execute_with_message_or_checkpoint( f"Expected input types: {start_executor.input_types}." ) + self._initialize_workflow_session_id() + if self._initial_state_initializer is not None: + self._initial_state_initializer(self._runner.state) + await self._runner.context.send_message(entry_message) # Record the entry checkpoint capturing the seeded input before any # executor runs. The runner only checkpoints after each superstep. @@ -871,6 +943,11 @@ def run( "Workflow is already running; concurrent runs are not allowed on the same instance." ) + if message is not None: + self._runner.state.set(WORKFLOW_AGENT_RUN_KEY, False) + self._runner.state.set(WORKFLOW_AGENT_SESSION_ID_KEY, None) + self._runner.state.set(WORKFLOW_PARENT_SESSION_ID_KEY, None) + # No run is active, so any runtime checkpoint storage override still set on the # context is stale - left over from a prior run whose stream was dropped before # its async-generator finalizer ran. Clear it so this run starts clean and does @@ -949,7 +1026,8 @@ async def _run_core( # NOTE: _validate_run_params already enforces that ``message`` is # mutually exclusive with both ``checkpoint_id`` and ``responses``, # so we don't need to re-check those here. - if message is not None and await self._runner.context.has_messages(): + has_messages = await self._runner.context.has_messages() + if message is not None and has_messages: raise RuntimeError( "Cannot start a new run with 'message' while in-flight executor " "messages remain from a prior run. Resume from a checkpoint " @@ -958,6 +1036,9 @@ async def _run_core( "checkpointing; there is no in-process recovery path." ) + if checkpoint_id is None and not has_messages and self._runner.state.has(WORKFLOW_RESPONSE_SCOPE_KEY): + self._runner.state.delete(WORKFLOW_RESPONSE_SCOPE_KEY) + # Warn (but don't block) when a fresh message or a checkpoint restore begins while the # workflow still has pending request_info events from an unfinished request/response # cycle. A fresh ``message`` does NOT drop those pending requests - they remain pending and @@ -1130,6 +1211,81 @@ async def _restore_and_send_responses( await self._runner.restore_from_checkpoint(checkpoint_id, checkpoint_storage) await self._send_responses_internal(responses) + def _validate_response_batch(self, requests: Sequence[WorkflowEvent]) -> None: + if self._response_batch_validator is not None: + migration = self._response_batch_validator(requests) + if migration is not None: + migration() + + async def _validate_response_requests( + self, + requests: Sequence[WorkflowEvent], + *, + wrapped_requests: Mapping[str, Sequence[WorkflowEvent]] | None = None, + scopes: set[str] | None = None, + ) -> None: + validations = await self._get_response_validation_batches( + requests, wrapped_requests=wrapped_requests, scopes=scopes + ) + migrations = [validator(batch) for validator, batch in validations] + for migration in migrations: + if migration is not None: + migration() + + async def _get_response_validation_batches( + self, + requests: Sequence[WorkflowEvent], + *, + wrapped_requests: Mapping[str, Sequence[WorkflowEvent]] | None = None, + scopes: set[str] | None = None, + ) -> list[_ResponseBatchValidation]: + messages = await self._runner._get_queued_messages() # pyright: ignore[reportPrivateUsage] + queued_requests = _get_queued_response_info_events(messages) + if wrapped_requests is None: + wrapped_requests = {} + if scopes is None: + scopes = get_response_scopes(self._runner.state) + batch = [*queued_requests, *requests] + parent_contexts = self._runner.state.get(WORKFLOW_REQUEST_PARENT_CONTEXTS_KEY, {}) + boundary_contexts = [ + ( + self._runner.state.get(request_context_state_key(executor_id), parent_contexts), + wrapped_requests.get(executor_id, []), + ) + for executor_id in self.executors + if executor_id in wrapped_requests or self._runner.state.has(request_context_state_key(executor_id)) + ] + + def validate_parent_contexts(_requests: Sequence[WorkflowEvent]) -> Callable[[], None]: + selected_scopes = validate_request_contexts( + parent_contexts, + batch, + boundary_contexts=boundary_contexts, + previous_scopes=scopes, + ) + + def accept_scope() -> None: + scopes.update(selected_scopes) + if selected_scopes: + self._runner.state.set(WORKFLOW_RESPONSE_SCOPE_KEY, next(iter(selected_scopes))) + + return accept_scope + + validations: list[_ResponseBatchValidation] = [] + if self._response_batch_validator is not None: + validations.append((self._response_batch_validator, batch)) + validations.append((validate_parent_contexts, batch)) + for executor in self.executors.values(): + delegated = [ + *(request for request in batch if request.source_executor_id == executor.id), + *wrapped_requests.get(executor.id, []), + ] + if delegated: + validations.extend( + await executor._get_response_validation_batches(delegated) # pyright: ignore[reportPrivateUsage] + ) + return validations + async def _validate_responses_internal(self, responses: Mapping[str, Any]) -> dict[str, Any]: """Validate responses before a run can consume pending requests.""" from ._agent_executor import _validate_computer_tool_result # pyright: ignore[reportPrivateUsage] @@ -1148,6 +1304,7 @@ async def _validate_responses_internal(self, responses: Mapping[str, Any]) -> di if isinstance(pending_request.data, Content) and pending_request.data.type == "computer_tool_call": _validate_computer_tool_result(pending_request.data, response) coerced_responses[request_id] = response + await self._validate_response_requests([pending_requests[request_id] for request_id in coerced_responses]) return coerced_responses async def _send_responses_internal(self, responses: Mapping[str, Any]) -> None: @@ -1531,6 +1688,9 @@ async def apply_cancellations() -> None: for request_id, event in pending.items() if event.source_executor_id in computer_executor_ids ) + await self._validate_response_requests([ + request for request_id, request in pending.items() if request_id in selected_ids + ]) cancelled_events = await self._runner.context.cancel_request_info_events(selected_ids) for request_id, request_event in cancelled_events.items(): source_executor_id = request_event.source_executor_id @@ -1542,6 +1702,7 @@ async def apply_cancellations() -> None: state=self._runner.state, runner_context=self._runner.context, ) + restore_request_context(self._runner.state, request_id) await executor._cancel_pending_request(request_id, context) # pyright: ignore[reportPrivateUsage] if checkpoint_storage is not None: diff --git a/python/packages/core/agent_framework/_workflows/_workflow_context.py b/python/packages/core/agent_framework/_workflows/_workflow_context.py index 0a887c20d9d..4d96f81d94f 100644 --- a/python/packages/core/agent_framework/_workflows/_workflow_context.py +++ b/python/packages/core/agent_framework/_workflows/_workflow_context.py @@ -20,6 +20,7 @@ WorkflowEventSource, _framework_event_origin, # type: ignore ) +from ._request_context import capture_request_context from ._runner_context import RunnerContext, WorkflowMessage from ._state import State from ._typing_utils import contains_typevar @@ -433,6 +434,7 @@ async def request_info(self, request_data: object, response_type: type, *, reque response_type=response_type, ) await self._runner_context.add_request_info_event(request_info_event) + capture_request_context(self._state, request_info_event.request_id) def get_state(self, key: str, default: Any = None) -> Any: """Get a value from the workflow state.""" diff --git a/python/packages/core/agent_framework/_workflows/_workflow_executor.py b/python/packages/core/agent_framework/_workflows/_workflow_executor.py index 7313e91f322..2feadaf43ae 100644 --- a/python/packages/core/agent_framework/_workflows/_workflow_executor.py +++ b/python/packages/core/agent_framework/_workflows/_workflow_executor.py @@ -1,9 +1,11 @@ # Copyright (c) Microsoft. All rights reserved. +from __future__ import annotations + import logging import sys import types -from collections.abc import Mapping +from collections.abc import Mapping, Sequence from dataclasses import dataclass from typing import TYPE_CHECKING, Any, cast @@ -24,7 +26,12 @@ WorkflowRunState, _framework_event_origin, # type: ignore[reportPrivateUsage] ) -from ._executor import Executor, handler +from ._executor import ( + Executor, + _ResponseBatchValidation, # pyright: ignore[reportPrivateUsage] + handler, +) +from ._request_context import capture_request_context, restore_request_context from ._request_info_mixin import response_handler from ._runner_context import WorkflowMessage from ._typing_utils import is_instance_of @@ -286,6 +293,8 @@ async def handle_subworkflow_request( - Event processing is ordered - outputs are forwarded before requests - Responses are forwarded to the sub-workflow as they arrive; the sub-workflow tracks its own pending requests and resumes when they are answered + - Forwarded requests retain their parent run context in checkpointed parent shared state. + Responses and cancellations restore it before forwarding resumed outputs. - The WorkflowExecutor keeps no per-execution bookkeeping; the sub-workflow is the single source of truth for its pending requests. Starting a new execution while the sub-workflow still has pending requests logs a warning and is only safe when the wrapped workflow is stateless @@ -293,7 +302,7 @@ async def handle_subworkflow_request( def __init__( self, - workflow: "Workflow", + workflow: Workflow, id: str, allow_direct_output: bool = False, propagate_request: bool = False, @@ -481,12 +490,15 @@ async def process_workflow(self, input_data: object, ctx: WorkflowContext[Any, A fi_kwargs, ci_kwargs = self._get_child_invocation_kwargs(ctx) # Run the sub-workflow and collect all events, passing parent kwargs - result = await self.workflow.run( + child_run = self.workflow.run( input_data, tools=ctx.get_runtime_tools(), function_invocation_kwargs=fi_kwargs, client_kwargs=ci_kwargs, ) + self.workflow._set_parent_run_context(ctx.state, self.id) # pyright: ignore[reportPrivateUsage] + result = await child_run + self._remove_reemitted_pending_requests(result, pending_requests) logger.debug(f"WorkflowExecutor {self.id} sub-workflow {self.workflow.id} completed with {len(result)} events") @@ -590,6 +602,21 @@ async def on_checkpoint_save(self) -> dict[str, Any]: "sub_workflow_checkpoint": await self.workflow._runner.build_checkpoint(), # pyright: ignore[reportPrivateUsage] } + @override + async def _get_response_validation_batches( + self, requests: Sequence[WorkflowEvent] + ) -> list[_ResponseBatchValidation]: + pending = await self.workflow._runner_context.get_pending_request_info_events() # pyright: ignore[reportPrivateUsage] + return await self.workflow._get_response_validation_batches([ # pyright: ignore[reportPrivateUsage] + pending[request.request_id] for request in requests if request.request_id in pending + ]) + + @override + async def _validate_checkpoint_restore(self, state: dict[str, Any], requests: Sequence[WorkflowEvent] = ()) -> None: + sub_workflow_checkpoint = state.get("sub_workflow_checkpoint") + if sub_workflow_checkpoint is not None: + await self.workflow._runner._validate_checkpoint_tree(sub_workflow_checkpoint, requests) # pyright: ignore[reportPrivateUsage] + @override async def on_checkpoint_restore(self, state: dict[str, Any]) -> None: """Restore the WorkflowExecutor state from a checkpoint snapshot.""" @@ -699,6 +726,7 @@ async def _forward_intermediate_output(output: Any) -> None: else: # In a workflow where the parent workflow has an executor that may intercept the # request and handle it directly, a message should be sent. + capture_request_context(ctx.state, request_id, boundary_id=self.id) await ctx.send_message(SubWorkflowRequestMessage(source_event=event, executor_id=self.id)) # Handle final state @@ -756,5 +784,8 @@ async def _handle_response( responses={request_id: response}, tools=ctx.get_runtime_tools(), ) + # Propagated replies already restored their context before their response handler. + if ctx.request_id != request_id: + restore_request_context(ctx.state, request_id, boundary_id=self.id) self._remove_reemitted_pending_requests(result, pending_requests) await self._process_workflow_result(result, ctx) diff --git a/python/packages/core/tests/workflow/conftest.py b/python/packages/core/tests/workflow/conftest.py new file mode 100644 index 00000000000..f9b7b3811f3 --- /dev/null +++ b/python/packages/core/tests/workflow/conftest.py @@ -0,0 +1,278 @@ +# Copyright (c) Microsoft. All rights reserved. + +import asyncio +from collections.abc import Callable +from typing import Any + +import pytest + +from agent_framework import ( + Case, + Default, + Executor, + SubWorkflowRequestMessage, + SubWorkflowResponseMessage, + Workflow, + WorkflowBuilder, + WorkflowContext, + WorkflowExecutor, + handler, + response_handler, +) + + +@pytest.fixture +def selected_reply_workflow() -> Callable[ + [str, bool, bool], tuple[Workflow, dict[str, Workflow], dict[str, SubWorkflowRequestMessage], list[str]] +]: + def build( + routing: str, select_both: bool = False, stateful: bool = False + ) -> tuple[Workflow, dict[str, Workflow], dict[str, SubWorkflowRequestMessage], list[str]]: + requests: dict[str, SubWorkflowRequestMessage] = {} + calls: list[str] = [] + + class Requester(Executor): + @handler + async def start(self, message: str, ctx: WorkflowContext) -> None: + await ctx.request_info(message, bool, request_id="shared") + + @response_handler + async def respond(self, original: str, approved: bool, ctx: WorkflowContext[Any, str]) -> None: + await ctx.yield_output(original) + + class Coordinator(Executor): + def __init__(self) -> None: + super().__init__(id="coordinator") + self.target = "first" + + async def on_checkpoint_save(self) -> dict[str, Any]: + return {"target": self.target} + + async def on_checkpoint_restore(self, state: dict[str, Any]) -> None: + self.target = state["target"] + + @handler + async def start(self, message: str, ctx: WorkflowContext[str]) -> None: + for target in ("first", "second") if message == "both" else (message,): + await ctx.send_message(target, target_id=target) + + @handler + async def reply( + self, messages: list[SubWorkflowResponseMessage], ctx: WorkflowContext[SubWorkflowResponseMessage] + ) -> None: + for message in messages: + await ctx.send_message(message) + + @handler + async def intercept(self, request: SubWorkflowRequestMessage, ctx: WorkflowContext) -> None: + requests[request.executor_id] = request + + @handler + async def queue_ordered( + self, + messages: tuple[str, SubWorkflowResponseMessage], + ctx: WorkflowContext[str | SubWorkflowResponseMessage], + ) -> None: + for message in messages: + await ctx.send_message(message, target_id="first") + + class RoutingWorkflowExecutor(WorkflowExecutor): + @handler + async def apply_routing(self, message: str, ctx: WorkflowContext[Any, Any]) -> None: + if message == "enable-first": + coordinator.target = "first" + else: + await self.process_workflow(message, ctx) + + def select(message: Any, targets: list[str]) -> list[str]: + if isinstance(message, SubWorkflowResponseMessage): + calls.append("selection") + return targets if select_both else [coordinator.target if stateful else "first"] + return targets + + async def first_condition(message: Any) -> bool: + if isinstance(message, SubWorkflowResponseMessage): + calls.append("first") + if stateful: + return coordinator.target == "first" + return True + + async def second_condition(message: Any) -> bool: + if isinstance(message, SubWorkflowResponseMessage): + calls.append("second") + return select_both or (stateful and coordinator.target == "second") + return coordinator.target == "first" if stateful else True + + def first_case(message: Any) -> bool: + if isinstance(message, SubWorkflowResponseMessage): + calls.append("switch") + return True + return message == "first" + + children = { + name: WorkflowBuilder(start_executor=Requester(id="requester")).build() for name in ("first", "second") + } + wrapper_type = RoutingWorkflowExecutor if stateful else WorkflowExecutor + wrappers = {name: wrapper_type(child, id=name, allow_direct_output=True) for name, child in children.items()} + coordinator = Coordinator() + builder = WorkflowBuilder(start_executor=coordinator) + if routing == "condition": + builder.add_edge(coordinator, wrappers["first"], condition=first_condition) + builder.add_edge(coordinator, wrappers["second"], condition=second_condition) + elif routing == "selection": + builder.add_multi_selection_edge_group(coordinator, list(wrappers.values()), select) + else: + builder.add_switch_case_edge_group( + coordinator, [Case(condition=first_case, target=wrappers["first"]), Default(target=wrappers["second"])] + ) + for wrapper in wrappers.values(): + builder.add_edge(wrapper, coordinator) + return builder.build(), children, requests, calls + + return build + + +@pytest.fixture +def waiting_reply_workflow() -> Callable[[bool, bool], tuple[Workflow, Workflow, list[str]]]: + def build(targeted: bool, preceding_reply: bool) -> tuple[Workflow, Workflow, list[str]]: + started = asyncio.Event() + released = asyncio.Event() + completed: list[str] = [] + requests: list[SubWorkflowRequestMessage] = [] + + class Requester(Executor): + @handler + async def start(self, message: str, ctx: WorkflowContext) -> None: + await ctx.request_info(message, bool, request_id="approval") + + @response_handler + async def respond(self, original: str, approved: bool, ctx: WorkflowContext[Any, str]) -> None: + await started.wait() + released.set() + await ctx.yield_output(original) + + class Waiter(Executor): + @handler + async def wait(self, message: str, ctx: WorkflowContext) -> None: + started.set() + await released.wait() + completed.append(message) + + class ReplyWaiter(Waiter): + @handler + async def receive_reply(self, message: SubWorkflowResponseMessage, ctx: WorkflowContext) -> None: + completed.append("reply") + + class Coordinator(Executor): + @handler + async def start(self, message: str, ctx: WorkflowContext[str]) -> None: + await ctx.send_message(message, target_id="child") + + @handler + async def intercept(self, request: SubWorkflowRequestMessage, ctx: WorkflowContext) -> None: + requests.append(request) + + @handler + async def reply(self, approved: bool, ctx: WorkflowContext[str | SubWorkflowResponseMessage]) -> None: + await ctx.send_message("wait", target_id="waiter") + if preceding_reply: + await ctx.send_message(requests[0].create_response(approved), target_id="waiter") + await ctx.send_message(requests[0].create_response(approved), target_id="child" if targeted else None) + + child = WorkflowBuilder(start_executor=Requester(id="requester")).build() + wrapper = WorkflowExecutor(child, id="child", allow_direct_output=True) + coordinator = Coordinator(id="coordinator") + parent = ( + WorkflowBuilder(start_executor=coordinator) + .add_edge(coordinator, ReplyWaiter(id="waiter") if targeted else Waiter(id="waiter")) + .add_edge(coordinator, wrapper) + .add_edge(wrapper, coordinator) + .build() + ) + return parent, child, completed + + return build + + +@pytest.fixture +def staggered_reply_workflow() -> Callable[ + [], tuple[Workflow, dict[str, Workflow], dict[str, SubWorkflowRequestMessage], list[tuple[str, str]]] +]: + def build() -> tuple[Workflow, dict[str, Workflow], dict[str, SubWorkflowRequestMessage], list[tuple[str, str]]]: + requests: dict[str, SubWorkflowRequestMessage] = {} + observed: list[tuple[str, str]] = [] + + class Requester(Executor): + @handler + async def start(self, message: str, ctx: WorkflowContext) -> None: + await ctx.request_info(message, bool, request_id=message) + + @response_handler + async def respond(self, original: str, approved: bool, ctx: WorkflowContext[Any, str]) -> None: + await ctx.yield_output(original) + + class Coordinator(Executor): + @handler + async def start(self, message: str, ctx: WorkflowContext[str]) -> None: + for target in ("first", "second") if message == "both" else (message,): + await ctx.send_message(target, target_id=target) + + @handler + async def intercept(self, request: SubWorkflowRequestMessage, ctx: WorkflowContext) -> None: + requests[request.executor_id] = request + await ctx.request_info(request.executor_id, bool, request_id=request.executor_id) + + @response_handler + async def respond( + self, original: str, approved: bool, ctx: WorkflowContext[SubWorkflowResponseMessage] + ) -> None: + await ctx.send_message(requests[original].create_response(approved), target_id=original) + + @handler + async def stagger( + self, + replies: list[SubWorkflowResponseMessage], + ctx: WorkflowContext[SubWorkflowResponseMessage | tuple[SubWorkflowResponseMessage]], + ) -> None: + await ctx.send_message(replies[0], target_id="first") + await ctx.send_message((replies[1],), target_id="delay") + + class ReplyHop(Executor): + @handler + async def forward( + self, reply: tuple[SubWorkflowResponseMessage], ctx: WorkflowContext[SubWorkflowResponseMessage] + ) -> None: + await ctx.send_message(reply[0]) + + class OutputHop(Executor): + @handler + async def forward(self, message: str, ctx: WorkflowContext[str]) -> None: + await ctx.send_message(message) + + class Recorder(Executor): + @handler + async def record(self, message: str, ctx: WorkflowContext[Any, str]) -> None: + from agent_framework._workflows._const import WORKFLOW_SESSION_ID_KEY + + observed.append((message, ctx.get_state(WORKFLOW_SESSION_ID_KEY))) + await ctx.yield_output(message) + + children = { + name: WorkflowBuilder(start_executor=Requester(id="requester")).build() for name in ("first", "second") + } + wrappers = {name: WorkflowExecutor(child, id=name) for name, child in children.items()} + coordinator = Coordinator(id="coordinator") + delay = ReplyHop(id="delay") + output = OutputHop(id="output") + recorder = Recorder(id="recorder") + builder = WorkflowBuilder(start_executor=coordinator) + for wrapper in wrappers.values(): + builder.add_edge(coordinator, wrapper).add_edge( + wrapper, coordinator, condition=lambda message: isinstance(message, SubWorkflowRequestMessage) + ).add_edge(wrapper, output) + parent = ( + builder.add_edge(coordinator, delay).add_edge(delay, wrappers["second"]).add_edge(output, recorder).build() + ) + return parent, children, requests, observed + + return build diff --git a/python/packages/core/tests/workflow/test_checkpoint.py b/python/packages/core/tests/workflow/test_checkpoint.py index e1c6a77f6ea..f23ec064b44 100644 --- a/python/packages/core/tests/workflow/test_checkpoint.py +++ b/python/packages/core/tests/workflow/test_checkpoint.py @@ -713,8 +713,13 @@ def _build_workflow() -> Any: # Resuming from the response-entry checkpoint replays the continuation and reproduces the output. replay_workflow = _build_workflow() + validated_requests: list[list[Any]] = [] + replay_workflow._response_batch_validator = lambda requests: validated_requests.append(list(requests)) replay_outputs = (await replay_workflow.run(checkpoint_id=response_entry.checkpoint_id)).get_outputs() assert replay_outputs == ["approve->yes"] + # Validate both the detached checkpoint candidate and the live dispatch batch. + assert [[event.request_id for event in batch] for batch in validated_requests] == [[request_id], [request_id]] + assert all(batch[0].data == "approve" for batch in validated_requests) async def test_memory_checkpoint_storage_roundtrip_json_native_types(): diff --git a/python/packages/core/tests/workflow/test_state.py b/python/packages/core/tests/workflow/test_state.py index 6b6a52c6084..8973dbbe830 100644 --- a/python/packages/core/tests/workflow/test_state.py +++ b/python/packages/core/tests/workflow/test_state.py @@ -2,7 +2,7 @@ """Unit tests for the State class superstep caching behavior.""" -from unittest.mock import MagicMock +from unittest.mock import MagicMock, patch import pytest @@ -10,6 +10,136 @@ from agent_framework._workflows._state import State +def test_request_origin_fan_out_copies_only_incremental_entries() -> None: + import copy + + from agent_framework._workflows._const import WORKFLOW_REQUEST_PARENT_CONTEXTS_KEY, WORKFLOW_SESSION_ID_KEY + from agent_framework._workflows._request_context import capture_request_context, restore_request_context + + state = State() + state.set(WORKFLOW_SESSION_ID_KEY, "origin") + request_count = 200 + with patch("agent_framework._workflows._state.copy.deepcopy", wraps=copy.deepcopy) as deepcopy: + for index in range(request_count): + capture_request_context(state, str(index)) + state.commit() + for index in range(request_count): + restore_request_context(state, str(index)) + assert state.get(WORKFLOW_SESSION_ID_KEY) == "origin" + copied_entries = sum(len(call.args[0]) for call in deepcopy.call_args_list if isinstance(call.args[0], dict)) + + assert state.get(WORKFLOW_REQUEST_PARENT_CONTEXTS_KEY) == {} + assert copied_entries < 40 * request_count + + +def test_incremental_mapping_updates_preserve_snapshots_and_rollback() -> None: + state = State() + state.set("origins", {"first": {"scope": "A"}, "second": {"scope": "B"}}) + state.commit() + snapshot = state.export_state() + update = {"scope": "C"} + + state._set_mapping_item("origins", "third", update) + update["scope"] = "mutated" + removed = state._pop_mapping_item("origins", "first") + removed["scope"] = "mutated" + assert state._get_mapping_item("origins", "third") == {"scope": "C"} + assert state.get("origins") == {"second": {"scope": "B"}, "third": {"scope": "C"}} + assert state.export_state() == snapshot + + state.discard() + assert state.get("origins") == snapshot["origins"] + state._pop_mapping_item("origins", "first") + state._set_mapping_item("origins", "second", {"scope": "D"}) + state.commit() + assert state.export_state()["origins"] == {"second": {"scope": "D"}} + assert snapshot["origins"]["second"] == {"scope": "B"} + + +@pytest.mark.parametrize("operation", ["set", "delete", "clear"]) +def test_whole_value_operations_replace_pending_mapping_updates(operation: str) -> None: + state = State() + state.set("origins", {"first": 1}) + state.commit() + state._set_mapping_item("origins", "second", 2) + + if operation == "set": + state.set("origins", {"replacement": 3}) + expected = {"replacement": 3} + elif operation == "delete": + state.delete("origins") + expected = None + else: + state.clear() + expected = None + assert state.get("origins") == expected + state.commit() + assert state.get("origins") == expected + + +def test_incremental_mapping_updates_support_new_deleted_and_null_entries() -> None: + state = State() + state._set_mapping_item("origins", "first", None) + assert state.has("origins") + state._pop_mapping_item("origins", "first") + assert state.get("origins") == {} + state.commit() + state.delete("origins") + state._set_mapping_item("origins", "second", {"scope": "B"}) + state.commit() + assert state.export_state()["origins"] == {"second": {"scope": "B"}} + + +def test_request_origin_owners_remain_independent_for_the_same_request_id() -> None: + from agent_framework._workflows._const import WORKFLOW_SESSION_ID_KEY + from agent_framework._workflows._request_context import capture_request_context, restore_request_context + + state = State() + state.set(WORKFLOW_SESSION_ID_KEY, "child-boundary") + capture_request_context(state, "same-id", boundary_id="child") + state.set(WORKFLOW_SESSION_ID_KEY, "coordinator-relay") + capture_request_context(state, "same-id") + state.commit() + + restore_request_context(state, "same-id") + assert state.get(WORKFLOW_SESSION_ID_KEY) == "coordinator-relay" + state.commit() + restore_request_context(state, "same-id", boundary_id="child") + assert state.get(WORKFLOW_SESSION_ID_KEY) == "child-boundary" + + +def test_wrapped_request_restores_legacy_checkpoint_origin() -> None: + from agent_framework._workflows._const import WORKFLOW_REQUEST_PARENT_CONTEXTS_KEY, WORKFLOW_SESSION_ID_KEY + from agent_framework._workflows._request_context import capture_request_context, restore_request_context + + legacy = State() + legacy.set(WORKFLOW_SESSION_ID_KEY, "legacy-boundary") + capture_request_context(legacy, "request") + legacy.commit() + restored = State() + restored.import_state(legacy.export_state()) + restored.set(WORKFLOW_SESSION_ID_KEY, "later-run") + + restore_request_context(restored, "request", boundary_id="child") + + assert restored.get(WORKFLOW_SESSION_ID_KEY) == "legacy-boundary" + assert restored.get(WORKFLOW_REQUEST_PARENT_CONTEXTS_KEY) == {} + + +@pytest.mark.parametrize("operation", ["capture", "restore"]) +def test_incremental_request_origin_operations_reject_malformed_entries(operation: str) -> None: + from agent_framework._workflows._const import WORKFLOW_REQUEST_PARENT_CONTEXTS_KEY + from agent_framework._workflows._request_context import capture_request_context, restore_request_context + + state = State() + state.set(WORKFLOW_REQUEST_PARENT_CONTEXTS_KEY, {"request": None}) + state.commit() + action = capture_request_context if operation == "capture" else restore_request_context + with pytest.raises(ValueError, match="Invalid parent workflow request context"): + action(state, "request") + assert state.get(WORKFLOW_REQUEST_PARENT_CONTEXTS_KEY) == {"request": None} + + class TestStateBasicOperations: """Tests for basic State get/set/has/delete operations.""" diff --git a/python/packages/core/tests/workflow/test_sub_workflow.py b/python/packages/core/tests/workflow/test_sub_workflow.py index cfd144b816b..bb4836b45ed 100644 --- a/python/packages/core/tests/workflow/test_sub_workflow.py +++ b/python/packages/core/tests/workflow/test_sub_workflow.py @@ -2,6 +2,7 @@ import asyncio import logging +from collections.abc import Sequence from dataclasses import dataclass, field from typing import Any from uuid import uuid4 @@ -220,6 +221,525 @@ async def test_basic_sub_workflow() -> None: assert parent.result.is_valid is True +@pytest.mark.parametrize("restore", [False, True]) +async def test_interleaved_sibling_responses_reject_mixed_parent_scopes(restore: bool) -> None: + from agent_framework._workflows._const import WORKFLOW_REQUEST_PARENT_CONTEXTS_KEY, WORKFLOW_SESSION_ID_KEY + + storage = InMemoryCheckpointStorage() + + def build() -> Workflow: + start = EmailFormatValidator() + children = [ + WorkflowExecutor( + WorkflowBuilder(start_executor=EmailDomainValidator()).build(), + id=name, + propagate_request=True, + allow_direct_output=True, + ) + for name in ("first", "second") + ] + return ( + WorkflowBuilder(start_executor=start, checkpoint_storage=storage).add_fan_out_edges(start, children).build() + ) + + workflow = build() + await workflow.run(EmailValidationRequest("a@first.example")) + requests_a = await workflow._runner_context.get_pending_request_info_events() + scope_a = workflow._runner.state.get(WORKFLOW_SESSION_ID_KEY) + await workflow.run(EmailValidationRequest("b@second.example")) + requests = await workflow._runner_context.get_pending_request_info_events() + requests_b = {key: request for key, request in requests.items() if key not in requests_a} + scope_b = workflow._runner.state.get(WORKFLOW_SESSION_ID_KEY) + if restore: + checkpoint_id = workflow.get_last_checkpoint_id() + assert checkpoint_id is not None + workflow = build() + await workflow.run(checkpoint_id=checkpoint_id) + + first_a = next(iter(requests_a.values())) + second_b = next( + request for request in requests_b.values() if request.source_executor_id != first_a.source_executor_id + ) + with pytest.raises(ValueError, match="multiple parent workflow scopes"): + await workflow.run(responses={first_a.request_id: True, second_b.request_id: True}) + with pytest.raises(ValueError, match="multiple parent workflow scopes"): + await workflow.cancel_pending_requests([first_a.request_id, second_b.request_id]) + assert set(await workflow._runner_context.get_pending_request_info_events()) == set(requests) + assert workflow._runner.state.get(WORKFLOW_SESSION_ID_KEY) == scope_b + + for batch, scope, email in ( + (requests_b, scope_b, "b@second.example"), + (requests_a, scope_a, "a@first.example"), + ): + result = await workflow.run(responses=dict.fromkeys(batch, True)) + outputs = result.get_outputs() + assert len(outputs) == 2 + assert {output.email for output in outputs} == {email} + assert workflow._runner.state.get(WORKFLOW_SESSION_ID_KEY) == scope + assert workflow._runner.state.get(WORKFLOW_REQUEST_PARENT_CONTEXTS_KEY) == {} + assert not await workflow._runner_context.get_pending_request_info_events() + + +class _WrappedReplyCoordinator(Executor): + def __init__(self) -> None: + super().__init__(id="coordinator") + self.requests: dict[str, SubWorkflowRequestMessage] = {} + self.outputs: list[tuple[str, str]] = [] + + @handler + async def start(self, email: str, ctx: WorkflowContext[DomainCheckRequest]) -> None: + await ctx.send_message(DomainCheckRequest(domain=email.split("@")[1], email=email)) + + @handler + async def intercept(self, request: SubWorkflowRequestMessage, ctx: WorkflowContext) -> None: + email = request.source_event.data.email + self.requests[email] = request + await ctx.request_info(email, bool, request_id=email) + + @response_handler + async def respond(self, email: str, approved: bool, ctx: WorkflowContext[SubWorkflowResponseMessage]) -> None: + request = self.requests[email] + await ctx.send_message(request.create_response(approved), target_id=request.executor_id) + + @handler + async def queue_replies( + self, replies: list[SubWorkflowResponseMessage], ctx: WorkflowContext[SubWorkflowResponseMessage] + ) -> None: + for reply in replies: + await ctx.send_message(reply) + + @handler + async def collect(self, result: ValidationResult, ctx: WorkflowContext) -> None: + from agent_framework._workflows._const import WORKFLOW_SESSION_ID_KEY + + self.outputs.append((result.email, ctx.get_state(WORKFLOW_SESSION_ID_KEY))) + + +def _build_wrapped_reply_workflow() -> tuple[Workflow, Workflow, _WrappedReplyCoordinator]: + coordinator = _WrappedReplyCoordinator() + child = WorkflowBuilder(start_executor=EmailDomainValidator()).build() + wrapper = WorkflowExecutor(child, id="child") + parent = ( + WorkflowBuilder(start_executor=coordinator) + .add_edge(coordinator, wrapper) + .add_edge(wrapper, coordinator) + .build() + ) + return parent, child, coordinator + + +@pytest.mark.parametrize("streaming", [False, True]) +@pytest.mark.parametrize("delivery", ["live", "queued_and_new"]) +async def test_wrapped_reply_delivery_rejects_later_mixed_origins(delivery: str, streaming: bool) -> None: + from agent_framework._workflows._const import WORKFLOW_REQUEST_PARENT_CONTEXTS_KEY, WORKFLOW_SESSION_ID_KEY + + parent, child, coordinator = _build_wrapped_reply_workflow() + emails = ["a@first.example", "b@second.example"] + scopes: list[str] = [] + for email in emails: + await parent.run(email) + scopes.append(parent._runner.state.get(WORKFLOW_SESSION_ID_KEY)) + child_pending = await child._runner_context.get_pending_request_info_events() + parent_pending = await parent._runner_context.get_pending_request_info_events() + origins = parent._runner.state.get(WORKFLOW_REQUEST_PARENT_CONTEXTS_KEY) + checkpoint = await parent._runner.build_checkpoint() + replies = [coordinator.requests[email].create_response(True) for email in emails] + if delivery == "queued_and_new": + await parent._runner_context.send_message( + WorkflowMessage(data=replies[0], source_id=coordinator.id, target_id="child") + ) + + with pytest.raises(ValueError, match="multiple parent workflow scopes"): + message = replies if delivery == "live" else None + responses = {emails[1]: True} if delivery == "queued_and_new" else None + if streaming: + _ = [event async for event in parent.run(message, responses=responses, stream=True)] + else: + await parent.run(message, responses=responses) + + assert coordinator.outputs == [] + expected_child_pending = dict(child_pending) + expected_parent_pending = dict(parent_pending) + if delivery == "live": + del expected_child_pending[replies[0].source_event.request_id] + else: + del expected_parent_pending[emails[1]] + assert await child._runner_context.get_pending_request_info_events() == expected_child_pending + assert await parent._runner_context.get_pending_request_info_events() == expected_parent_pending + assert parent._runner.state.get(WORKFLOW_REQUEST_PARENT_CONTEXTS_KEY) == origins + queued = await parent._runner_context.drain_messages() + assert [ + message.data.source_event.request_id + for batch in queued.values() + for message in batch + if isinstance(message.data, SubWorkflowResponseMessage) + ] == [reply.source_event.request_id for reply in (replies if delivery == "queued_and_new" else replies[1:])] + + await parent._runner.restore_checkpoint(checkpoint) + restored_child = await child._runner_context.get_pending_request_info_events() + restored_parent = await parent._runner_context.get_pending_request_info_events() + assert {key: event.to_dict() for key, event in restored_child.items()} == { + key: event.to_dict() for key, event in child_pending.items() + } + assert {key: event.to_dict() for key, event in restored_parent.items()} == { + key: event.to_dict() for key, event in parent_pending.items() + } + for email, scope in reversed(list(zip(emails, scopes, strict=True))): + await parent.run(responses={email: True}) + assert coordinator.outputs[-1] == (email, scope) + assert not await child._runner_context.get_pending_request_info_events() + assert not await parent._runner_context.get_pending_request_info_events() + assert parent._runner.state.get(WORKFLOW_REQUEST_PARENT_CONTEXTS_KEY) == {} + + +@pytest.mark.parametrize("delivery", ["live", "queued_and_new"]) +async def test_wrapped_reply_delivery_applies_child_policy(delivery: str) -> None: + parent, child, coordinator = _build_wrapped_reply_workflow() + email = "a@first.example" + await parent.run(email) + parent_pending = await parent._runner_context.get_pending_request_info_events() + child_pending = await child._runner_context.get_pending_request_info_events() + reply = coordinator.requests[email].create_response(True) + + def reject(requests: Sequence[WorkflowEvent]) -> None: + if requests: + raise ValueError("child reply policy") + + child._response_batch_validator = reject + if delivery == "queued_and_new": + await parent._runner_context.send_message( + WorkflowMessage(data=reply, source_id=coordinator.id, target_id="child") + ) + run = parent.run(responses={email: True}) + else: + run = parent.run([reply]) + with pytest.raises(ValueError, match="child reply policy"): + await run + + assert coordinator.outputs == [] + assert await parent._runner_context.get_pending_request_info_events() == ( + parent_pending if delivery == "live" else {} + ) + assert await child._runner_context.get_pending_request_info_events() == child_pending + queued = await parent._runner_context.drain_messages() + assert [ + message.data.source_event.request_id + for batch in queued.values() + for message in batch + if isinstance(message.data, SubWorkflowResponseMessage) + ] == [reply.source_event.request_id] * (2 if delivery == "queued_and_new" else 1) + + +@pytest.mark.parametrize("routing", ["condition", "selection", "switch"]) +@pytest.mark.parametrize("restore", [False, True]) +@pytest.mark.parametrize("streaming", [False, True]) +async def test_wrapped_reply_delivery_uses_selected_routes( + selected_reply_workflow, routing: str, restore: bool, streaming: bool +) -> None: + from agent_framework._workflows._const import WORKFLOW_SESSION_ID_KEY + + parent, children, requests, calls = selected_reply_workflow(routing, False, False) + await parent.run("first") + first_scope = parent._runner.state.get(WORKFLOW_SESSION_ID_KEY) + await parent.run("second") + assert parent._runner.state.get(WORKFLOW_SESSION_ID_KEY) != first_scope + pending_second = await children["second"]._runner_context.get_pending_request_info_events() + assert requests["first"].source_event.request_id == requests["second"].source_event.request_id == "shared" + reply = requests["first"].create_response(True) + checkpoint_id: str | None = None + storage: InMemoryCheckpointStorage | None = None + if restore: + await parent._runner_context.send_message(WorkflowMessage(data=reply, source_id="coordinator")) + checkpoint = await parent._runner.build_checkpoint() + storage = InMemoryCheckpointStorage() + checkpoint_id = await storage.save(checkpoint) + + message = None if restore else [reply] + if streaming: + events = [ + event + async for event in parent.run(message, checkpoint_id=checkpoint_id, checkpoint_storage=storage, stream=True) + ] + outputs = [event.data for event in events if event.type == "output"] + else: + outputs = (await parent.run(message, checkpoint_id=checkpoint_id, checkpoint_storage=storage)).get_outputs() + + assert outputs == ["first"] + assert not await children["first"]._runner_context.get_pending_request_info_events() + remaining = await children["second"]._runner_context.get_pending_request_info_events() + assert {key: event.to_dict() for key, event in remaining.items()} == { + key: event.to_dict() for key, event in pending_second.items() + } + assert parent._runner.state.get(WORKFLOW_SESSION_ID_KEY) == first_scope + assert ( + sorted(calls) == {"condition": ["first", "second"], "selection": ["selection"], "switch": ["switch"]}[routing] + ) + + +@pytest.mark.parametrize("routing", ["condition", "selection"]) +@pytest.mark.parametrize("restore", [False, True]) +async def test_wrapped_reply_selected_mixed_scopes_preserve_queue( + selected_reply_workflow, routing: str, restore: bool +) -> None: + from agent_framework._workflows._const import WORKFLOW_SESSION_ID_KEY + + parent, children, requests, _ = selected_reply_workflow(routing, True, False) + await parent.run("first") + await parent.run("second") + pending = {name: await child._runner_context.get_pending_request_info_events() for name, child in children.items()} + message = WorkflowMessage(data=requests["first"].create_response(True), source_id="coordinator") + await parent._runner_context.send_message(message) + scope = parent._runner.state.get(WORKFLOW_SESSION_ID_KEY) + if restore: + checkpoint = await parent._runner.build_checkpoint() + await parent._runner.restore_checkpoint(checkpoint) + [restored_message] = (await parent._runner._get_queued_messages())["coordinator"] + assert restored_message.data.source_event.to_dict() == message.data.source_event.to_dict() + message = restored_message + pending = { + name: await child._runner_context.get_pending_request_info_events() for name, child in children.items() + } + with pytest.raises(ValueError, match="multiple parent workflow scopes"): + await parent._runner._run_iteration() + assert parent._runner.state.get(WORKFLOW_SESSION_ID_KEY) == scope + assert await parent._runner._get_queued_messages() == {"coordinator": [message]} + for name, child in children.items(): + assert await child._runner_context.get_pending_request_info_events() == pending[name] + + +@pytest.mark.parametrize("routing", ["condition", "selection"]) +async def test_wrapped_reply_validates_all_selected_child_policies_before_delivery( + selected_reply_workflow, routing: str +) -> None: + parent, children, requests, _ = selected_reply_workflow(routing, True, False) + await parent.run("both") + pending = {name: await child._runner_context.get_pending_request_info_events() for name, child in children.items()} + migrations: list[str] = [] + + def stage(requests: Sequence[WorkflowEvent]): + return lambda: migrations.append("first") + + def reject(requests: Sequence[WorkflowEvent]) -> None: + if requests: + raise ValueError("second child policy") + + children["first"]._response_batch_validator = stage + children["second"]._response_batch_validator = reject + with pytest.raises(ValueError, match="second child policy"): + await parent.run([requests["first"].create_response(True)]) + assert migrations == [] + for name, child in children.items(): + assert await child._runner_context.get_pending_request_info_events() == pending[name] + [queued] = (await parent._runner._get_queued_messages())["coordinator"] + assert queued.data.source_event.request_id == "shared" + + +@pytest.mark.parametrize("routing", ["condition", "selection"]) +@pytest.mark.parametrize("streaming", [False, True]) +async def test_wrapped_reply_routing_observes_earlier_same_edge_handler( + selected_reply_workflow, routing: str, streaming: bool +) -> None: + parent, children, requests, calls = selected_reply_workflow(routing, False, True) + await parent.run("first") + coordinator = parent.get_start_executor() + await coordinator.on_checkpoint_restore({"target": "second"}) + message = ("enable-first", requests["first"].create_response(True)) + + if streaming: + events = [event async for event in parent.run(message, stream=True)] + outputs = [event.data for event in events if event.type == "output"] + else: + outputs = (await parent.run(message)).get_outputs() + + assert await coordinator.on_checkpoint_save() == {"target": "first"} + assert outputs == ["first"] + assert not await children["first"]._runner_context.get_pending_request_info_events() + assert calls == {"condition": ["first"], "selection": ["selection"]}[routing] + + +@pytest.mark.parametrize(("targeted", "preceding_reply"), [(False, False), (True, False), (True, True)]) +async def test_wrapped_reply_does_not_wait_for_unrelated_outgoing_handler( + waiting_reply_workflow, targeted: bool, preceding_reply: bool +) -> None: + parent, child, completed = waiting_reply_workflow(targeted, preceding_reply) + await parent.run("start") + + result = await asyncio.wait_for(parent.run(True), timeout=2.0) + + assert result.get_outputs() == ["start"] + assert completed == (["wait", "reply"] if preceding_reply else ["wait"]) + assert not await child._runner_context.get_pending_request_info_events() + + +@pytest.mark.parametrize("routing", ["condition", "selection"]) +@pytest.mark.parametrize("streaming", [False, True]) +async def test_wrapped_reply_routing_uses_checkpointed_executor_fields( + selected_reply_workflow, routing: str, streaming: bool +) -> None: + parent, children, requests, calls = selected_reply_workflow(routing, False, True) + await parent.run("first") + await parent.run("second") + coordinator = parent.get_start_executor() + await coordinator.on_checkpoint_restore({"target": "second"}) + await parent._runner_context.send_message( + WorkflowMessage(data=requests["second"].create_response(True), source_id="coordinator") + ) + storage = InMemoryCheckpointStorage() + checkpoint_id = await storage.save(await parent._runner.build_checkpoint()) + await coordinator.on_checkpoint_restore({"target": "first"}) + + if streaming: + events = [ + event async for event in parent.run(checkpoint_id=checkpoint_id, checkpoint_storage=storage, stream=True) + ] + outputs = [event.data for event in events if event.type == "output"] + else: + outputs = (await parent.run(checkpoint_id=checkpoint_id, checkpoint_storage=storage)).get_outputs() + + assert await coordinator.on_checkpoint_save() == {"target": "second"} + assert outputs == ["second"] + assert not await children["second"]._runner_context.get_pending_request_info_events() + assert set(await children["first"]._runner_context.get_pending_request_info_events()) == {"shared"} + assert sorted(calls) == {"condition": ["first", "second"], "selection": ["selection"]}[routing] + + +@pytest.mark.parametrize("streaming", [False, True]) +@pytest.mark.parametrize("restore", [False, True]) +async def test_staggered_wrapped_replies_reject_cross_scope_continuation( + staggered_reply_workflow, streaming: bool, restore: bool +) -> None: + parent, children, requests, observed = staggered_reply_workflow() + await parent.run("first") + await parent.run("second") + replies = [requests[name].create_response(True) for name in ("first", "second")] + checkpoint_id: str | None = None + storage: InMemoryCheckpointStorage | None = None + if restore: + await parent._runner_context.send_message( + WorkflowMessage(data=replies[0], source_id="coordinator", target_id="first") + ) + await parent._runner_context.send_message( + WorkflowMessage(data=(replies[1],), source_id="coordinator", target_id="delay") + ) + await parent._runner._run_iteration() + storage = InMemoryCheckpointStorage() + checkpoint_id = await storage.save(await parent._runner.build_checkpoint()) + parent, children, _, observed = staggered_reply_workflow() + + with pytest.raises(ValueError, match="multiple parent workflow scopes"): + if streaming: + _ = [ + event + async for event in parent.run( + None if restore else replies, + checkpoint_id=checkpoint_id, + checkpoint_storage=storage, + stream=True, + ) + ] + else: + await parent.run(None if restore else replies, checkpoint_id=checkpoint_id, checkpoint_storage=storage) + assert not await children["first"]._runner_context.get_pending_request_info_events() + assert set(await children["second"]._runner_context.get_pending_request_info_events()) == {"second"} + assert observed == [] + + +@pytest.mark.parametrize("streaming", [False, True]) +async def test_staggered_wrapped_replies_allow_one_parent_scope(staggered_reply_workflow, streaming: bool) -> None: + from agent_framework._workflows._const import WORKFLOW_SESSION_ID_KEY + + parent, children, requests, observed = staggered_reply_workflow() + await parent.run("both") + scope = parent._runner.state.get(WORKFLOW_SESSION_ID_KEY) + replies = [requests[name].create_response(True) for name in ("first", "second")] + if streaming: + events = [event async for event in parent.run(replies, stream=True)] + outputs = [event.data for event in events if event.type == "output"] + else: + outputs = (await parent.run(replies)).get_outputs() + assert sorted(outputs) == ["first", "second"] + assert sorted(observed) == [("first", scope), ("second", scope)] + for child in children.values(): + assert not await child._runner_context.get_pending_request_info_events() + + +@pytest.mark.parametrize("restore", [False, True]) +async def test_separate_response_runs_can_restore_different_parent_scopes( + staggered_reply_workflow, restore: bool +) -> None: + from agent_framework._workflows._const import WORKFLOW_SESSION_ID_KEY + + parent, children, _, observed = staggered_reply_workflow() + scopes: list[str] = [] + for name in ("first", "second"): + await parent.run(name) + scopes.append(parent._runner.state.get(WORKFLOW_SESSION_ID_KEY)) + assert scopes[0] != scopes[1] + for name, scope in zip(("first", "second"), scopes, strict=True): + if restore and name == "second": + storage = InMemoryCheckpointStorage() + checkpoint_id = await storage.save(await parent._runner.build_checkpoint()) + result = await parent.run(checkpoint_id=checkpoint_id, checkpoint_storage=storage, responses={name: True}) + else: + result = await parent.run(responses={name: True}) + assert result.get_outputs() == [name] + assert observed[-1] == (name, scope) + assert not await children[name]._runner_context.get_pending_request_info_events() + + +@pytest.mark.parametrize("scope", ["", False, ["first"]]) +async def test_checkpoint_rejects_malformed_response_scope_without_restoring( + staggered_reply_workflow, scope: Any +) -> None: + from agent_framework import WorkflowCheckpointException + from agent_framework._workflows._checkpoint_encoding import encode_checkpoint_value + from agent_framework._workflows._const import WORKFLOW_RESPONSE_SCOPE_KEY + + parent, _, _, _ = staggered_reply_workflow() + await parent.run("first") + checkpoint = await parent._runner.build_checkpoint() + checkpoint.state[WORKFLOW_RESPONSE_SCOPE_KEY] = scope + before = encode_checkpoint_value(parent._runner.state.export_state()) + with pytest.raises(WorkflowCheckpointException) as error: + await parent._runner.restore_checkpoint(checkpoint) + assert isinstance(error.value.__cause__, ValueError) + assert "Invalid parent workflow response scope" in str(error.value.__cause__) + assert encode_checkpoint_value(parent._runner.state.export_state()) == before + + +@pytest.mark.parametrize("streaming", [False, True]) +async def test_response_validation_preserves_noncopyable_request_payload(streaming: bool) -> None: + from threading import Lock + + payload: dict[str, Any] = {"lock": Lock()} + + class Requester(Executor): + @handler + async def start(self, message: str, ctx: WorkflowContext) -> None: + await ctx.request_info(payload, bool) + + @response_handler + async def respond( + self, + original: dict[str, Any], + approved: bool, + ctx: WorkflowContext[Never, str], # type: ignore[valid-type] + ) -> None: + assert original is payload + assert original["lock"] is payload["lock"] + await ctx.yield_output("approved" if approved else "denied") + + workflow = WorkflowBuilder(start_executor=Requester(id="requester")).build() + await workflow.run("start") + [request_id] = (await workflow._runner_context.get_pending_request_info_events()).keys() + if streaming: + result = await workflow.run(responses={request_id: True}, stream=True).get_final_response() + else: + result = await workflow.run(responses={request_id: True}) + assert result.get_outputs() == ["approved"] + assert not await workflow._runner_context.get_pending_request_info_events() + + async def test_propagated_sub_workflow_response_receives_runtime_tools() -> None: """Runtime tools supplied on parent resume reach the child response handler.""" captured_runtime_tools: list[Any] = [] diff --git a/python/packages/core/tests/workflow/test_workflow.py b/python/packages/core/tests/workflow/test_workflow.py index 76a285e6d93..99d75f83dd5 100644 --- a/python/packages/core/tests/workflow/test_workflow.py +++ b/python/packages/core/tests/workflow/test_workflow.py @@ -26,6 +26,7 @@ InProcRunnerContext, Message, ResponseStream, + RunnerContext, WorkflowBuilder, WorkflowCheckpoint, WorkflowCheckpointException, @@ -131,6 +132,146 @@ def test_coerce_request_info_response_rejects_mismatched_type() -> None: _coerce_request_info_response("not-an-int", int, "typed") +class LegacyRunnerContext: + """Delegate the original context contract without the new queue-inspection method.""" + + def __init__(self) -> None: + self.inner = InProcRunnerContext() + + def __getattr__(self, name: str) -> Any: + if name in {"get_queued_response_info_events", "get_queued_messages"}: + raise AttributeError(name) + return getattr(self.inner, name) + + +@pytest.mark.parametrize("streaming", [False, True]) +async def test_response_continuation_supports_legacy_runner_context(streaming: bool) -> None: + context = LegacyRunnerContext() + executor = MockExecutorRequestApproval(id="approver") + workflow = WorkflowBuilder(start_executor=executor).build() + workflow._runner._ctx = cast(RunnerContext, context) + result = await workflow.run(NumberMessage(data=7)) + requests = result.get_request_info_events() + assert len(requests) == 1 + request_id = requests[0].request_id + assert request_id is not None + + if streaming: + events = [event async for event in workflow.run(responses={request_id: ApprovalMessage(True)}, stream=True)] + assert [event.data for event in events if event.type == "output"] == [7] + else: + resumed = await workflow.run(responses={request_id: ApprovalMessage(True)}) + assert resumed.get_outputs() == [7] + assert await context.inner.get_pending_request_info_events() == {} + + +async def test_legacy_runner_context_validates_queued_responses_without_consuming_them() -> None: + class BatchApprover(Executor): + @handler + async def start(self, message: str, ctx: WorkflowContext) -> None: + await ctx.request_info("first", str, request_id="first") + await ctx.request_info("second", str, request_id="second") + + @response_handler + async def respond(self, original_request: str, response: str, ctx: WorkflowContext[Any, str]) -> None: + await ctx.yield_output(f"{original_request}:{response}") + + context = LegacyRunnerContext() + workflow = WorkflowBuilder(start_executor=BatchApprover(id="approver")).build() + workflow._runner._ctx = cast(RunnerContext, context) + await workflow.run("start") + await context.inner.send_request_info_response("first", "yes") + pending_before = await context.inner.get_pending_request_info_events() + validated: list[list[str | None]] = [] + + def validate(requests: Sequence[WorkflowEvent]) -> None: + validated.append([request.request_id for request in requests]) + if len(requests) > 1: + raise ValueError("mixed response batch") + + workflow._response_batch_validator = validate + with pytest.raises(ValueError, match="mixed response batch"): + await workflow.run(responses={"second": "yes"}) + + assert validated == [["first", "second"]] + assert await context.inner.get_pending_request_info_events() == pending_before + queued = await context.inner.get_queued_response_info_events() + assert [request.request_id for request in queued] == ["first"] + + workflow._response_batch_validator = None + resumed = await workflow.run(responses={"second": "yes"}) + assert sorted(resumed.get_outputs()) == ["first:yes", "second:yes"] + + +@pytest.mark.parametrize("legacy", [False, True]) +@pytest.mark.parametrize("mixed", [False, True]) +async def test_cancellation_preflight_checks_queued_reply_origins(legacy: bool, mixed: bool) -> None: + from agent_framework._workflows._const import WORKFLOW_REQUEST_PARENT_CONTEXTS_KEY, WORKFLOW_SESSION_ID_KEY + + cancelled: list[str] = [] + handled: list[str] = [] + delivered: list[tuple[str, str]] = [] + + class Requester(Executor): + @handler + async def start(self, message: str, ctx: WorkflowContext) -> None: + for request_id in ("a", "b") if message == "same" else (message,): + await ctx.request_info(request_id, str, request_id=request_id) + + @response_handler + async def respond(self, request: str, response: str, ctx: WorkflowContext[str]) -> None: + handled.append(request) + await ctx.send_message(f"response:{request}") + + async def _cancel_pending_request(self, request_id: str, ctx: WorkflowContext[Any, Any]) -> None: + cancelled.append(request_id) + await ctx.send_message(f"cancelled:{request_id}") + + class Collector(Executor): + @handler + async def collect(self, message: str, ctx: WorkflowContext) -> None: + delivered.append((message, ctx.get_state(WORKFLOW_SESSION_ID_KEY))) + + requester = Requester(id="requester") + workflow = WorkflowBuilder(start_executor=requester).add_edge(requester, Collector(id="collector")).build() + if legacy: + workflow._runner._ctx = cast(RunnerContext, LegacyRunnerContext()) + context = workflow._runner.context + if mixed: + await workflow.run("a") + await workflow.run("b") + else: + await workflow.run("same") + scope = workflow._runner.state.get(WORKFLOW_SESSION_ID_KEY) + await context.send_request_info_response("a", "yes") + checkpoint = await workflow._runner.build_checkpoint() + await workflow._runner.restore_checkpoint(checkpoint) + before = await context.build_checkpoint( + workflow.name, workflow.graph_signature_hash, workflow._runner.state, None, 0 + ) + pending = await context.get_pending_request_info_events() + origins = workflow._runner.state.get(WORKFLOW_REQUEST_PARENT_CONTEXTS_KEY) + + if mixed: + with pytest.raises(ValueError, match="multiple parent workflow scopes"): + await workflow.cancel_pending_requests(["b"]) + assert cancelled == handled == delivered == [] + assert await context.get_pending_request_info_events() == pending + snapshot = await context.build_checkpoint( + workflow.name, workflow.graph_signature_hash, workflow._runner.state, None, 0 + ) + assert snapshot.messages == before.messages + assert workflow._runner.state.get(WORKFLOW_REQUEST_PARENT_CONTEXTS_KEY) == origins + assert workflow._runner.state.get(WORKFLOW_SESSION_ID_KEY) == scope + else: + await workflow.cancel_pending_requests(["b"]) + assert cancelled == ["b"] + assert handled == ["a"] + assert sorted(delivered) == [("cancelled:b", scope), ("response:a", scope)] + assert not await context.get_pending_request_info_events() + assert workflow._runner.state.get(WORKFLOW_REQUEST_PARENT_CONTEXTS_KEY) == {} + + async def test_public_stream_close_immediately_drains_live_executor_without_gc() -> None: started, cancelled, closed, release = (asyncio.Event() for _ in range(4)) cleanup: list[str] = [] diff --git a/python/packages/core/tests/workflow/test_workflow_agent.py b/python/packages/core/tests/workflow/test_workflow_agent.py index 93a1db2a807..1a6a8bf2834 100644 --- a/python/packages/core/tests/workflow/test_workflow_agent.py +++ b/python/packages/core/tests/workflow/test_workflow_agent.py @@ -1,5 +1,6 @@ # Copyright (c) Microsoft. All rights reserved. +import asyncio import uuid from collections.abc import Awaitable, Sequence from dataclasses import dataclass @@ -18,6 +19,7 @@ Executor, FinishReason, HistoryProvider, + InMemoryCheckpointStorage, InMemoryHistoryProvider, Message, ResponseStream, @@ -27,13 +29,16 @@ UsageDetails, WorkflowAgent, WorkflowBuilder, + WorkflowCheckpointException, WorkflowContext, WorkflowEvent, + WorkflowException, executor, handler, response_handler, ) from agent_framework._types import _get_operation_state, _set_operation_state +from agent_framework._workflows._const import WORKFLOW_AGENT_RUN_KEY, WORKFLOW_AGENT_SESSION_ID_KEY from agent_framework._workflows._typing_utils import deserialize_type @@ -194,6 +199,176 @@ async def handle_message( class TestWorkflowAgent: """Test cases for WorkflowAgent end-to-end functionality.""" + @pytest.mark.parametrize("streaming", [False, True]) + @pytest.mark.parametrize("explicit_session", [False, True]) + async def test_direct_fresh_runs_clear_agent_session_id(self, streaming: bool, explicit_session: bool) -> None: + workflow = WorkflowBuilder(start_executor=SimpleExecutor(id="start", response_text="Accepted")).build() + agent = workflow.as_agent() + session = AgentSession() if explicit_session else None + await agent.run("Agent turn", session=session) + assert workflow._runner.state.get(WORKFLOW_AGENT_SESSION_ID_KEY) == ( + session.session_id if session is not None else None + ) + assert workflow._runner.state.get(WORKFLOW_AGENT_RUN_KEY) is True + + for text in ("First direct run", "Second direct run"): + messages = [Message(role="user", contents=[Content.from_text(text)])] + if streaming: + await workflow.run(messages, stream=True).get_final_response() + else: + await workflow.run(messages) + assert workflow._runner.state.get(WORKFLOW_AGENT_SESSION_ID_KEY) is None + assert workflow._runner.state.get(WORKFLOW_AGENT_RUN_KEY) is False + + @pytest.mark.parametrize("streaming", [False, True]) + async def test_direct_continuations_preserve_agent_session_id(self, streaming: bool) -> None: + workflow = WorkflowBuilder(start_executor=RequestingExecutor(id="request")).build() + storage = InMemoryCheckpointStorage() + session = AgentSession() + await workflow.as_agent().run("Agent turn", session=session, checkpoint_storage=storage) + checkpoint_id = workflow.get_last_checkpoint_id() + assert checkpoint_id is not None + pending_requests = await workflow._runner_context.get_pending_request_info_events() + assert len(pending_requests) == 1 + responses = {request_id: "Answer" for request_id in pending_requests} + + if streaming: + await workflow.run(responses=responses, stream=True).get_final_response() + assert workflow._runner.state.get(WORKFLOW_AGENT_SESSION_ID_KEY) == session.session_id + assert workflow._runner.state.get(WORKFLOW_AGENT_RUN_KEY) is True + await workflow.run( + checkpoint_id=checkpoint_id, checkpoint_storage=storage, stream=True + ).get_final_response() + else: + await workflow.run(responses=responses) + assert workflow._runner.state.get(WORKFLOW_AGENT_SESSION_ID_KEY) == session.session_id + assert workflow._runner.state.get(WORKFLOW_AGENT_RUN_KEY) is True + await workflow.run(checkpoint_id=checkpoint_id, checkpoint_storage=storage) + assert workflow._runner.state.get(WORKFLOW_AGENT_SESSION_ID_KEY) == session.session_id + assert workflow._runner.state.get(WORKFLOW_AGENT_RUN_KEY) is True + + @pytest.mark.parametrize("streaming", [False, True]) + @pytest.mark.parametrize("prior_agent_run", [False, True]) + async def test_rejected_checkpoint_preserves_agent_context(self, streaming: bool, prior_agent_run: bool) -> None: + workflow = WorkflowBuilder(start_executor=SimpleExecutor(id="start", response_text="Accepted")).build() + agent = workflow.as_agent() + storage = InMemoryCheckpointStorage() + await agent.run("Checkpoint turn", session=AgentSession(), checkpoint_storage=storage) + checkpoint_id = workflow.get_last_checkpoint_id() + assert checkpoint_id is not None + checkpoint = await storage.load(checkpoint_id) + assert checkpoint is not None + checkpoint.graph_signature_hash = "different-graph" + await storage.save(checkpoint) + + if prior_agent_run: + await agent.run("Live turn", session=AgentSession()) + else: + await workflow.run([Message(role="user", contents=[Content.from_text("Live turn")])]) + live_state = workflow._runner.state.export_state() + + with pytest.raises(WorkflowCheckpointException, match="graph has changed"): + if streaming: + await agent.run( + session=AgentSession(), + checkpoint_id=checkpoint_id, + checkpoint_storage=storage, + stream=True, + ).get_final_response() + else: + await agent.run(session=AgentSession(), checkpoint_id=checkpoint_id, checkpoint_storage=storage) + + assert workflow._runner.state.export_state() == live_state + + @pytest.mark.parametrize("streaming", [False, True]) + @pytest.mark.parametrize("fresh_input", [False, True]) + async def test_checkpoint_restore_uses_checkpoint_agent_context(self, streaming: bool, fresh_input: bool) -> None: + workflow = WorkflowBuilder(start_executor=SimpleExecutor(id="start", response_text="Accepted")).build() + agent = workflow.as_agent() + storage = InMemoryCheckpointStorage() + checkpoint_session = AgentSession() + await agent.run("Checkpoint turn", session=checkpoint_session, checkpoint_storage=storage) + checkpoint_id = workflow.get_last_checkpoint_id() + assert checkpoint_id is not None + await agent.run("Live turn", session=AgentSession()) + caller_session = AgentSession() + message = "Fresh turn" if fresh_input else None + + if streaming: + await agent.run( + message, + session=caller_session, + checkpoint_id=checkpoint_id, + checkpoint_storage=storage, + stream=True, + ).get_final_response() + else: + await agent.run(message, session=caller_session, checkpoint_id=checkpoint_id, checkpoint_storage=storage) + + expected_session = caller_session if fresh_input else checkpoint_session + assert workflow._runner.state.get(WORKFLOW_AGENT_SESSION_ID_KEY) == expected_session.session_id + assert workflow._runner.state.get(WORKFLOW_AGENT_RUN_KEY) is True + + @pytest.mark.parametrize("streaming", [False, True]) + @pytest.mark.parametrize("rejected_streaming", [False, True]) + @pytest.mark.parametrize("first_agent", [False, True]) + async def test_rejected_concurrent_run_preserves_agent_session_id( + self, streaming: bool, rejected_streaming: bool, first_agent: bool, monkeypatch: pytest.MonkeyPatch + ) -> None: + started = asyncio.Event() + release = asyncio.Event() + workflow = WorkflowBuilder(start_executor=SimpleExecutor(id="start", response_text="Accepted")).build() + agent = workflow.as_agent() + session = AgentSession() + rejected_session = AgentSession() + original_has_messages = workflow._runner.context.has_messages + + async def gated_has_messages() -> bool: + started.set() + await release.wait() + return await original_has_messages() + + monkeypatch.setattr(workflow._runner.context, "has_messages", gated_has_messages) + + async def run_agent(agent_session: AgentSession, stream: bool) -> AgentResponse: + if stream: + return await agent.run("Hello", session=agent_session, stream=True).get_final_response() + return await agent.run("Hello", session=agent_session) + + async def run_first() -> AgentResponse: + if first_agent: + return await run_agent(session, streaming) + messages = [Message(role="user", contents=[Content.from_text("Hello")])] + if streaming: + result = await workflow.run(messages, stream=True).get_final_response() + else: + result = await workflow.run(messages) + [output] = result.get_outputs() + assert isinstance(output, AgentResponse) + return output + + expected_session_id = session.session_id if first_agent else None + task = asyncio.create_task(run_first()) + try: + await started.wait() + assert workflow._runner.state.get(WORKFLOW_AGENT_SESSION_ID_KEY) == expected_session_id + assert workflow._runner.state.get(WORKFLOW_AGENT_RUN_KEY) is first_agent + with pytest.raises(WorkflowException, match="Workflow is already running"): + await run_agent(rejected_session, rejected_streaming) + with pytest.raises(WorkflowException, match="Workflow is already running"): + await workflow.run([Message(role="user", contents=[Content.from_text("Rejected")])]) + assert workflow._runner.state.get(WORKFLOW_AGENT_SESSION_ID_KEY) == expected_session_id + assert workflow._runner.state.get(WORKFLOW_AGENT_RUN_KEY) is first_agent + finally: + release.set() + result = await task + + assert result.text == "Accepted: Hello" + assert workflow._runner.state.get(WORKFLOW_AGENT_SESSION_ID_KEY) == expected_session_id + assert workflow._runner.state.get(WORKFLOW_AGENT_RUN_KEY) is first_agent + await run_agent(rejected_session, rejected_streaming) + assert workflow._runner.state.get(WORKFLOW_AGENT_SESSION_ID_KEY) == rejected_session.session_id + async def test_end_to_end_basic_workflow(self): """Test basic end-to-end workflow execution with 2 executors emitting AgentResponse.""" # Create workflow with two executors diff --git a/python/packages/declarative/AGENTS.md b/python/packages/declarative/AGENTS.md index 059c1c605d7..19607e06674 100644 --- a/python/packages/declarative/AGENTS.md +++ b/python/packages/declarative/AGENTS.md @@ -23,7 +23,31 @@ series because response caching and default-header integration use its internals ## MCP Handler Lifetimes -`DefaultMCPToolHandler` caches/coalesces sessions only without a `client_provider`. +`DefaultMCPToolHandler` caches/coalesces sessions only without a `client_provider` +and with a non-empty `workflow_session_id`. Missing or empty IDs use the same +per-invocation lifecycle as provider-backed calls, without cache reuse or coalescing. +Cache identity includes a framework-owned workflow session ID in addition to +endpoint, label, connection, and headers, so separate fresh runs do not share a +stateful MCP protocol session while continuations and checkpoint restores do. +Direct fresh runs reset this scope even for message-list inputs; only +framework-marked WorkflowAgent turns retain the agent continuation scope. +Executors manually composed with Core `WorkflowBuilder` use Core's active run scope, +including inherited child scopes, rather than a cached declarative scope. +Run scopes are selected before the initial entry checkpoint and preserved when +Entry replays. Sessionless WorkflowAgent turns select their own checkpointed scope, +independent of explicit agent sessions, direct runs, and restored approval scopes. +Response-only continuations retain the originating workflow scope even when supplied +a new `AgentSession` after checkpoint restoration. +Nested workflow bindings derive their scopes from the checkpointed parent scope and +binding ID, so replaying a parent entry checkpoint preserves child scopes without +sharing them across sibling bindings. +Forwarded child requests also checkpoint the originating context at each parent +boundary. Interleaved continuations restore those ancestor scopes before forwarding +outputs to subsequent bindings; responses from different parent scopes must be +submitted separately. +Ordinary Core `ctx.request_info()` requests before a child binding retain and +restore their origins too, including when a queued response is replayed. +Approval continuations keep their originating active scope until the next run with new input. With a provider, every invocation (including `tools/list`) gets a fresh tool/session, even if the provider returns `None` or a shared HTTP client. Invocation cleanup closes the session and any internally owned fallback client, never caller-owned HTTP clients. @@ -32,6 +56,13 @@ from an active invocation's context (including inherited child tasks and cleanup raises `RuntimeError` before changing handler state; shutdown must run outside that context or after invocation completion. +For cache-managed scopes, `cache_max_size` bounds live sessions including connecting and +closing clients, not only idle cache entries. At capacity, the least-recently-used +idle client is fully closed before its replacement connects. If every slot is +busy, new scopes receive an explicit capacity error; same-scope cache reuse +remains available. +Provider-backed and unscoped per-invocation sessions do not use this cache bound. + Per-invocation sessions intentionally add connection overhead and lose server session continuity between invocations; shared session ownership requires an explicitly scoped custom `MCPToolHandler`. @@ -43,6 +74,13 @@ workflow-local HMAC key held separately in trusted host checkpoint state. Only the opaque binding and header names enter the approval payload; raw headers are not checkpointed. Changed or unverifiable headers produce a replacement request for the same pinned operation, with a fresh request ID and no dispatch. +Pending requests also retain their originating workflow session ID across later +fresh runs and checkpoint restores. Resuming an MCP approval restores that scope +for downstream actions instead of inheriting a later run's MCP session. +Factory-built workflows reject response batches containing MCP approvals from +different originating scopes before dispatch; resume each scope in a separate run. +Response policies stage only the legacy scope-field updates until all nested +policies pass; live request payloads are not deep-copied for preflight. Fresh executors verify unchanged approvals using the checkpointed key; legacy requests or missing verification state require reapproval for non-empty headers. Custom handlers remain responsible for identity changes @@ -57,6 +95,37 @@ arguments. Explicit caller tool arguments in options are preserved. - **`ExternalInputRequest`** / **`ExternalInputResponse`** - Human-in-the-loop support - **`AgentExternalInputRequest`** / **`AgentExternalInputResponse`** - Agent-level input requests +`Question`, `RequestExternalInput`, function approvals, MCP approvals, and agent +external-input requests retain their originating workflow scope across fresh runs +and checkpoint restores. Responses restore that scope before updating state or +continuing execution. Agent external-input requests also retain their own +`ExternalLoopState` rather than overwriting another pending request's configuration. +Factory-built workflows reject response batches spanning different scopes across +all four request families before consuming any responses. + +Legacy requests without an originating scope are rejected by default with recovery +instructions. Hosts can opt in through +`WorkflowFactory(migrate_legacy_request_scopes=True)`: each submitted legacy request +receives a stable `legacy-request:` scope. Migration never +infers an origin from current state or recovers the old MCP session. Resume requests +separately, preserve their framework IDs, and checkpoint the migrated continuation. +The same validation and migration run before restoring ordinary queued checkpoint responses, +including embedded child checkpoints. The complete nested tree is validated on a +detached snapshot before any parent or child state, executor state, or queue is +replaced; a rejected descendant does not partially restore earlier siblings. +Newly submitted responses are validated together with ordinary queued responses +before any pending request is consumed. +Wrapped child replies instead select and validate their actual routes immediately +before delivery, after earlier handlers and checkpoint restoration. All selected +children must validate before that reply is delivered. A later rejected reply does +not roll back earlier continuations or restored state; recover a partially executed +failed run from a known-good checkpoint rather than retrying its remaining queue. +The accepted parent scope persists across supersteps and mid-run checkpoint replay +until all in-flight messages drain; later response-only runs may select other origins. +Rejected mixed batches do not mutate request origins. Malformed stored origins +are rejected even with migration enabled. Legacy agent loop configuration is read +from the old checkpoint slot; newly emitted requests retain their own snapshot. + ## Usage ```python diff --git a/python/packages/declarative/README.md b/python/packages/declarative/README.md index 0f502ccd003..35a0e66418e 100644 --- a/python/packages/declarative/README.md +++ b/python/packages/declarative/README.md @@ -72,7 +72,10 @@ on HTTPX internals. Compatibility tests cover HTTPX 0.27.0 and 0.28.1. ## MCP handler session lifetime -`DefaultMCPToolHandler()` reuses MCP sessions through a bounded LRU cache. +`DefaultMCPToolHandler()` reuses MCP sessions through a bounded LRU cache only +when the invocation has a non-empty `workflow_session_id`. Direct calls with a +missing or empty ID instead create and close a fresh session per invocation, +including `tools/list`; concurrent unscoped calls never share a connection. When configured with `client_provider`, it instead invokes the provider and creates a fresh MCP tool/session for every invocation, including `tools/list`. Sessions are never reused in this mode, even if the provider returns `None` @@ -81,13 +84,57 @@ or the same `httpx.AsyncClient` instance. Each invocation closes its tool/session on success, failure, or cancellation. Internally created fallback HTTP clients are also closed; caller-supplied HTTP clients remain caller-owned and are never closed by the handler. Calling -`aclose()` rejects new invocations and waits for active provider-backed +`aclose()` rejects new invocations and waits for active provider-backed or unscoped invocations to finish cleaning up. Calling it from an active invocation's context (including provider callbacks, inherited child tasks, and cleanup) raises `RuntimeError` before changing handler state, rather than waiting on itself. Close the handler outside that context or after the invocation completes. This intentionally incurs connection and initialization overhead and does not -retain server session state between provider-backed invocations. Applications +retain server session state between provider-backed or unscoped invocations. +`cache_max_size` bounds cache-managed clients, not these per-invocation sessions. Applications that require shared session ownership must implement an explicitly scoped custom `MCPToolHandler`; there is no provider-backed session-cache opt-in. + +## Resuming pending workflow requests + +Questions, external-input requests, function approvals, MCP approvals, and agent +external-input requests retain the workflow scope in which they were issued. +Responses restore that scope even after another fresh run or a checkpoint restore. +Submit responses from different scopes in separate runs; mixed-scope batches are +rejected before any request is consumed. + +Older checkpoints may contain requests without an originating scope. By default, +resuming these requests raises an error without consuming them. To recover, the +host can explicitly opt in when recreating the workflow: + +```python +from agent_framework.declarative import WorkflowFactory + +factory = WorkflowFactory( + checkpoint_storage=storage, + mcp_tool_handler=mcp_handler, + migrate_legacy_request_scopes=True, +) +workflow = factory.create_workflow_from_yaml_path("workflow.yaml") +await workflow.run(checkpoint_id=checkpoint_id) +await workflow.run(responses={pending_request_id: response}) +``` + +Migration gives each legacy request a stable, separately isolated scope derived +from its framework request ID. It never reuses the latest run's scope and cannot +recover the old MCP session or its server-side history. Preserve request IDs, +resume legacy requests individually, and retain the resulting checkpoint. Newly +emitted requests carry the migrated scope, so subsequent restores no longer +require the opt-in for those requests. Any remaining unmigrated requests still do. +The same policy applies when a checkpoint already contains queued responses: +restore validates and migrates their original requests before restoring workflow +state or replaying the continuation. Nested workflows preflight the entire embedded +checkpoint tree using each child's own policy before replacing any live state, +executor state, or queues. Preflight rejection leaves the running workflows and supplied +checkpoint unchanged, including staged migration metadata. +Responses submitted alongside queued responses must belong to the same scope; +mixed batches leave pending requests unconsumed. +Invalid stored scopes are rejected even with migration enabled. Header verification +and any required MCP reapproval remain unchanged. Hosts that require the old +server-side session must restart the workflow instead of migrating it. diff --git a/python/packages/declarative/agent_framework_declarative/_workflows/_declarative_base.py b/python/packages/declarative/agent_framework_declarative/_workflows/_declarative_base.py index 7338aa62260..fbceec45fc4 100644 --- a/python/packages/declarative/agent_framework_declarative/_workflows/_declarative_base.py +++ b/python/packages/declarative/agent_framework_declarative/_workflows/_declarative_base.py @@ -1080,6 +1080,15 @@ def _get_state(self, state: State) -> DeclarativeWorkflowState: """Get the declarative workflow state wrapper.""" return DeclarativeWorkflowState(state, env_config=self._declarative_env_config) + @staticmethod + def _is_agent_continuation(state: DeclarativeWorkflowState, trigger: Any) -> bool: + """Return whether an agent-facing message list continues initialized state.""" + return ( + isinstance(trigger, list) + and all(isinstance(message, Message) for message in trigger) # pyright: ignore[reportUnknownVariableType] + and state.is_initialized() + ) + async def _ensure_state_initialized( self, ctx: WorkflowContext[Any, Any], @@ -1132,7 +1141,7 @@ async def _ensure_state_initialized( # reuse this workflow instance across unrelated conversations. # State left on the instance is not evidence of session ownership. # Explicit checkpoint/HIL restores use the runner's resume path. - is_continuation = not isinstance(trigger, Message) and state.is_initialized() + is_continuation = self._is_agent_continuation(state, trigger) # Locate the trailing user message in the trigger. last_user_index = -1 diff --git a/python/packages/declarative/agent_framework_declarative/_workflows/_declarative_builder.py b/python/packages/declarative/agent_framework_declarative/_workflows/_declarative_builder.py index 9e7a18711e6..a298170f3b9 100644 --- a/python/packages/declarative/agent_framework_declarative/_workflows/_declarative_builder.py +++ b/python/packages/declarative/agent_framework_declarative/_workflows/_declarative_builder.py @@ -14,12 +14,17 @@ from __future__ import annotations import logging +from collections.abc import Callable, Sequence +from functools import partial from typing import Any, cast from agent_framework import ( Workflow, WorkflowBuilder, + WorkflowEvent, ) +from agent_framework._workflows._const import WORKFLOW_AGENT_RUN_KEY +from agent_framework._workflows._state import State from ._declarative_base import ( ConditionResult, @@ -28,7 +33,7 @@ LoopIterationResult, ) from ._errors import DeclarativeWorkflowError -from ._executors_agents import AGENT_ACTION_EXECUTORS, InvokeAzureAgentExecutor +from ._executors_agents import AGENT_ACTION_EXECUTORS, AgentExternalInputRequest, InvokeAzureAgentExecutor from ._executors_basic import BASIC_ACTION_EXECUTORS from ._executors_control_flow import ( CONTROL_FLOW_EXECUTORS, @@ -39,12 +44,19 @@ IfConditionEvaluatorExecutor, JoinExecutor, ) -from ._executors_external_input import EXTERNAL_INPUT_EXECUTORS +from ._executors_external_input import ( + EXTERNAL_INPUT_EXECUTORS, + ExternalInputRequest, +) from ._executors_http import HTTP_ACTION_EXECUTORS, HttpRequestActionExecutor -from ._executors_mcp import MCP_ACTION_EXECUTORS, InvokeMcpToolActionExecutor -from ._executors_tools import TOOL_ACTION_EXECUTORS, InvokeFunctionToolExecutor +from ._executors_mcp import ( + MCP_ACTION_EXECUTORS, + InvokeMcpToolActionExecutor, + MCPToolApprovalRequest, +) +from ._executors_tools import TOOL_ACTION_EXECUTORS, InvokeFunctionToolExecutor, ToolApprovalRequest from ._http_handler import HttpRequestHandler -from ._mcp_handler import MCPToolHandler +from ._mcp_handler import MCPToolHandler, activate_workflow_session_id, get_request_workflow_session_id logger = logging.getLogger(__name__) @@ -108,6 +120,40 @@ } +def _validate_response_batch( + requests: Sequence[WorkflowEvent[Any]], *, migrate_legacy_request_scopes: bool = False +) -> Callable[[], None]: + scoped_requests: list[ + tuple[ExternalInputRequest | MCPToolApprovalRequest | ToolApprovalRequest | AgentExternalInputRequest, str] + ] = [] + for request in requests: + data = request.data + if not isinstance( + data, (ExternalInputRequest, MCPToolApprovalRequest, ToolApprovalRequest, AgentExternalInputRequest) + ): + continue + scope = data.workflow_session_id + if scope is None and migrate_legacy_request_scopes: + if not isinstance(request.request_id, str) or not request.request_id: + raise DeclarativeWorkflowError("Cannot migrate a legacy request without a valid request ID.") + scope = f"legacy-request:{request.request_id}" + scoped_requests.append((data, get_request_workflow_session_id(scope))) + scopes = {scope for _, scope in scoped_requests} + if len(scopes) > 1: + raise DeclarativeWorkflowError("Responses from different workflow sessions must be submitted in separate runs.") + + def apply_migration() -> None: + for data, scope in scoped_requests: + if data.workflow_session_id is None: + data.workflow_session_id = scope + + return apply_migration + + +def _initialize_mcp_run_state(state: State) -> None: + activate_workflow_session_id(state, reset_unscoped=True, is_agent_run=state.get(WORKFLOW_AGENT_RUN_KEY) is True) + + class DeclarativeWorkflowBuilder: """Builds a Workflow graph from declarative YAML actions. @@ -141,6 +187,7 @@ def __init__( http_request_handler: HttpRequestHandler | None = None, mcp_tool_handler: MCPToolHandler | None = None, env_config: DeclarativeEnvConfig | None = None, + migrate_legacy_request_scopes: bool = False, ): """Initialize the builder. @@ -163,6 +210,8 @@ def __init__( how the ``Env`` PowerFx symbol is populated for every executor built by this builder. Defaults to an empty configuration (``Env`` not exposed). + migrate_legacy_request_scopes: Opt in to assigning legacy pending requests + stable, isolated scopes instead of rejecting missing origins. """ self._yaml_def = yaml_definition self._workflow_id = workflow_id or yaml_definition.get("name", "declarative_workflow") @@ -178,6 +227,7 @@ def __init__( self._http_request_handler = http_request_handler self._mcp_tool_handler = mcp_tool_handler self._env_config: DeclarativeEnvConfig = env_config if env_config is not None else DeclarativeEnvConfig() + self._migrate_legacy_request_scopes = migrate_legacy_request_scopes # Resolve max_iterations: explicit arg > YAML maxTurns > core default resolved = max_iterations if max_iterations is not None else yaml_definition.get("maxTurns") if resolved is not None and (not isinstance(resolved, int) or resolved <= 0): @@ -240,7 +290,12 @@ def build(self) -> Workflow: if isinstance(executor, DeclarativeActionExecutor): executor.set_declarative_env_config(self._env_config) - return builder.build() + workflow = builder.build() + workflow._response_batch_validator = partial( # pyright: ignore[reportPrivateUsage] + _validate_response_batch, migrate_legacy_request_scopes=self._migrate_legacy_request_scopes + ) + workflow._initial_state_initializer = _initialize_mcp_run_state # pyright: ignore[reportPrivateUsage] + return workflow def _collect_explicit_ids(self, actions: list[dict[str, Any]] | None) -> None: """Reserve authored executor IDs, excluding virtual structures and skipped actions.""" diff --git a/python/packages/declarative/agent_framework_declarative/_workflows/_executors_agents.py b/python/packages/declarative/agent_framework_declarative/_workflows/_executors_agents.py index 30de4b7eb00..c0d6cc4acc8 100644 --- a/python/packages/declarative/agent_framework_declarative/_workflows/_executors_agents.py +++ b/python/packages/declarative/agent_framework_declarative/_workflows/_executors_agents.py @@ -12,7 +12,8 @@ - AutoSend behavior control """ -import contextlib +from __future__ import annotations + import json import logging import uuid @@ -34,6 +35,11 @@ DeclarativeActionExecutor, DeclarativeWorkflowState, ) +from ._mcp_handler import ( + get_or_create_workflow_session_id, + get_request_workflow_session_id, + restore_workflow_session_id, +) logger = logging.getLogger(__name__) @@ -432,6 +438,8 @@ async def on_request(request: AgentExternalInputRequest) -> ExternalInputRespons iteration: int = 0 messages: list[Message] = field(default_factory=lambda: cast(list[Message], [])) function_calls: list[Content] = field(default_factory=lambda: cast(list[Content], [])) + workflow_session_id: str | None = None + loop_state: ExternalLoopState | None = None @dataclass @@ -469,8 +477,8 @@ class AgentExternalInputResponse: class ExternalLoopState: """State saved for external loop resumption. - Stored in workflow state to allow the response_handler to - continue the loop with the same configuration. + Retained by each pending request so interleaved continuations + preserve their own configuration. """ agent_name: str @@ -1006,8 +1014,6 @@ async def handle_action( messages_path=messages_path, max_iterations=max_iterations, ) - ctx.state.set(EXTERNAL_LOOP_STATE_KEY, loop_state) - # Emit request for external input - workflow will yield here request = AgentExternalInputRequest( request_id=str(uuid.uuid4()), @@ -1016,6 +1022,8 @@ async def handle_action( iteration=iteration, messages=all_messages, function_calls=tool_calls, + workflow_session_id=get_or_create_workflow_session_id(ctx.state), + loop_state=loop_state, ) logger.info(f"InvokeAzureAgent: yielding for external input (iteration {iteration})") await ctx.request_info(request, AgentExternalInputResponse) @@ -1037,14 +1045,17 @@ async def handle_external_input_response( This is called when the workflow resumes after yielding for external input. It continues the agent invocation loop with the user's new input. """ + restore_workflow_session_id(ctx.state, get_request_workflow_session_id(original_request.workflow_session_id)) logger.debug( "handle_external_input_response: resuming with user_input='%s'", response.user_input[:100] if response.user_input else None, ) state = self._get_state(ctx.state) - # Retrieve saved loop state - loop_state: ExternalLoopState | None = ctx.state.get(EXTERNAL_LOOP_STATE_KEY) + # Older checkpoints store the loop configuration in a single shared slot. + loop_state = original_request.loop_state + if loop_state is None: + loop_state = ctx.state.get(EXTERNAL_LOOP_STATE_KEY) if loop_state is None: logger.error("InvokeAzureAgent: external loop state not found, cannot resume") await ctx.send_message(ActionComplete()) @@ -1081,9 +1092,7 @@ async def handle_external_input_response( ) if not should_continue: - # User input caused loop to exit - clean up and complete - with contextlib.suppress(KeyError): - ctx.state.delete(EXTERNAL_LOOP_STATE_KEY) + # User input caused loop to exit. await ctx.send_message(ActionComplete()) return @@ -1141,8 +1150,6 @@ async def handle_external_input_response( "InvokeAzureAgent: external loop exited due to condition=false " "(sending ActionComplete to continue workflow)" ) - with contextlib.suppress(KeyError): - ctx.state.delete(EXTERNAL_LOOP_STATE_KEY) await ctx.send_message(ActionComplete()) return @@ -1150,8 +1157,6 @@ async def handle_external_input_response( if iteration < max_iterations: # Update loop state for next iteration loop_state.iteration = iteration + 1 - ctx.state.set(EXTERNAL_LOOP_STATE_KEY, loop_state) - # Emit another request for external input request = AgentExternalInputRequest( request_id=str(uuid.uuid4()), @@ -1160,6 +1165,8 @@ async def handle_external_input_response( iteration=iteration, messages=all_messages, function_calls=tool_calls, + workflow_session_id=get_or_create_workflow_session_id(ctx.state), + loop_state=loop_state, ) logger.info(f"InvokeAzureAgent: yielding for external input (iteration {iteration})") await ctx.request_info(request, AgentExternalInputResponse) @@ -1167,10 +1174,6 @@ async def handle_external_input_response( logger.warning(f"InvokeAzureAgent: external loop exceeded max iterations ({max_iterations})") - # Loop complete - clean up and send completion - with contextlib.suppress(KeyError): - ctx.state.delete(EXTERNAL_LOOP_STATE_KEY) - await ctx.send_message(ActionComplete()) diff --git a/python/packages/declarative/agent_framework_declarative/_workflows/_executors_external_input.py b/python/packages/declarative/agent_framework_declarative/_workflows/_executors_external_input.py index 44bd6aad8da..edf8b78eaa8 100644 --- a/python/packages/declarative/agent_framework_declarative/_workflows/_executors_external_input.py +++ b/python/packages/declarative/agent_framework_declarative/_workflows/_executors_external_input.py @@ -21,6 +21,11 @@ ActionComplete, DeclarativeActionExecutor, ) +from ._mcp_handler import ( + get_or_create_workflow_session_id, + get_request_workflow_session_id, + restore_workflow_session_id, +) def _get_prompt_text(action_def: dict[str, Any], primary_key: str, fallback_key: str) -> Any: @@ -67,12 +72,14 @@ class ExternalInputRequest: ``RequestExternalInput`` action (e.g. ``"approval"``) and it is propagated unchanged. metadata: Additional context (choices, output_property, timeout, etc.). + workflow_session_id: Framework-owned scope of the workflow that issued the request. """ request_id: str message: str request_type: str = "external" metadata: dict[str, Any] = field(default_factory=dict) # type: ignore + workflow_session_id: str | None = None @dataclass @@ -144,6 +151,7 @@ async def handle_action( "allow_free_text": allow_free_text, "default_value": default_value, }, + workflow_session_id=get_or_create_workflow_session_id(ctx.state), ), ExternalInputResponse, ) @@ -156,6 +164,7 @@ async def handle_response( ctx: WorkflowContext[ActionComplete], ) -> None: """Handle the user's response to the question.""" + restore_workflow_session_id(ctx.state, get_request_workflow_session_id(original_request.workflow_session_id)) state = self._get_state(ctx.state) output_property = original_request.metadata.get("output_property", "Local.answer") @@ -212,6 +221,7 @@ async def handle_action( message=str(evaluated_message), request_type=request_type, metadata=request_metadata, + workflow_session_id=get_or_create_workflow_session_id(ctx.state), ), ExternalInputResponse, ) @@ -224,6 +234,7 @@ async def handle_response( ctx: WorkflowContext[ActionComplete], ) -> None: """Handle the external input response.""" + restore_workflow_session_id(ctx.state, get_request_workflow_session_id(original_request.workflow_session_id)) state = self._get_state(ctx.state) output_property = original_request.metadata.get("output_property", "Local.externalInput") diff --git a/python/packages/declarative/agent_framework_declarative/_workflows/_executors_mcp.py b/python/packages/declarative/agent_framework_declarative/_workflows/_executors_mcp.py index 28ff521b0c4..53420b251ab 100644 --- a/python/packages/declarative/agent_framework_declarative/_workflows/_executors_mcp.py +++ b/python/packages/declarative/agent_framework_declarative/_workflows/_executors_mcp.py @@ -48,7 +48,14 @@ DeclarativeWorkflowState, ) from ._executors_tools import ToolApprovalResponse -from ._mcp_handler import MCPToolHandler, MCPToolInvocation, MCPToolResult +from ._mcp_handler import ( + MCPToolHandler, + MCPToolInvocation, + MCPToolResult, + get_or_create_workflow_session_id, + get_request_workflow_session_id, + restore_workflow_session_id, +) __all__ = [ "MCP_ACTION_EXECUTORS", @@ -81,6 +88,8 @@ class MCPToolApprovalRequest: (e.g. ``conversation_id``) for use by the resume handler. header_binding: Opaque binding of the reviewed headers. The verification key is retained separately in trusted workflow state. + workflow_session_id: Framework-owned MCP session identifier from the + run that originated the approval request. """ request_id: str @@ -92,6 +101,7 @@ class MCPToolApprovalRequest: connection_name: str | None = None metadata: dict[str, Any] = field(default_factory=lambda: {}) header_binding: str | None = None + workflow_session_id: str | None = None # --------------------------------------------------------------------------- @@ -236,6 +246,7 @@ async def handle_action( arguments=arguments, headers=headers, connection_name=connection_name, + workflow_session_id=get_or_create_workflow_session_id(ctx.state), ) if require_approval: await self._request_approval( @@ -301,6 +312,7 @@ async def _request_approval( header_binding=( self._bind_headers(ctx, request_id, invocation.headers, create_key=True) if invocation.headers else None ), + workflow_session_id=invocation.workflow_session_id, ) logger.info("%s: requesting approval for MCP tool '%s'", self.__class__.__name__, invocation.tool_name) await ctx.request_info(request, ToolApprovalResponse, request_id=request_id) @@ -313,6 +325,8 @@ async def handle_approval_response( ctx: WorkflowContext[ActionComplete, str], ) -> None: """Resume the pinned operation only with the reviewed header context.""" + workflow_session_id = get_request_workflow_session_id(original_request.workflow_session_id) + restore_workflow_session_id(ctx.state, workflow_session_id) state = self._get_state(ctx.state) tool_name = original_request.tool_name @@ -342,6 +356,7 @@ async def handle_approval_response( arguments=original_request.arguments, headers=self._evaluate_headers(state, self._action_def.get("headers")), connection_name=getattr(original_request, "connection_name", None), + workflow_session_id=workflow_session_id, ) if invocation.headers or original_request.header_names: binding = getattr(original_request, "header_binding", None) diff --git a/python/packages/declarative/agent_framework_declarative/_workflows/_executors_tools.py b/python/packages/declarative/agent_framework_declarative/_workflows/_executors_tools.py index 16c2f3a025c..ab111cdd052 100644 --- a/python/packages/declarative/agent_framework_declarative/_workflows/_executors_tools.py +++ b/python/packages/declarative/agent_framework_declarative/_workflows/_executors_tools.py @@ -34,6 +34,11 @@ DeclarativeWorkflowState, ) from ._executors_agents import TOOL_REGISTRY_KEY +from ._mcp_handler import ( + get_or_create_workflow_session_id, + get_request_workflow_session_id, + restore_workflow_session_id, +) logger = logging.getLogger(__name__) @@ -61,11 +66,13 @@ class ToolApprovalRequest: request_id: Unique identifier for this approval request. function_name: Evaluated function name to be invoked. arguments: Evaluated arguments to be passed to the function. + workflow_session_id: Framework-owned scope of the workflow that issued the request. """ request_id: str function_name: str arguments: dict[str, Any] + workflow_session_id: str | None = None @dataclass @@ -511,6 +518,7 @@ async def handle_action( request_id=request_id, function_name=function_name, arguments=arguments, + workflow_session_id=get_or_create_workflow_session_id(ctx.state), ) logger.info(f"{self.__class__.__name__}: requesting approval for '{function_name}'") await ctx.request_info(request, ToolApprovalResponse, request_id=request_id) @@ -547,6 +555,7 @@ async def handle_approval_response( Rejected calls store the rejection and complete without evaluating ``autoSend``. """ + restore_workflow_session_id(ctx.state, get_request_workflow_session_id(original_request.workflow_session_id)) state = self._get_state(ctx.state) function_name = original_request.function_name diff --git a/python/packages/declarative/agent_framework_declarative/_workflows/_factory.py b/python/packages/declarative/agent_framework_declarative/_workflows/_factory.py index a8329a59734..b80062659b0 100644 --- a/python/packages/declarative/agent_framework_declarative/_workflows/_factory.py +++ b/python/packages/declarative/agent_framework_declarative/_workflows/_factory.py @@ -98,6 +98,7 @@ def __init__( mcp_tool_handler: MCPToolHandler | None = None, configuration: Mapping[str, str] | None = None, restrict_env_to_configuration: bool = True, + migrate_legacy_request_scopes: bool = False, ) -> None: """Initialize the workflow factory. @@ -141,6 +142,12 @@ def __init__( discovered in PowerFx expressions inside the workflow definition so unrelated environment variables never enter the PowerFx scope. + migrate_legacy_request_scopes: Opt in to resuming legacy pending requests + without an originating workflow scope. Each request receives a stable, + isolated scope derived from its framework request ID, never the active + run's scope. The old MCP session cannot be recovered; resume requests + separately and checkpoint the migrated continuation. Defaults to + ``False``, which rejects unknown origins before consuming responses. Examples: .. code-block:: python @@ -196,6 +203,7 @@ def __init__( self._mcp_tool_handler = mcp_tool_handler self._configuration: dict[str, str] = dict(configuration) if configuration else {} self._restrict_env_to_configuration = restrict_env_to_configuration + self._migrate_legacy_request_scopes = migrate_legacy_request_scopes def create_workflow_from_yaml_path( self, @@ -455,6 +463,7 @@ def _create_workflow( http_request_handler=self._http_request_handler, mcp_tool_handler=self._mcp_tool_handler, env_config=env_config, + migrate_legacy_request_scopes=self._migrate_legacy_request_scopes, ) workflow = graph_builder.build() except ValueError as e: diff --git a/python/packages/declarative/agent_framework_declarative/_workflows/_mcp_handler.py b/python/packages/declarative/agent_framework_declarative/_workflows/_mcp_handler.py index 3dee698de35..c3a3d16511e 100644 --- a/python/packages/declarative/agent_framework_declarative/_workflows/_mcp_handler.py +++ b/python/packages/declarative/agent_framework_declarative/_workflows/_mcp_handler.py @@ -31,16 +31,26 @@ import hashlib import json import logging +import uuid from collections import OrderedDict from collections.abc import Awaitable, Callable +from contextlib import suppress from contextvars import ContextVar, Token from dataclasses import dataclass, field from typing import TYPE_CHECKING, Any, ClassVar, Protocol, cast, runtime_checkable import httpx2 as httpx +from agent_framework._workflows._const import ( + WORKFLOW_AGENT_RUN_KEY, + WORKFLOW_PARENT_SESSION_ID_KEY, + WORKFLOW_SESSION_ID_KEY, +) + +from ._errors import DeclarativeWorkflowError if TYPE_CHECKING: from agent_framework import Content + from agent_framework._workflows._state import State __all__ = [ "ClientProvider", @@ -53,6 +63,93 @@ logger = logging.getLogger(__name__) _DEFAULT_CACHE_MAX_SIZE = 32 +_WORKFLOW_SESSION_ID_KEY = "_declarative_mcp_workflow_session_id" +_AGENT_WORKFLOW_SESSION_NAMESPACE_KEY = "_declarative_mcp_agent_workflow_session_namespace" +_SESSIONLESS_AGENT_WORKFLOW_SESSION_ID_KEY = "_declarative_mcp_sessionless_agent_workflow_session_id" +_WORKFLOW_AGENT_SESSION_ID_KEY = "_workflow_agent_session_id" + + +def get_or_create_workflow_session_id(state: State) -> str: + # Core runs own the active scope even without the declarative factory initializer. + if state.has(WORKFLOW_AGENT_RUN_KEY): + workflow_session_id = state.get(WORKFLOW_SESSION_ID_KEY) + if workflow_session_id is not None: + if not isinstance(workflow_session_id, str) or not workflow_session_id: + raise ValueError("Invalid MCP workflow session state.") + state.set(_WORKFLOW_SESSION_ID_KEY, workflow_session_id) + return workflow_session_id + return _get_or_create_declarative_session_id(state) + + +def _get_or_create_declarative_session_id(state: State) -> str: + agent_session_id = state.get(_WORKFLOW_AGENT_SESSION_ID_KEY) + if agent_session_id is not None and state.get(WORKFLOW_PARENT_SESSION_ID_KEY) is None: + if not isinstance(agent_session_id, str) or not agent_session_id: + raise ValueError("Invalid agent session state.") + namespace_value = state.get(_AGENT_WORKFLOW_SESSION_NAMESPACE_KEY) + if namespace_value is None: + namespace_value = uuid.uuid4().hex + state.set(_AGENT_WORKFLOW_SESSION_NAMESPACE_KEY, namespace_value) + if not isinstance(namespace_value, str): + raise ValueError("Invalid MCP agent workflow session state.") + try: + namespace = uuid.UUID(namespace_value) + except ValueError as exc: + raise ValueError("Invalid MCP agent workflow session state.") from exc + workflow_session_id = uuid.uuid5(namespace, agent_session_id).hex + state.set(_WORKFLOW_SESSION_ID_KEY, workflow_session_id) + return workflow_session_id + + workflow_session_id = state.get(_WORKFLOW_SESSION_ID_KEY) + if workflow_session_id is None: + workflow_session_id = uuid.uuid4().hex + state.set(_WORKFLOW_SESSION_ID_KEY, workflow_session_id) + if not isinstance(workflow_session_id, str) or not workflow_session_id: + raise ValueError("Invalid MCP workflow session state.") + return workflow_session_id + + +def activate_workflow_session_id(state: State, *, reset_unscoped: bool, is_agent_run: bool = False) -> None: + parent_scope = state.get(WORKFLOW_PARENT_SESSION_ID_KEY) + if parent_scope is not None: + if not isinstance(parent_scope, str) or not parent_scope: + raise ValueError("Invalid parent workflow session scope.") + state.set(_WORKFLOW_SESSION_ID_KEY, parent_scope) + elif state.get(_WORKFLOW_AGENT_SESSION_ID_KEY) is not None: + _get_or_create_declarative_session_id(state) + elif is_agent_run: + workflow_session_id = state.get(_SESSIONLESS_AGENT_WORKFLOW_SESSION_ID_KEY) + if workflow_session_id is None: + workflow_session_id = uuid.uuid4().hex + state.set(_SESSIONLESS_AGENT_WORKFLOW_SESSION_ID_KEY, workflow_session_id) + if not isinstance(workflow_session_id, str) or not workflow_session_id: + raise ValueError("Invalid MCP sessionless agent workflow session state.") + state.set(_WORKFLOW_SESSION_ID_KEY, workflow_session_id) + elif reset_unscoped: + state.set(_WORKFLOW_SESSION_ID_KEY, uuid.uuid4().hex) + state.set(WORKFLOW_SESSION_ID_KEY, _get_or_create_declarative_session_id(state)) + + +def restore_workflow_session_id(state: State, workflow_session_id: str) -> None: + if not isinstance(workflow_session_id, str) or not workflow_session_id: + raise ValueError("Invalid MCP approval workflow session state.") + # A response resumes its originating scope, not the latest agent turn's scope. + state.set(_WORKFLOW_AGENT_SESSION_ID_KEY, None) + state.set(_WORKFLOW_SESSION_ID_KEY, workflow_session_id) + state.set(WORKFLOW_SESSION_ID_KEY, workflow_session_id) + + +def get_request_workflow_session_id(workflow_session_id: str | None) -> str: + """Validate a pending request's scope without inferring it from the active run.""" + if workflow_session_id is None: + raise DeclarativeWorkflowError( + "Request is missing a valid originating workflow session. " + "Restore the checkpoint with WorkflowFactory(migrate_legacy_request_scopes=True) " + "to resume legacy requests in isolated scopes, or restart the workflow." + ) + if not isinstance(workflow_session_id, str) or not workflow_session_id: + raise DeclarativeWorkflowError("Invalid originating workflow session state.") + return workflow_session_id @dataclass @@ -74,6 +171,10 @@ class MCPToolInvocation: - ``connection_name``: Optional Foundry connection name forwarded for handlers that resolve auth/credentials by connection. The default handler does not consume this field. + - ``workflow_session_id``: Framework-owned identifier for the current + workflow session. The default handler uses it to prevent separate + workflows from sharing one stateful MCP protocol session. Missing or + empty identifiers use a fresh session that is closed after each invocation. """ server_url: str @@ -82,6 +183,7 @@ class MCPToolInvocation: arguments: dict[str, Any] = field(default_factory=dict) # type: ignore[reportUnknownVariableType] headers: dict[str, str] = field(default_factory=dict) # type: ignore[reportUnknownVariableType] connection_name: str | None = None + workflow_session_id: str | None = None def _empty_outputs() -> list[Any]: @@ -154,23 +256,46 @@ class _CacheEntry: tool: Any # MCPStreamableHTTPTool — typed Any to avoid import at module load owned_httpx_client: httpx.AsyncClient | None + active_users: int = 0 + evicted: bool = False + disposal_claimed: bool = False + closed: asyncio.Event = field(default_factory=asyncio.Event) + close_exception: BaseException | None = None + + +class _EntryCreationCancelled(Exception): + """Signal waiters to retry after the task creating their entry was cancelled.""" + + +class _EntryCreationCleanupFailure(BaseException): + """Keep failed-handshake cleanup separate from the connection failure.""" + + def __init__(self, creation_exception: BaseException, cleanup_exception: BaseException) -> None: + self.creation_exception = creation_exception + self.cleanup_exception = cleanup_exception class DefaultMCPToolHandler: """Default :class:`MCPToolHandler` backed by :class:`agent_framework.MCPStreamableHTTPTool`. - Without a ``client_provider``, caches one + Without a ``client_provider`` and with a non-empty workflow session ID, caches one :class:`agent_framework.MCPStreamableHTTPTool` instance per - ``(server_url, server_label, connection_name, headers_hash)`` in a bounded - LRU. The cache prevents re-establishing an MCP session for every - invocation while ensuring different header sets (auth tokens) cannot - share a session — matches the .NET design intent while bounding - cardinality. ``server_label`` and ``connection_name`` also participate - in the key to distinguish logical connections. + ``(workflow_session_id, server_url, server_label, connection_name, + headers_hash)`` in a bounded LRU. The cache prevents re-establishing an MCP + session for every invocation while ensuring separate workflow sessions and + different header sets (auth tokens) cannot share a session — matches the + .NET design intent while bounding cardinality. ``server_label`` and + ``connection_name`` also participate in the key to distinguish logical + connections. Active invocations lease their cache entry. Admission counts + cached clients, connections in progress, and clients still closing. At + capacity, an idle entry is closed before its replacement connects; if no + idle entry is available, a new scope receives a capacity error. Header *names* are lower-cased inside the hash payload only — the headers passed on the wire keep the caller's original casing — so two YAML actions that spell ``Authorization`` differently still share a cache entry. + Invocations without a workflow session ID bypass the cache and close their + tool/session after each invocation, using the same lifecycle as provider-backed calls. Construction modes: @@ -198,10 +323,12 @@ class DefaultMCPToolHandler: Args: client_provider: Optional per-invocation ``httpx.AsyncClient`` provider. - cache_max_size: Maximum number of cached MCP clients in no-provider mode. - When exceeded, the least-recently-used entry is evicted and its - owned client closed. Defaults to ``32``. Does not enable session - caching when a provider is configured. + cache_max_size: Maximum number of live cache-managed MCP clients, + including connecting and closing clients. At capacity, the + least-recently-used idle entry is closed before connecting a new + client. If every slot is busy, new scopes receive a capacity error; + existing scopes can still reuse their cached client. Defaults to + ``32``. Does not limit provider-backed or unscoped invocations. """ LIST_TOOLS_TOOL_NAME: ClassVar[str] = "tools/list" @@ -229,15 +356,18 @@ def __init__( raise ValueError(f"cache_max_size must be positive, got {cache_max_size}") self._client_provider = client_provider self._cache_max_size = cache_max_size - self._cache: OrderedDict[tuple[str, str, str, str], _CacheEntry] = OrderedDict() + self._cache: OrderedDict[tuple[str, str, str, str, str], _CacheEntry] = OrderedDict() + self._retired: dict[int, _CacheEntry] = {} # Outer lock guards the cache + in-flight-future map only — never # held across network I/O. self._cache_lock = asyncio.Lock() + self._creation_semaphore = asyncio.Semaphore(cache_max_size) # Per-key in-flight futures: while one task is connecting, other # tasks awaiting the same key will await the same future and share # the resulting cache entry. - self._inflight: dict[tuple[str, str, str, str], asyncio.Future[_CacheEntry]] = {} - # Completion signals only: provider-backed calls never share entries. + self._inflight: dict[tuple[str, str, str, str, str], asyncio.Future[_CacheEntry]] = {} + self._inflight_cleanup: dict[tuple[str, str, str, str, str], asyncio.Future[None]] = {} + # Completion signals only: non-cacheable calls never share entries. self._active_invocations: set[asyncio.Future[None]] = set() # Keep ancestry so a completed nested call cannot hide an active parent # in the context inherited by child tasks, including cleanup tasks. @@ -247,6 +377,7 @@ def __init__( # Set by ``aclose`` to prevent post-close cache insertions and to # reject new ``invoke_tool`` calls. Once set, never cleared. self._closed = False + self._shutdown_task: asyncio.Task[None] | None = None async def invoke_tool(self, invocation: MCPToolInvocation) -> MCPToolResult: """Invoke ``invocation.tool_name`` on an MCP client for the server. @@ -272,9 +403,10 @@ async def invoke_tool(self, invocation: MCPToolInvocation) -> MCPToolResult: entry: _CacheEntry | None = None completion: asyncio.Future[None] | None = None context_token: Token[tuple[asyncio.Future[None], ...]] | None = None + use_cache = self._client_provider is None and bool(invocation.workflow_session_id) try: try: - if self._client_provider is None: + if use_cache: entry = await self._get_or_create_entry(invocation) else: async with self._cache_lock: @@ -304,7 +436,9 @@ async def invoke_tool(self, invocation: MCPToolInvocation) -> MCPToolResult: return await self._invoke_entry(entry, invocation) finally: - if completion is not None: + if use_cache and entry is not None: + await self._release_entry(entry) + elif completion is not None: try: if entry is not None: await self._close_invocation_entry(entry) @@ -433,42 +567,79 @@ async def _invoke_list_tools(entry: _CacheEntry) -> MCPToolResult: return MCPToolResult(outputs=[Content.from_text(json.dumps(payload, indent=2, allow_nan=False))]) async def aclose(self) -> None: - """Close cached clients and wait for provider-backed invocations to clean up. + """Close cached clients and wait for non-cacheable invocations to clean up. Caller-supplied :class:`httpx.AsyncClient` instances (returned by the ``client_provider`` callback) are NOT closed. - Provider-backed calls already executing may finish; pending connections + Provider-backed and unscoped calls already executing may finish; pending connections are rejected after connecting and cleaned up by their invocation. Concurrent shutdown calls wait for those invocations as well. - Idempotent. In no-provider mode, a second call returns immediately. + Idempotent. Concurrent and subsequent calls await the same shutdown task. Drains any in-flight ``_create_entry`` tasks before returning so their resources are cleaned up; the in-flight tasks see ``self._closed`` in phase 3 of :meth:`_get_or_create_entry`, close their own entry, and resolve their future with ``RuntimeError("DefaultMCPToolHandler is closed")``. Raises: - RuntimeError: If called from an active provider-backed invocation's + RuntimeError: If called from an active provider-backed or unscoped invocation's context, including inherited child tasks and cleanup. Rejected before changing handler state to avoid waiting on itself. """ if any(not completion.done() for completion in self._invocation_context.get()): raise RuntimeError( "DefaultMCPToolHandler.aclose() cannot be called from an active provider-backed invocation" + " or unscoped invocation" ) async with self._cache_lock: - if self._closed and self._client_provider is None: - return - self._closed = True - entries = list(self._cache.values()) - self._cache.clear() - inflight_futures = list(self._inflight.values()) - active_invocations = list(self._active_invocations) - - for completion in active_invocations: - # Cancelling shutdown must not cancel an invocation's completion signal. - await asyncio.shield(completion) + if self._shutdown_task is None: + self._closed = True + entries = list(self._cache.values()) + entry_ids = {id(entry) for entry in entries} + entries.extend(entry for entry in self._retired.values() if id(entry) not in entry_ids) + self._cache.clear() + entries_to_close: list[_CacheEntry] = [] + for entry in entries: + entry.evicted = True + self._retired[id(entry)] = entry + if entry.active_users == 0 and not entry.disposal_claimed: + entry.disposal_claimed = True + entries_to_close.append(entry) + inflight_futures = list(self._inflight.values()) + inflight_cleanup_futures = list(self._inflight_cleanup.values()) + active_invocations = list(self._active_invocations) + self._shutdown_task = asyncio.create_task( + self._drain_shutdown( + entries, + entries_to_close, + inflight_futures, + inflight_cleanup_futures, + active_invocations, + ) + ) + shutdown_task = self._shutdown_task + + await asyncio.shield(shutdown_task) + + async def _drain_shutdown( + self, + entries: list[_CacheEntry], + entries_to_close: list[_CacheEntry], + inflight_futures: list[asyncio.Future[_CacheEntry]], + inflight_cleanup_futures: list[asyncio.Future[None]], + active_invocations: list[asyncio.Future[None]], + ) -> None: + if active_invocations: + invocation_results = await asyncio.gather(*active_invocations, return_exceptions=True) + for result in invocation_results: + if isinstance(result, BaseException): + if isinstance(result, (KeyboardInterrupt, SystemExit, GeneratorExit)): + raise result + logger.debug( + "DefaultMCPToolHandler: active invocation raised during aclose", + exc_info=(type(result), result, result.__traceback__), + ) # Wait for in-flight creations to finish their self-cleanup. Each # in-flight task self-closes its entry under the closed-flag branch @@ -480,9 +651,23 @@ async def aclose(self) -> None: except BaseException: logger.debug("DefaultMCPToolHandler: in-flight future raised during aclose", exc_info=True) continue + cleanup_results = await asyncio.gather(*inflight_cleanup_futures, return_exceptions=True) + close_results = await asyncio.gather( + *(self._close_claimed_entry(entry) for entry in entries_to_close), + return_exceptions=True, + ) + for entry in entries: + await entry.closed.wait() for entry in entries: - await self._close_entry(entry) + if entry.close_exception is not None: + raise entry.close_exception + for result in cleanup_results: + if isinstance(result, BaseException): + raise result + for result in close_results: + if isinstance(result, BaseException): + raise result async def __aenter__(self) -> DefaultMCPToolHandler: return self @@ -497,6 +682,7 @@ async def __aexit__(self, exc_type: Any, exc: Any, tb: Any) -> None: async def _get_or_create_entry(self, invocation: MCPToolInvocation) -> _CacheEntry: """Look up (or create) the cached MCP client for this invocation.""" key = self._cache_key( + invocation.workflow_session_id, invocation.server_url, invocation.server_label, invocation.connection_name, @@ -506,81 +692,215 @@ async def _get_or_create_entry(self, invocation: MCPToolInvocation) -> _CacheEnt # Phase 1: check the cache and either claim creation or wait for an # already in-flight creation. creating = False + evicted: _CacheEntry | None = None async with self._cache_lock: if self._closed: raise RuntimeError("DefaultMCPToolHandler is closed") existing = self._cache.get(key) if existing is not None: self._cache.move_to_end(key) + existing.active_users += 1 return existing inflight = self._inflight.get(key) if inflight is None: + if len(self._cache) + len(self._retired) + len(self._inflight) >= self._cache_max_size: + idle = next( + ((cached_key, entry) for cached_key, entry in self._cache.items() if entry.active_users == 0), + None, + ) + if idle is None: + raise RuntimeError( + "MCP client capacity is exhausted; retry when an active invocation completes" + ) + evicted_key, evicted = idle + del self._cache[evicted_key] + evicted.evicted = True + evicted.disposal_claimed = True + self._retired[id(evicted)] = evicted inflight = asyncio.get_running_loop().create_future() self._inflight[key] = inflight + self._inflight_cleanup[key] = asyncio.get_running_loop().create_future() creating = True if not creating: - return await inflight + with suppress(_EntryCreationCancelled): + _ = await asyncio.shield(inflight) + return await self._get_or_create_entry(invocation) # Phase 2: we own creation. Build the entry outside the lock. try: - entry = await self._create_entry(invocation) + if evicted is not None: + await self._close_claimed_entry(evicted) + async with self._creation_semaphore: + async with self._cache_lock: + if self._closed: + raise RuntimeError("DefaultMCPToolHandler is closed") + entry = await self._create_entry(invocation) except BaseException as exc: - async with self._cache_lock: - self._inflight.pop(key, None) - if not inflight.done(): - inflight.set_exception(exc if isinstance(exc, BaseException) else RuntimeError(str(exc))) - # Mark the exception retrieved to suppress noisy "Future exception - # was never retrieved" warnings when there are no other awaiters - # (other awaiters still see the exception through their ``await``). - inflight.exception() + await self._complete_inflight_failure(key, inflight, exc) + if isinstance(exc, _EntryCreationCleanupFailure): + raise exc.cleanup_exception from exc.creation_exception raise - # Phase 3: insert with LRU eviction; resolve the in-flight future. + # Phase 3: insert and resolve the in-flight future. # If ``aclose`` ran while we were connecting, ``_closed`` is now # True; don't insert into the cache (it has been drained), close # the just-built entry, and surface the closed-handler error to # all awaiters of the future. - evicted: _CacheEntry | None = None duplicate: _CacheEntry | None = None handler_closed = False - async with self._cache_lock: - self._inflight.pop(key, None) - if self._closed: - handler_closed = True - else: - existing = self._cache.get(key) - if existing is not None: - # Another writer beat us; prefer the existing entry and - # discard ours after the lock is released. - self._cache.move_to_end(key) - duplicate = entry - entry = existing + try: + async with self._cache_lock: + self._inflight.pop(key, None) + if self._closed: + handler_closed = True else: - self._cache[key] = entry - self._cache.move_to_end(key) - if len(self._cache) > self._cache_max_size: - _evicted_key, evicted = self._cache.popitem(last=False) - if not inflight.done(): - inflight.set_result(entry) + existing = self._cache.get(key) + if existing is not None: + # Another writer beat us; prefer the existing entry and + # discard ours after the lock is released. + self._cache.move_to_end(key) + duplicate = entry + entry = existing + entry.active_users += 1 + else: + entry.active_users = 1 + self._cache[key] = entry + self._cache.move_to_end(key) + if not inflight.done(): + inflight.set_result(entry) + except BaseException as exc: + await self._abort_entry_creation(key, inflight, entry, exc) + raise if handler_closed: # Close our orphaned entry; resolve the future with a clear # error so the caller (and any other awaiters) surface a # consistent "handler is closed" failure rather than receiving # an entry we are about to close behind their back. - await self._close_entry(entry) err = RuntimeError("DefaultMCPToolHandler is closed") - if not inflight.done(): - inflight.set_exception(err) - inflight.exception() + cleanup = self._inflight_cleanup.pop(key) + try: + await self._close_invocation_entry(entry, cleanup) + finally: + if not inflight.done(): + inflight.set_exception(err) + inflight.exception() raise err - if duplicate is not None: - await self._close_entry(duplicate) - if evicted is not None: - await self._close_entry(evicted) + cleanup = self._inflight_cleanup.pop(key) + cleanup.set_result(None) + try: + if duplicate is not None: + await self._close_entry(duplicate) + except BaseException: + await self._release_entry(entry) + raise return entry + async def _abort_entry_creation( + self, + key: tuple[str, str, str, str, str], + inflight: asyncio.Future[_CacheEntry], + entry: _CacheEntry, + exc: BaseException, + ) -> None: + async def cleanup() -> None: + async with self._cache_lock: + cleanup_outcome = self._inflight_cleanup[key] + try: + await self._close_entry(entry) + except BaseException as cleanup_exc: + cleanup_outcome.set_exception(cleanup_exc) + cleanup_outcome.exception() + raise + else: + cleanup_outcome.set_result(None) + finally: + async with self._cache_lock: + self._inflight.pop(key, None) + self._inflight_cleanup.pop(key, None) + if not inflight.done(): + inflight.set_exception(self._entry_creation_exception(exc)) + inflight.exception() + + cleanup_task = asyncio.create_task(cleanup()) + while not cleanup_task.done(): + try: + await asyncio.shield(cleanup_task) + except asyncio.CancelledError: + continue + cleanup_task.result() + + async def _complete_inflight_failure( + self, + key: tuple[str, str, str, str, str], + inflight: asyncio.Future[_CacheEntry], + exc: BaseException, + ) -> None: + async def cleanup() -> None: + async with self._cache_lock: + self._inflight.pop(key, None) + cleanup_outcome = self._inflight_cleanup.pop(key) + creation_exception = exc + if isinstance(exc, _EntryCreationCleanupFailure): + creation_exception = exc.creation_exception + cleanup_outcome.set_exception(exc.cleanup_exception) + cleanup_outcome.exception() + else: + cleanup_outcome.set_result(None) + if not inflight.done(): + inflight.set_exception(self._entry_creation_exception(creation_exception)) + inflight.exception() + + await self._await_cleanup_task(asyncio.create_task(cleanup())) + + @staticmethod + def _entry_creation_exception(exc: BaseException) -> BaseException: + if isinstance(exc, asyncio.CancelledError): + return _EntryCreationCancelled() + return exc + + async def _release_entry(self, entry: _CacheEntry) -> None: + await self._await_cleanup_task(asyncio.create_task(self._release_entry_core(entry))) + + @staticmethod + async def _await_cleanup_task(cleanup: asyncio.Task[None]) -> None: + cancelled = False + while not cleanup.done(): + try: + await asyncio.shield(cleanup) + except asyncio.CancelledError: + cancelled = True + cleanup.result() + if cancelled: + raise asyncio.CancelledError + + async def _release_entry_core(self, entry: _CacheEntry) -> None: + close_entry = False + async with self._cache_lock: + entry.active_users -= 1 + if entry.active_users == 0 and entry.evicted and not entry.disposal_claimed: + entry.disposal_claimed = True + close_entry = True + + if close_entry: + await self._close_claimed_entry(entry) + + async def _close_claimed_entry(self, entry: _CacheEntry) -> None: + await self._await_cleanup_task(asyncio.create_task(self._close_claimed_entry_core(entry))) + + async def _close_claimed_entry_core(self, entry: _CacheEntry) -> None: + cleanup_outcome: asyncio.Future[None] = asyncio.get_running_loop().create_future() + try: + await self._close_invocation_entry(entry, cleanup_outcome) + except BaseException as exc: + entry.close_exception = cleanup_outcome.exception() if cleanup_outcome.done() else exc + raise + finally: + async with self._cache_lock: + self._retired.pop(id(entry), None) + entry.closed.set() + async def _create_entry(self, invocation: MCPToolInvocation) -> _CacheEntry: """Construct (and connect) a fresh MCP client for ``invocation``.""" from agent_framework import MCPStreamableHTTPTool @@ -606,7 +926,7 @@ def _header_provider(_kwargs: dict[str, Any]) -> dict[str, str]: ) try: await tool.connect() - except BaseException: + except BaseException as creation_exception: failed_entry = _CacheEntry( tool=tool, owned_httpx_client=( @@ -615,10 +935,17 @@ def _header_provider(_kwargs: dict[str, Any]) -> dict[str, str]: else None ), ) - if self._client_provider is not None: - await self._close_invocation_entry(failed_entry) - else: - await self._close_entry(failed_entry) + cleanup_outcome: asyncio.Future[None] = asyncio.get_running_loop().create_future() + try: + await self._close_invocation_entry(failed_entry, cleanup_outcome) + except BaseException as cleanup_exception: + if ( + self._client_provider is not None + or not invocation.workflow_session_id + or (cleanup_outcome.done() and cleanup_outcome.exception() is None) + ): + raise + raise _EntryCreationCleanupFailure(creation_exception, cleanup_exception) from cleanup_exception raise # ``MCPStreamableHTTPTool.get_mcp_client`` lazily creates an @@ -631,7 +958,11 @@ def _header_provider(_kwargs: dict[str, Any]) -> dict[str, str]: owned_client = cast("httpx.AsyncClient | None", getattr(tool, "_httpx_client", None)) return _CacheEntry(tool=tool, owned_httpx_client=owned_client) - async def _close_invocation_entry(self, entry: _CacheEntry) -> None: + async def _close_invocation_entry( + self, + entry: _CacheEntry, + cleanup_outcome: asyncio.Future[None] | None = None, + ) -> None: """Finish invocation cleanup even if the caller is cancelled again.""" # MCPStreamableHTTPTool dispatches connect/close to its lifecycle owner, # keeping the SDK's cancel-scope entry and exit on that same task. @@ -642,7 +973,16 @@ async def _close_invocation_entry(self, entry: _CacheEntry) -> None: await asyncio.shield(cleanup) except asyncio.CancelledError: cancelled = True - cleanup.result() # Propagate cancellation/errors from cleanup itself. + try: + cleanup.result() # Propagate cancellation/errors from cleanup itself. + except BaseException as exc: + if cleanup_outcome is not None and not cleanup_outcome.done(): + cleanup_outcome.set_exception(exc) + cleanup_outcome.exception() + raise + else: + if cleanup_outcome is not None and not cleanup_outcome.done(): + cleanup_outcome.set_result(None) if cancelled: raise asyncio.CancelledError @@ -662,25 +1002,35 @@ async def _close_entry(self, entry: _CacheEntry) -> None: @staticmethod def _cache_key( + workflow_session_id: str | None, server_url: str, server_label: str | None, connection_name: str | None, headers: dict[str, str] | None, - ) -> tuple[str, str, str, str]: + ) -> tuple[str, str, str, str, str]: """Build an order-independent cache key for the invocation identity. Used only without a ``client_provider``. The key includes - ``server_label`` and ``connection_name`` to distinguish logical - connections. + ``workflow_session_id`` to isolate stateful MCP protocol sessions + between workflows, plus ``server_label`` and ``connection_name`` to + distinguish logical connections. Header *names* are lower-cased inside the hash payload only so that ``Authorization`` and ``authorization`` map to the same cache entry. Header values remain case-sensitive (per RFC 7235). """ + if not workflow_session_id: + raise ValueError("MCP session caching requires a non-empty workflow session ID") if not headers: headers_hash = "0" else: normalized = sorted((k.lower(), v) for k, v in headers.items()) payload = json.dumps(normalized, ensure_ascii=False) headers_hash = hashlib.sha256(payload.encode("utf-8")).hexdigest() - return (server_url, server_label or "", connection_name or "", headers_hash) + return ( + workflow_session_id, + server_url, + server_label or "", + connection_name or "", + headers_hash, + ) diff --git a/python/packages/declarative/tests/conftest.py b/python/packages/declarative/tests/conftest.py index 083de75660f..d6b7f7859ae 100644 --- a/python/packages/declarative/tests/conftest.py +++ b/python/packages/declarative/tests/conftest.py @@ -2,9 +2,25 @@ """Pytest configuration for declarative tests.""" +from __future__ import annotations + import sys +from collections.abc import Callable +from typing import Any import pytest +from agent_framework import ( + CheckpointStorage, + Executor, + SubWorkflowRequestMessage, + SubWorkflowResponseMessage, + Workflow, + WorkflowBuilder, + WorkflowContext, + WorkflowExecutor, + handler, + response_handler, +) # Skip all tests in this directory on Python 3.14+ because powerfx doesn't support it yet if sys.version_info >= (3, 14): @@ -18,3 +34,54 @@ def pytest_collection_modifyitems(config: pytest.Config, items: list[pytest.Item for item in items: if "declarative" in str(item.fspath): item.add_marker(skip_marker) + + +@pytest.fixture +def coordinator_relay_workflow() -> Callable[[Workflow, CheckpointStorage], Workflow]: + """Build a parent that relays child requests without changing their IDs.""" + + class Requester(Executor): + @handler + async def start(self, message: str, ctx: WorkflowContext) -> None: + await ctx.request_info(message, str) + + @response_handler + async def respond(self, original: str, response: str, ctx: WorkflowContext[Any, dict[str, str]]) -> None: + await ctx.yield_output({"response": response}) + + class Coordinator(Executor): + @handler + async def start(self, message: int, ctx: WorkflowContext[Any]) -> None: + await ctx.send_message({}, target_id="mcp") + await ctx.send_message(str(message), target_id="requester") + + @handler + async def relay(self, request: SubWorkflowRequestMessage, ctx: WorkflowContext) -> None: + await ctx.request_info(request, str, request_id=request.source_event.request_id) + + @response_handler + async def respond( + self, + original: SubWorkflowRequestMessage, + response: str, + ctx: WorkflowContext[SubWorkflowResponseMessage], + ) -> None: + await ctx.send_message(original.create_response(response), target_id="requester") + + @handler + async def collect(self, message: dict[str, str], ctx: WorkflowContext[dict[str, Any]]) -> None: + await ctx.send_message({}, target_id="mcp") + + def build(mcp_workflow: Workflow, storage: CheckpointStorage) -> Workflow: + coordinator = Coordinator(id="coordinator") + requester = WorkflowExecutor(WorkflowBuilder(start_executor=Requester(id="pause")).build(), id="requester") + mcp = WorkflowExecutor(mcp_workflow, id="mcp") + return ( + WorkflowBuilder(start_executor=coordinator, checkpoint_storage=storage) + .add_edge(coordinator, requester) + .add_edge(requester, coordinator) + .add_edge(coordinator, mcp) + .build() + ) + + return build diff --git a/python/packages/declarative/tests/test_declarative_approval_binding.py b/python/packages/declarative/tests/test_declarative_approval_binding.py index 441f73891d0..a12ce501d1c 100644 --- a/python/packages/declarative/tests/test_declarative_approval_binding.py +++ b/python/packages/declarative/tests/test_declarative_approval_binding.py @@ -177,7 +177,9 @@ def my_tool(x: int) -> int: executor = InvokeFunctionToolExecutor(self._action(), tools={"my_tool": my_tool}) - request = ToolApprovalRequest(request_id="r-1", function_name="my_tool", arguments={"x": 1}) + request = ToolApprovalRequest( + request_id="r-1", function_name="my_tool", arguments={"x": 1}, workflow_session_id="workflow-a" + ) await executor.handle_approval_response(request, ToolApprovalResponse(approved=True), mock_context) assert call_log == [1] @@ -193,7 +195,9 @@ def my_tool(x: int) -> int: return x executor = InvokeFunctionToolExecutor(self._action(), tools={"my_tool": my_tool}) - request = ToolApprovalRequest(request_id="r-bool", function_name="my_tool", arguments={"x": 1}) + request = ToolApprovalRequest( + request_id="r-bool", function_name="my_tool", arguments={"x": 1}, workflow_session_id="workflow-a" + ) response = object.__new__(ToolApprovalResponse) response.approved = "true" # type: ignore[assignment] # ty: ignore[invalid-assignment] response.reason = None @@ -214,8 +218,12 @@ def my_tool(x: int) -> int: executor = InvokeFunctionToolExecutor(self._action(), tools={"my_tool": my_tool}) - request_a = ToolApprovalRequest(request_id="r-A", function_name="my_tool", arguments={"x": 1}) - request_b = ToolApprovalRequest(request_id="r-B", function_name="my_tool", arguments={"x": 999}) + request_a = ToolApprovalRequest( + request_id="r-A", function_name="my_tool", arguments={"x": 1}, workflow_session_id="workflow-a" + ) + request_b = ToolApprovalRequest( + request_id="r-B", function_name="my_tool", arguments={"x": 999}, workflow_session_id="workflow-b" + ) # Deliver response for B first, then for A. Each invocation must use its own payload. await executor.handle_approval_response(request_b, ToolApprovalResponse(approved=True), mock_context) @@ -238,7 +246,9 @@ def my_tool(x: int) -> int: # Poison the old key shape (no longer read by the executor). mock_state._data["_tool_approval_state_fn_action"] = {"function_name": "other", "arguments": {"x": 999}} - request = ToolApprovalRequest(request_id="r-3", function_name="my_tool", arguments={"x": 7}) + request = ToolApprovalRequest( + request_id="r-3", function_name="my_tool", arguments={"x": 7}, workflow_session_id="workflow-a" + ) await executor.handle_approval_response(request, ToolApprovalResponse(approved=True), mock_context) assert call_log == [7] @@ -258,7 +268,9 @@ def my_tool(x: int) -> int: # Pretend the executor that emitted the request is gone; a fresh one handles the response. fresh = InvokeFunctionToolExecutor(self._action(), tools={"my_tool": my_tool}) - request = ToolApprovalRequest(request_id="r-4", function_name="my_tool", arguments={"x": 42}) + request = ToolApprovalRequest( + request_id="r-4", function_name="my_tool", arguments={"x": 42}, workflow_session_id="workflow-a" + ) await fresh.handle_approval_response(request, ToolApprovalResponse(approved=True), mock_context) assert call_log == [42] @@ -275,7 +287,9 @@ def my_tool(x: int) -> int: executor = InvokeFunctionToolExecutor(self._action(), tools={"my_tool": my_tool}) - request = ToolApprovalRequest(request_id="r-5", function_name="my_tool", arguments={"x": 3}) + request = ToolApprovalRequest( + request_id="r-5", function_name="my_tool", arguments={"x": 3}, workflow_session_id="workflow-a" + ) await executor.handle_approval_response( request, ToolApprovalResponse(approved=False, reason="not authorized"), mock_context ) @@ -333,6 +347,7 @@ async def test_resume_uses_request_payload_fields(self, mock_state, mock_context server_label="prod", arguments={"q": "x"}, connection_name="conn-A", + workflow_session_id="workflow-a", ) await executor.handle_approval_response(request, ToolApprovalResponse(approved=True), mock_context) @@ -358,6 +373,7 @@ async def test_truthy_malformed_response_does_not_invoke(self, mock_state, mock_ server_label=None, arguments={"q": "x"}, connection_name=None, + workflow_session_id="workflow-a", ) response = object.__new__(ToolApprovalResponse) response.approved = 1 # type: ignore[assignment] # ty: ignore[invalid-assignment] @@ -380,6 +396,7 @@ async def test_concurrent_pending_mcp_approvals_do_not_swap(self, mock_state, mo server_label=None, arguments={"q": "alpha"}, connection_name="conn-A", + workflow_session_id="workflow-a", ) request_b = MCPToolApprovalRequest( request_id="r-B", @@ -388,6 +405,7 @@ async def test_concurrent_pending_mcp_approvals_do_not_swap(self, mock_state, mo server_label=None, arguments={"q": "beta"}, connection_name="conn-B", + workflow_session_id="workflow-b", ) await executor.handle_approval_response(request_b, ToolApprovalResponse(approved=True), mock_context) @@ -416,6 +434,7 @@ async def test_legacy_headers_require_fresh_approval(self, mock_state, mock_cont server_label=None, arguments={"q": "x"}, connection_name=None, + workflow_session_id="legacy-request:r-1", ) await executor.handle_approval_response(request, ToolApprovalResponse(approved=True), mock_context) @@ -534,6 +553,7 @@ async def test_mcp_resume_ignores_stale_state_at_old_approval_key(self, mock_sta server_label=None, arguments={"q": "real"}, connection_name=None, + workflow_session_id="workflow-a", ) await executor.handle_approval_response(request, ToolApprovalResponse(approved=True), mock_context) @@ -557,6 +577,7 @@ async def test_fresh_mcp_executor_resume_works(self, mock_state, mock_context) - server_label=None, arguments={"q": "fresh"}, connection_name=None, + workflow_session_id="workflow-a", ) await fresh.handle_approval_response(request, ToolApprovalResponse(approved=True), mock_context) @@ -625,6 +646,7 @@ async def test_resume_routes_output_to_pinned_conversation_not_mutated_state( arguments={"q": "x"}, connection_name=None, metadata={"conversation_id": "conv-original"}, + workflow_session_id="workflow-a", ) await executor.handle_approval_response(original_request, ToolApprovalResponse(approved=True), mock_context) @@ -632,8 +654,10 @@ async def test_resume_routes_output_to_pinned_conversation_not_mutated_state( assert state.get("System.conversations.conv-mutated.messages") == [] @pytest.mark.asyncio - async def test_resume_handles_legacy_request_without_new_fields(self, mock_state, mock_context) -> None: - """Resume tolerates payloads lacking ``connection_name`` / ``metadata`` (legacy pickle shape).""" + async def test_resume_handles_migrated_legacy_request_without_routing_fields( + self, mock_state, mock_context + ) -> None: + """Migrated legacy payloads may still lack ``connection_name`` / ``metadata``.""" @dataclass class _LegacyMCPApprovalRequest: @@ -643,6 +667,7 @@ class _LegacyMCPApprovalRequest: server_label: str | None arguments: dict[str, Any] header_names: list[str] + workflow_session_id: str = "legacy-request:r-1" _seed_state(mock_state) handler = _RecordingMcpHandler() diff --git a/python/packages/declarative/tests/test_default_mcp_tool_handler.py b/python/packages/declarative/tests/test_default_mcp_tool_handler.py index 387744d84c7..5e0fb9543bb 100644 --- a/python/packages/declarative/tests/test_default_mcp_tool_handler.py +++ b/python/packages/declarative/tests/test_default_mcp_tool_handler.py @@ -17,17 +17,21 @@ import asyncio import json import sys -from typing import Any +from typing import Any, Literal from unittest.mock import AsyncMock, patch import httpx2 as httpx import pytest from agent_framework import Content +from agent_framework._workflows._state import State from agent_framework.exceptions import ToolExecutionException from agent_framework_declarative._workflows._mcp_handler import ( DefaultMCPToolHandler, MCPToolInvocation, + activate_workflow_session_id, + get_or_create_workflow_session_id, + restore_workflow_session_id, ) pytestmark = pytest.mark.skipif( @@ -140,22 +144,123 @@ def _patch_tool() -> Any: def _invocation( - *, server_url: str = "https://mcp.example/api", tool_name: str = "search", **overrides: Any + *, + server_url: str = "https://mcp.example/api", + tool_name: str = "search", + workflow_session_id: str | None = "test-workflow", + **overrides: Any, ) -> MCPToolInvocation: return MCPToolInvocation( server_url=server_url, tool_name=tool_name, + workflow_session_id=workflow_session_id, **overrides, ) +class TestUnscopedInvocations: + @pytest.mark.parametrize("workflow_session_id", [None, ""]) + @pytest.mark.parametrize("tool_name", ["search", "tools/list"]) + async def test_sequential_calls_do_not_reuse_sessions( + self, workflow_session_id: str | None, tool_name: str + ) -> None: + with _patch_tool(): + async with DefaultMCPToolHandler() as handler: + for index in range(1, 3): + result = await handler.invoke_tool( + _invocation( + workflow_session_id=workflow_session_id, tool_name=tool_name, headers={"X-Test": "1"} + ) + ) + assert not result.is_error + assert len(FakeTool.instances) == index + tool = FakeTool.instances[-1] + assert tool.connect_count == 1 + assert tool.close_count == 1 + assert tool._httpx_client is not None + assert tool._httpx_client.is_closed + assert not handler._cache + assert not handler._inflight + assert all(tool.close_count == 1 for tool in FakeTool.instances) + + @pytest.mark.parametrize("workflow_session_id", [None, ""]) + async def test_concurrent_calls_do_not_coalesce_connections(self, workflow_session_id: str | None) -> None: + all_connecting = asyncio.Event() + release = asyncio.Event() + original_connect = FakeTool.connect + + async def gated_connect(tool: FakeTool) -> None: + await original_connect(tool) + if len(FakeTool.instances) == 2: + all_connecting.set() + await release.wait() + + with _patch_tool(), patch.object(FakeTool, "connect", gated_connect): + async with DefaultMCPToolHandler() as handler: + tasks = [ + asyncio.create_task( + handler.invoke_tool(_invocation(workflow_session_id=workflow_session_id, tool_name=name)) + ) + for name in ("search", "tools/list") + ] + try: + await asyncio.wait_for(all_connecting.wait(), timeout=5) + assert len({id(tool.session) for tool in FakeTool.instances}) == 2 + assert all(tool.close_count == 0 for tool in FakeTool.instances) + assert not handler._cache + assert not handler._inflight + finally: + release.set() + results = await asyncio.gather(*tasks) + assert all(not result.is_error for result in results) + assert all(tool.close_count == 1 for tool in FakeTool.instances) + + @pytest.mark.parametrize("workflow_session_id", [None, ""]) + @pytest.mark.parametrize("cancel_active", [False, True]) + async def test_shutdown_waits_for_invocation_cleanup( + self, workflow_session_id: str | None, cancel_active: bool + ) -> None: + started = asyncio.Event() + release = asyncio.Event() + + async def gated_call(_tool: FakeTool, tool_name: str, **arguments: Any) -> Any: + started.set() + await release.wait() + return [Content.from_text("ok")] + + with _patch_tool(), patch.object(FakeTool, "call_tool", gated_call): + handler = DefaultMCPToolHandler() + invocation = asyncio.create_task( + handler.invoke_tool(_invocation(workflow_session_id=workflow_session_id, headers={"X-Test": "1"})) + ) + await started.wait() + shutdown = asyncio.create_task(handler.aclose()) + try: + await asyncio.sleep(0) + assert not shutdown.done() + assert FakeTool.instances[0].close_count == 0 + if cancel_active: + invocation.cancel() + finally: + release.set() + results = await asyncio.gather(invocation, shutdown, return_exceptions=True) + if cancel_active: + assert isinstance(results[0], asyncio.CancelledError) + else: + assert not invocation.result().is_error + shutdown.result() + assert FakeTool.instances[0].close_count == 1 + assert FakeTool.instances[0]._httpx_client is not None + assert FakeTool.instances[0]._httpx_client.is_closed + + # ---------- Provider-backed invocation lifetimes -------------------------- class TestProviderLifetimes: - @pytest.mark.parametrize("returns_none", [False, True]) + @pytest.mark.parametrize("client_source", ["provider", "fallback", "unscoped"]) async def test_actual_sdk_lifecycle_and_hook_cleanup( - self, returns_none: bool, caplog: pytest.LogCaptureFixture + self, client_source: str, caplog: pytest.LogCaptureFixture ) -> None: import agent_framework from agent_framework import MCPStreamableHTTPTool @@ -222,17 +327,23 @@ def create_tool(**kwargs: Any) -> MCPStreamableHTTPTool: tools.append(tool) return tool - caller_client = None if returns_none else RecordingClient() + caller_client = RecordingClient() if client_source == "provider" else None provider = AsyncMock(return_value=caller_client) try: with ( patch.object(httpx, "AsyncClient", RecordingClient), patch.object(agent_framework, "MCPStreamableHTTPTool", side_effect=create_tool), ): - async with DefaultMCPToolHandler(client_provider=provider) as handler: + async with DefaultMCPToolHandler( + client_provider=None if client_source == "unscoped" else provider + ) as handler: for index, tool_name in enumerate(("ok", "tools/list", "ok", "tools/list"), start=1): result = await handler.invoke_tool( - _invocation(tool_name=tool_name, headers={"X-Test": "lifecycle"}) + _invocation( + tool_name=tool_name, + headers={"X-Test": "lifecycle"}, + workflow_session_id=None if client_source == "unscoped" else "test-workflow", + ) ) assert not result.is_error if tool_name == "ok": @@ -240,7 +351,7 @@ def create_tool(**kwargs: Any) -> MCPStreamableHTTPTool: else: assert result.outputs[0].text is not None assert json.loads(result.outputs[0].text)["tools"][0]["name"] == "ok" - assert provider.await_count == index + assert provider.await_count == (0 if client_source == "unscoped" else index) assert len(tools) == index assert len(initialized) == index assert not terminated @@ -252,7 +363,7 @@ def create_tool(**kwargs: Any) -> MCPStreamableHTTPTool: assert tool._lifecycle_owner_task is None assert tool._header_hook_client is None assert all(client.event_hooks["request"] == [caller_hook] for client in clients) - if returns_none: + if caller_client is None: assert len(clients) == index assert all(client.is_closed and client.close_count == 1 for client in clients) else: @@ -620,7 +731,8 @@ async def close(tool: FakeTool) -> None: assert tool._httpx_client.is_closed @pytest.mark.parametrize("connect_fails", [False, True]) - async def test_cleanup_cancellation_signals_completion(self, connect_fails: bool) -> None: + @pytest.mark.parametrize("unscoped", [False, True]) + async def test_cleanup_cancellation_signals_completion(self, connect_fails: bool, unscoped: bool) -> None: original_connect = FakeTool.connect completions: list[asyncio.Future[None]] = [] @@ -639,10 +751,12 @@ async def close(tool: FakeTool) -> None: patch.object(FakeTool, "connect", connect), patch.object(FakeTool, "close", close), ): - handler = DefaultMCPToolHandler(client_provider=AsyncMock(return_value=None)) + handler = DefaultMCPToolHandler(client_provider=None if unscoped else AsyncMock(return_value=None)) # Python 3.10 may drop the message when a cancelled task's result is retrieved again. with pytest.raises(asyncio.CancelledError): - await handler.invoke_tool(_invocation(headers={"X-Test": "1"})) + await handler.invoke_tool( + _invocation(headers={"X-Test": "1"}, workflow_session_id=None if unscoped else "test-workflow") + ) assert handler._invocation_context.get() == () assert not handler._active_invocations assert len(completions) == 1 @@ -830,6 +944,84 @@ def test_invalid_cache_size_raises(self) -> None: DefaultMCPToolHandler(cache_max_size=-3) +class TestWorkflowSessionId: + def test_sessionless_activation_does_not_adopt_direct_scope(self) -> None: + state = State() + direct_scope = get_or_create_workflow_session_id(state) + + activate_workflow_session_id(state, reset_unscoped=False, is_agent_run=True) + sessionless_scope = get_or_create_workflow_session_id(state) + assert sessionless_scope != direct_scope + + activate_workflow_session_id(state, reset_unscoped=True) + assert get_or_create_workflow_session_id(state) not in (direct_scope, sessionless_scope) + state.commit() + + activate_workflow_session_id(state, reset_unscoped=False, is_agent_run=True) + assert get_or_create_workflow_session_id(state) == sessionless_scope + + def test_approval_restoration_preserves_sessionless_scope(self) -> None: + state = State() + activate_workflow_session_id(state, reset_unscoped=True, is_agent_run=True) + sessionless_scope = get_or_create_workflow_session_id(state) + state.set("_workflow_agent_session_id", "session-a") + explicit_scope = get_or_create_workflow_session_id(state) + + restore_workflow_session_id(state, explicit_scope) + state.commit() + assert get_or_create_workflow_session_id(state) == explicit_scope + assert explicit_scope != sessionless_scope + + activate_workflow_session_id(state, reset_unscoped=False, is_agent_run=True) + assert get_or_create_workflow_session_id(state) == sessionless_scope + + @pytest.mark.parametrize("scope", ["", 42, False]) + def test_invalid_sessionless_scope_preserves_current_scope(self, scope: Any) -> None: + state = State() + current_scope = get_or_create_workflow_session_id(state) + state.set("_declarative_mcp_sessionless_agent_workflow_session_id", scope) + + with pytest.raises(ValueError, match="Invalid MCP sessionless agent workflow session state"): + activate_workflow_session_id(state, reset_unscoped=True, is_agent_run=True) + + assert get_or_create_workflow_session_id(state) == current_scope + + def test_restored_approval_scope_overrides_current_agent_session(self) -> None: + state = State() + state.set("_workflow_agent_session_id", "session-b") + agent_scope = get_or_create_workflow_session_id(state) + + restore_workflow_session_id(state, "approval-scope-a") + state.commit() + + assert get_or_create_workflow_session_id(state) == "approval-scope-a" + assert get_or_create_workflow_session_id(state) != agent_scope + + @pytest.mark.parametrize("workflow_session_id", ["", None, 42]) + def test_invalid_restored_approval_scope_preserves_current_scope(self, workflow_session_id: Any) -> None: + state = State() + current_scope = get_or_create_workflow_session_id(state) + + with pytest.raises(ValueError, match="Invalid MCP approval workflow session state"): + restore_workflow_session_id(state, workflow_session_id) + + assert get_or_create_workflow_session_id(state) == current_scope + + def test_separate_workflow_states_get_separate_ids(self) -> None: + first = get_or_create_workflow_session_id(State()) + second = get_or_create_workflow_session_id(State()) + + assert first != second + + def test_same_workflow_state_reuses_id(self) -> None: + state = State() + + first = get_or_create_workflow_session_id(state) + second = get_or_create_workflow_session_id(state) + + assert first == second + + # ---------- Tool kwargs ---------------------------------------------------- @@ -878,16 +1070,122 @@ async def test_header_provider_returns_captured_headers(self) -> None: class TestCache: + @pytest.mark.parametrize("cancel_active", [False, True]) + async def test_active_clients_reject_new_scopes_at_capacity(self, cancel_active: bool) -> None: + handler = DefaultMCPToolHandler(cache_max_size=1) + started = asyncio.Event() + release = asyncio.Event() + calls = 0 + + async def gated_call(_tool: FakeTool, tool_name: str, **arguments: Any) -> Any: + nonlocal calls + calls += 1 + if calls == 1: + started.set() + await release.wait() + return [Content.from_text("ok")] + + with _patch_tool(), patch.object(FakeTool, "call_tool", gated_call): + first = asyncio.create_task(handler.invoke_tool(_invocation(workflow_session_id="workflow-a"))) + await started.wait() + try: + reused = await handler.invoke_tool(_invocation(workflow_session_id="workflow-a")) + assert not reused.is_error + rejected = await handler.invoke_tool(_invocation(workflow_session_id="workflow-b")) + assert rejected.is_error + assert rejected.error_message is not None + assert "capacity" in rejected.error_message + assert len(FakeTool.instances) == 1 + assert FakeTool.instances[0].close_count == 0 + finally: + if cancel_active: + first.cancel() + release.set() + await asyncio.gather(first, return_exceptions=True) + + recovered = await handler.invoke_tool(_invocation(workflow_session_id="workflow-b")) + assert not recovered.is_error + await handler.aclose() + assert len(FakeTool.instances) == 2 + assert all(tool.close_count == 1 for tool in FakeTool.instances) + + async def test_replacement_waits_for_idle_client_cleanup_before_connecting(self) -> None: + handler = DefaultMCPToolHandler(cache_max_size=1) + close_started = asyncio.Event() + release_close = asyncio.Event() + original_close = FakeTool.close + + async def gated_close(tool: FakeTool) -> None: + if tool is FakeTool.instances[0]: + close_started.set() + await release_close.wait() + await original_close(tool) + + with _patch_tool(), patch.object(FakeTool, "close", gated_close): + await handler.invoke_tool(_invocation(workflow_session_id="workflow-a")) + replacement = asyncio.create_task(handler.invoke_tool(_invocation(workflow_session_id="workflow-b"))) + await close_started.wait() + try: + assert len(FakeTool.instances) == 1 + rejected = await handler.invoke_tool(_invocation(workflow_session_id="workflow-c")) + assert rejected.is_error + assert rejected.error_message is not None + assert "capacity" in rejected.error_message + finally: + release_close.set() + try: + replacement_result = await replacement + finally: + await handler.aclose() + assert not replacement_result.is_error + assert len(FakeTool.instances) == 2 + assert all(tool.close_count == 1 for tool in FakeTool.instances) + + @pytest.mark.parametrize("explicit_session", [False, True]) + async def test_sessionless_scope_reuses_its_own_cached_client(self, explicit_session: bool) -> None: + state = State() + with _patch_tool(): + async with DefaultMCPToolHandler() as handler: + activate_workflow_session_id(state, reset_unscoped=True, is_agent_run=True) + await handler.invoke_tool(_invocation(workflow_session_id=get_or_create_workflow_session_id(state))) + FakeTool.instances[0].call_handler = lambda **_: [Content.from_text("sessionless client")] + + if explicit_session: + state.set("_workflow_agent_session_id", "session-a") + activate_workflow_session_id(state, reset_unscoped=True, is_agent_run=explicit_session) + await handler.invoke_tool(_invocation(workflow_session_id=get_or_create_workflow_session_id(state))) + + state.set("_workflow_agent_session_id", None) + activate_workflow_session_id(state, reset_unscoped=False, is_agent_run=True) + result = await handler.invoke_tool( + _invocation(workflow_session_id=get_or_create_workflow_session_id(state)) + ) + + assert not result.is_error + assert len(result.outputs) == 1 + assert result.outputs[0].text == "sessionless client" + assert len(FakeTool.instances) == 2 + assert all(tool.connect_count == 1 for tool in FakeTool.instances) + assert all(tool.close_count == 1 for tool in FakeTool.instances) + @pytest.mark.asyncio async def test_same_url_and_headers_hit_cache(self) -> None: handler = DefaultMCPToolHandler() with _patch_tool(): - await handler.invoke_tool(_invocation(headers={"X": "1"})) - await handler.invoke_tool(_invocation(headers={"X": "1"})) + await handler.invoke_tool(_invocation(headers={"X": "1"}, workflow_session_id="workflow-a")) + await handler.invoke_tool(_invocation(headers={"X": "1"}, workflow_session_id="workflow-a")) # One tool created, connect called once. assert len(FakeTool.instances) == 1 assert FakeTool.instances[0].connect_count == 1 + @pytest.mark.asyncio + async def test_separate_workflow_sessions_use_separate_entries(self) -> None: + handler = DefaultMCPToolHandler() + with _patch_tool(): + await handler.invoke_tool(_invocation(workflow_session_id="workflow-a")) + await handler.invoke_tool(_invocation(workflow_session_id="workflow-b")) + assert len(FakeTool.instances) == 2 + @pytest.mark.asyncio async def test_different_headers_create_separate_entries(self) -> None: handler = DefaultMCPToolHandler() @@ -920,6 +1218,378 @@ async def test_lru_eviction_closes_old_entry(self) -> None: assert FakeTool.instances[1].close_count == 0 assert FakeTool.instances[2].close_count == 0 + @pytest.mark.asyncio + async def test_lru_eviction_skips_active_client_until_invocation_finishes(self) -> None: + handler = DefaultMCPToolHandler(cache_max_size=2) + first_started = asyncio.Event() + release_first = asyncio.Event() + + async def gated_call(tool: FakeTool, tool_name: str, **arguments: Any) -> Any: + if tool.kwargs["url"] == "https://a/": + first_started.set() + await release_first.wait() + return [Content.from_text("ok")] + + with _patch_tool(), patch.object(FakeTool, "call_tool", gated_call): + first = asyncio.create_task(handler.invoke_tool(_invocation(server_url="https://a/"))) + await first_started.wait() + second = await handler.invoke_tool(_invocation(server_url="https://b/")) + third = await handler.invoke_tool(_invocation(server_url="https://c/")) + + assert not second.is_error + assert not third.is_error + assert FakeTool.instances[0].close_count == 0 + assert FakeTool.instances[1].close_count == 1 + close_task = asyncio.create_task(handler.aclose()) + await asyncio.sleep(0) + assert not close_task.done() + + release_first.set() + first_result = await first + await asyncio.gather(close_task) + + assert not first_result.is_error + assert FakeTool.instances[0].close_count == 1 + assert FakeTool.instances[1].close_count == 1 + assert FakeTool.instances[2].close_count == 1 + + @pytest.mark.asyncio + async def test_lru_eviction_cleanup_cancellation_releases_new_entry(self) -> None: + handler = DefaultMCPToolHandler(cache_max_size=1) + cancel_first_close = True + original_close = FakeTool.close + + async def close(tool: FakeTool) -> None: + if cancel_first_close and tool.kwargs["url"] == "https://a/": + tool.close_count += 1 + raise asyncio.CancelledError + await original_close(tool) + + with _patch_tool(), patch.object(FakeTool, "close", close): + await handler.invoke_tool(_invocation(server_url="https://a/")) + with pytest.raises(asyncio.CancelledError): + await handler.invoke_tool(_invocation(server_url="https://b/")) + + assert all(entry.active_users == 0 for entry in handler._cache.values()) + cancel_first_close = False + await asyncio.wait_for(handler.aclose(), timeout=1) + + @pytest.mark.parametrize("cancellation_count", [1, 2]) + async def test_cancellation_during_failed_creation_bookkeeping_propagates(self, cancellation_count: int) -> None: + handler = DefaultMCPToolHandler() + connect_started = asyncio.Event() + release_connect = asyncio.Event() + bookkeeping_waiting = asyncio.Event() + + class ObservedLock(asyncio.Lock): + async def acquire(self) -> Literal[True]: + if self.locked(): + bookkeeping_waiting.set() + return await super().acquire() + + handler._cache_lock = ObservedLock() + + async def failed_connect(_tool: FakeTool) -> None: + connect_started.set() + await release_connect.wait() + raise RuntimeError("connection failed") + + with _patch_tool(), patch.object(FakeTool, "connect", failed_connect): + invocation = asyncio.create_task(handler.invoke_tool(_invocation())) + try: + await asyncio.wait_for(connect_started.wait(), timeout=1) + await handler._cache_lock.acquire() + try: + release_connect.set() + await asyncio.wait_for(bookkeeping_waiting.wait(), timeout=1) + for _ in range(cancellation_count): + invocation.cancel() + await asyncio.sleep(0) + assert not invocation.done() + finally: + handler._cache_lock.release() + + with pytest.raises(asyncio.CancelledError): + await asyncio.wait_for(invocation, timeout=1) + assert not handler._inflight + assert not handler._inflight_cleanup + finally: + release_connect.set() + await asyncio.wait_for(handler.aclose(), timeout=1) + + assert FakeTool.instances[0].close_count == 1 + + async def test_repeated_cancellation_during_eviction_finalization_completes_shutdown(self) -> None: + handler = DefaultMCPToolHandler(cache_max_size=1) + close_started = asyncio.Event() + release_close = asyncio.Event() + finalization_waiting = asyncio.Event() + original_close = FakeTool.close + + class ObservedLock(asyncio.Lock): + async def acquire(self) -> Literal[True]: + if self.locked(): + finalization_waiting.set() + return await super().acquire() + + handler._cache_lock = ObservedLock() + + async def gated_close(tool: FakeTool) -> None: + if tool.kwargs["url"] == "https://a/": + close_started.set() + await release_close.wait() + await original_close(tool) + + with _patch_tool(), patch.object(FakeTool, "close", gated_close): + await handler.invoke_tool(_invocation(server_url="https://a/")) + retired_entry = next(iter(handler._cache.values())) + invocation = asyncio.create_task(handler.invoke_tool(_invocation(server_url="https://b/"))) + await close_started.wait() + invocation.cancel() + await handler._cache_lock.acquire() + try: + release_close.set() + await asyncio.wait_for(finalization_waiting.wait(), timeout=1) + invocation.cancel() + await asyncio.sleep(0) + finally: + handler._cache_lock.release() + + await asyncio.wait({invocation}, timeout=1) + assert invocation.done() + with pytest.raises(asyncio.CancelledError): + invocation.result() + assert retired_entry.closed.is_set() + assert id(retired_entry) not in handler._retired + await asyncio.wait_for(handler.aclose(), timeout=1) + + assert all(tool.close_count == 1 for tool in FakeTool.instances) + + @pytest.mark.asyncio + async def test_repeated_cancellation_while_releasing_entry_does_not_leak_lease(self) -> None: + handler = DefaultMCPToolHandler() + invocation_started = asyncio.Event() + release_invocation = asyncio.Event() + + async def gated_call(_tool: FakeTool, _tool_name: str, **_arguments: Any) -> Any: + invocation_started.set() + await release_invocation.wait() + return [Content.from_text("ok")] + + with _patch_tool(), patch.object(FakeTool, "call_tool", gated_call): + invocation = asyncio.create_task(handler.invoke_tool(_invocation(headers={"X": "1"}))) + await invocation_started.wait() + await handler._cache_lock.acquire() + invocation.cancel() + await asyncio.sleep(0) + invocation.cancel() + handler._cache_lock.release() + + with pytest.raises(asyncio.CancelledError): + await asyncio.gather(invocation) + await asyncio.wait_for(handler.aclose(), timeout=1) + + assert FakeTool.instances[0].close_count == 1 + assert all(entry.active_users == 0 for entry in handler._cache.values()) + + @pytest.mark.asyncio + async def test_entry_creation_rejects_new_scopes_at_capacity(self) -> None: + handler = DefaultMCPToolHandler(cache_max_size=2) + connecting = 0 + max_connecting = 0 + capacity_reached = asyncio.Event() + release = asyncio.Event() + original_connect = FakeTool.connect + + async def gated_connect(tool: FakeTool) -> None: + nonlocal connecting, max_connecting + connecting += 1 + max_connecting = max(max_connecting, connecting) + if connecting == 2: + capacity_reached.set() + try: + await release.wait() + await original_connect(tool) + finally: + connecting -= 1 + + with _patch_tool(), patch.object(FakeTool, "connect", gated_connect): + first = asyncio.create_task(handler.invoke_tool(_invocation(workflow_session_id="workflow-a"))) + second = asyncio.create_task(handler.invoke_tool(_invocation(workflow_session_id="workflow-b"))) + await capacity_reached.wait() + try: + rejected = await handler.invoke_tool(_invocation(workflow_session_id="workflow-c")) + assert rejected.is_error + assert rejected.error_message is not None + assert "capacity" in rejected.error_message + assert len(FakeTool.instances) == 2 + finally: + release.set() + results = await asyncio.gather(first, second) + await handler.aclose() + + assert all(not result.is_error for result in results) + assert max_connecting == 2 + assert not handler._inflight + + @pytest.mark.asyncio + async def test_cancelled_waiter_does_not_cancel_shared_inflight_creation(self) -> None: + handler = DefaultMCPToolHandler() + connect_started = asyncio.Event() + release_connect = asyncio.Event() + original_connect = FakeTool.connect + + async def gated_connect(tool: FakeTool) -> None: + connect_started.set() + await release_connect.wait() + await original_connect(tool) + + with _patch_tool(), patch.object(FakeTool, "connect", gated_connect): + creator = asyncio.create_task(handler.invoke_tool(_invocation(workflow_session_id="workflow-a"))) + await connect_started.wait() + waiter = asyncio.create_task(handler.invoke_tool(_invocation(workflow_session_id="workflow-a"))) + await asyncio.sleep(0) + + waiter.cancel() + with pytest.raises(asyncio.CancelledError): + waiter_result = await waiter + assert waiter_result is None + + release_connect.set() + creator_result = await creator + follow_up_result = await handler.invoke_tool(_invocation(workflow_session_id="workflow-a")) + + assert not creator_result.is_error + assert not follow_up_result.is_error + assert len(FakeTool.instances) == 1 + assert FakeTool.instances[0].connect_count == 1 + + @pytest.mark.asyncio + async def test_cancelled_creator_does_not_cancel_shared_inflight_waiter(self) -> None: + handler = DefaultMCPToolHandler() + first_connect_started = asyncio.Event() + release_first_connect = asyncio.Event() + original_connect = FakeTool.connect + + async def gated_connect(tool: FakeTool) -> None: + if len(FakeTool.instances) == 1: + first_connect_started.set() + await release_first_connect.wait() + await original_connect(tool) + + with _patch_tool(), patch.object(FakeTool, "connect", gated_connect): + creator = asyncio.create_task(handler.invoke_tool(_invocation(workflow_session_id="workflow-a"))) + await first_connect_started.wait() + waiter = asyncio.create_task(handler.invoke_tool(_invocation(workflow_session_id="workflow-a"))) + await asyncio.sleep(0) + + creator.cancel() + release_first_connect.set() + with pytest.raises(asyncio.CancelledError): + _ = await creator + + waiter_result = await waiter + follow_up_result = await handler.invoke_tool(_invocation(workflow_session_id="workflow-a")) + + assert not waiter_result.is_error + assert not follow_up_result.is_error + assert len(FakeTool.instances) == 2 + assert FakeTool.instances[0].close_count == 1 + assert FakeTool.instances[1].connect_count == 1 + + @pytest.mark.asyncio + async def test_repeated_connect_cancellation_retains_slot_until_cleanup(self) -> None: + handler = DefaultMCPToolHandler(cache_max_size=1) + connect_started = asyncio.Event() + close_started = asyncio.Event() + release_close = asyncio.Event() + close_finished = asyncio.Event() + close_tasks: list[asyncio.Task[None]] = [] + + async def connect(_tool: FakeTool) -> None: + connect_started.set() + await asyncio.Event().wait() + + async def finish_close() -> None: + close_started.set() + await release_close.wait() + close_finished.set() + + async def close(_tool: FakeTool) -> None: + task = asyncio.create_task(finish_close()) + close_tasks.append(task) + await asyncio.shield(task) + + with _patch_tool(), patch.object(FakeTool, "connect", connect), patch.object(FakeTool, "close", close): + creator = asyncio.create_task(handler.invoke_tool(_invocation(workflow_session_id="first"))) + shutdown: asyncio.Task[None] | None = None + try: + await connect_started.wait() + creator.cancel() + await close_started.wait() + creator.cancel() + for _ in range(10): + await asyncio.sleep(0) + + assert not creator.done() + result = await handler.invoke_tool(_invocation(workflow_session_id="second")) + assert result.is_error + assert len(FakeTool.instances) == 1 + shutdown = asyncio.create_task(handler.aclose()) + await asyncio.sleep(0) + assert not shutdown.done() + assert not close_finished.is_set() + finally: + release_close.set() + with pytest.raises(asyncio.CancelledError): + await creator + await asyncio.gather(*close_tasks) + if shutdown is not None: + await shutdown + else: + await handler.aclose() + + @pytest.mark.asyncio + async def test_repeated_creator_cancellation_completes_inflight_failure(self) -> None: + handler = DefaultMCPToolHandler() + + class GatedCreationSemaphore: + def __init__(self) -> None: + self.entered = asyncio.Event() + self.release = asyncio.Event() + + async def __aenter__(self) -> None: + self.entered.set() + await self.release.wait() + + async def __aexit__(self, *_args: Any) -> None: + pass + + creation_gate: Any = GatedCreationSemaphore() + handler._creation_semaphore = creation_gate + + with _patch_tool(): + creator = asyncio.create_task(handler.invoke_tool(_invocation(workflow_session_id="workflow-a"))) + await creation_gate.entered.wait() + assert handler._inflight + waiter = asyncio.create_task(handler.invoke_tool(_invocation(workflow_session_id="workflow-a"))) + await asyncio.sleep(0) + + await handler._cache_lock.acquire() + creator.cancel() + await asyncio.sleep(0) + creator.cancel() + handler._cache_lock.release() + creation_gate.release.set() + + with pytest.raises(asyncio.CancelledError): + await asyncio.gather(creator) + waiter_result = await asyncio.wait_for(waiter, timeout=1) + await asyncio.wait_for(handler.aclose(), timeout=1) + + assert not waiter_result.is_error + assert not handler._inflight + @pytest.mark.asyncio async def test_repeated_use_keeps_lru_alive(self) -> None: handler = DefaultMCPToolHandler(cache_max_size=2) @@ -953,10 +1623,10 @@ async def slow_connect(self: FakeTool) -> None: with _patch_tool(), patch.object(FakeTool, "connect", slow_connect): results = await asyncio.gather( - handler.invoke_tool(_invocation(headers={"X": "1"})), - handler.invoke_tool(_invocation(headers={"X": "1"})), - handler.invoke_tool(_invocation(headers={"X": "1"})), - handler.invoke_tool(_invocation(headers={"X": "1"})), + handler.invoke_tool(_invocation(headers={"X": "1"}, workflow_session_id="workflow-a")), + handler.invoke_tool(_invocation(headers={"X": "1"}, workflow_session_id="workflow-a")), + handler.invoke_tool(_invocation(headers={"X": "1"}, workflow_session_id="workflow-a")), + handler.invoke_tool(_invocation(headers={"X": "1"}, workflow_session_id="workflow-a")), ) assert all(not r.is_error for r in results) # Only one tool was created and connected, despite 4 concurrent calls. @@ -1064,6 +1734,93 @@ async def test_aclose_is_idempotent(self) -> None: await handler.aclose() assert FakeTool.instances[0].close_count == 1 + @pytest.mark.asyncio + async def test_cancelled_aclose_continues_draining_cached_entries(self) -> None: + handler = DefaultMCPToolHandler() + first_close_started = asyncio.Event() + release_first_close = asyncio.Event() + original_close = FakeTool.close + + async def gated_close(tool: FakeTool) -> None: + if tool.kwargs["url"] == "https://a/": + first_close_started.set() + await release_first_close.wait() + await original_close(tool) + + with _patch_tool(), patch.object(FakeTool, "close", gated_close): + await handler.invoke_tool(_invocation(server_url="https://a/", headers={"X": "1"})) + await handler.invoke_tool(_invocation(server_url="https://b/", headers={"X": "1"})) + + shutdown = asyncio.create_task(handler.aclose()) + await first_close_started.wait() + shutdown.cancel() + with pytest.raises(asyncio.CancelledError): + await asyncio.gather(shutdown) + + release_first_close.set() + await asyncio.wait_for(handler.aclose(), timeout=1) + + assert [tool.close_count for tool in FakeTool.instances] == [1, 1] + assert all(tool._httpx_client is not None and tool._httpx_client.is_closed for tool in FakeTool.instances) + + @pytest.mark.asyncio + async def test_cleanup_cancellation_continues_draining_cached_entries(self) -> None: + handler = DefaultMCPToolHandler() + original_close = FakeTool.close + + async def cancel_first_close(tool: FakeTool) -> None: + if tool.kwargs["url"] == "https://a/": + tool.close_count += 1 + raise asyncio.CancelledError + await original_close(tool) + + with _patch_tool(), patch.object(FakeTool, "close", cancel_first_close): + await handler.invoke_tool(_invocation(server_url="https://a/", headers={"X": "1"})) + await handler.invoke_tool(_invocation(server_url="https://b/", headers={"X": "1"})) + + with pytest.raises(asyncio.CancelledError): + await handler.aclose() + with pytest.raises(asyncio.CancelledError): + await handler.aclose() + + assert [tool.close_count for tool in FakeTool.instances] == [1, 1] + assert all(tool._httpx_client is not None and tool._httpx_client.is_closed for tool in FakeTool.instances) + + @pytest.mark.asyncio + async def test_active_entry_cleanup_cancellation_is_reported_by_aclose(self) -> None: + handler = DefaultMCPToolHandler() + invocation_started = asyncio.Event() + release_invocation = asyncio.Event() + + async def gated_call(_tool: FakeTool, _tool_name: str, **_arguments: Any) -> Any: + invocation_started.set() + await release_invocation.wait() + return [Content.from_text("ok")] + + async def cancelled_close(tool: FakeTool) -> None: + tool.close_count += 1 + raise asyncio.CancelledError + + with ( + _patch_tool(), + patch.object(FakeTool, "call_tool", gated_call), + patch.object(FakeTool, "close", cancelled_close), + ): + invocation = asyncio.create_task(handler.invoke_tool(_invocation(headers={"X": "1"}))) + await invocation_started.wait() + shutdown = asyncio.create_task(handler.aclose()) + await asyncio.sleep(0) + assert handler._closed + release_invocation.set() + + with pytest.raises(asyncio.CancelledError): + await asyncio.gather(invocation) + with pytest.raises(asyncio.CancelledError): + await asyncio.gather(shutdown) + + assert FakeTool.instances[0].close_count == 1 + assert not handler._retired + @pytest.mark.asyncio async def test_invoke_after_close_returns_error_result(self) -> None: """Post-close ``invoke_tool`` surfaces a tool error rather than crashing.""" @@ -1114,6 +1871,194 @@ async def gated_connect(self: FakeTool) -> None: assert result.is_error is True assert "closed" in (result.error_message or "").lower() + async def test_aclose_reports_inflight_creation_cleanup_cancellation(self) -> None: + handler = DefaultMCPToolHandler() + connect_started = asyncio.Event() + release_connect = asyncio.Event() + original_connect = FakeTool.connect + + async def gated_connect(self: FakeTool) -> None: + connect_started.set() + await release_connect.wait() + await original_connect(self) + + async def cancelled_close(tool: FakeTool) -> None: + tool.close_count += 1 + raise asyncio.CancelledError + + with ( + _patch_tool(), + patch.object(FakeTool, "connect", gated_connect), + patch.object(FakeTool, "close", cancelled_close), + ): + invocation = asyncio.create_task(handler.invoke_tool(_invocation(headers={"X": "1"}))) + await connect_started.wait() + shutdown = asyncio.create_task(handler.aclose()) + await asyncio.sleep(0) + release_connect.set() + + with pytest.raises(asyncio.CancelledError): + _ = await invocation + await asyncio.wait({shutdown}) + with pytest.raises(asyncio.CancelledError): + shutdown.result() + + assert FakeTool.instances[0].close_count == 1 + assert not handler._inflight + assert not handler._inflight_cleanup + + async def test_aclose_reports_failed_handshake_cleanup_cancellation(self) -> None: + handler = DefaultMCPToolHandler() + connect_started = asyncio.Event() + release_connect = asyncio.Event() + + async def failed_connect(_tool: FakeTool) -> None: + connect_started.set() + await release_connect.wait() + raise RuntimeError("handshake failed") + + async def cancelled_close(tool: FakeTool) -> None: + tool.close_count += 1 + raise asyncio.CancelledError + + with ( + _patch_tool(), + patch.object(FakeTool, "connect", failed_connect), + patch.object(FakeTool, "close", cancelled_close), + ): + invocation = asyncio.create_task(handler.invoke_tool(_invocation(headers={"X": "1"}))) + await connect_started.wait() + shutdown = asyncio.create_task(handler.aclose()) + await asyncio.sleep(0) + release_connect.set() + + with pytest.raises(asyncio.CancelledError): + _ = await invocation + await asyncio.wait({shutdown}) + with pytest.raises(asyncio.CancelledError): + shutdown.result() + + assert FakeTool.instances[0].close_count == 1 + assert not handler._inflight + assert not handler._inflight_cleanup + + @pytest.mark.asyncio + async def test_cancelled_creator_does_not_block_aclose(self) -> None: + handler = DefaultMCPToolHandler() + connect_started = asyncio.Event() + release_connect = asyncio.Event() + close_started = asyncio.Event() + release_close = asyncio.Event() + original_connect = FakeTool.connect + original_close_entry = handler._close_entry + + async def gated_connect(self: FakeTool) -> None: + connect_started.set() + await release_connect.wait() + await original_connect(self) + + async def gated_close_entry(entry: Any) -> None: + close_started.set() + await release_close.wait() + await original_close_entry(entry) + + with ( + _patch_tool(), + patch.object(FakeTool, "connect", gated_connect), + patch.object(handler, "_close_entry", gated_close_entry), + ): + invoke_task = asyncio.create_task(handler.invoke_tool(_invocation(headers={"X": "1"}))) + await connect_started.wait() + close_task = asyncio.create_task(handler.aclose()) + await asyncio.sleep(0) + release_connect.set() + await close_started.wait() + invoke_task.cancel() + release_close.set() + + with pytest.raises(asyncio.CancelledError): + _ = await invoke_task + await asyncio.wait_for(close_task, timeout=1) + + assert FakeTool.instances[0].close_count == 1 + assert not handler._inflight + + @pytest.mark.asyncio + async def test_cancellation_while_waiting_for_phase_three_lock_cleans_up(self) -> None: + handler = DefaultMCPToolHandler() + entry_created = asyncio.Event() + release_creation = asyncio.Event() + original_create_entry = handler._create_entry + + async def gated_create_entry(invocation: MCPToolInvocation) -> Any: + entry = await original_create_entry(invocation) + entry_created.set() + await release_creation.wait() + return entry + + with _patch_tool(), patch.object(handler, "_create_entry", gated_create_entry): + invocation = asyncio.create_task(handler.invoke_tool(_invocation(headers={"X": "1"}))) + await entry_created.wait() + await handler._cache_lock.acquire() + release_creation.set() + await asyncio.sleep(0) + invocation.cancel() + handler._cache_lock.release() + + with pytest.raises(asyncio.CancelledError): + _ = await invocation + + assert FakeTool.instances[0].close_count == 1 + assert not handler._inflight + + @pytest.mark.asyncio + async def test_aclose_reports_cancelled_abort_cleanup(self) -> None: + handler = DefaultMCPToolHandler() + entry_created = asyncio.Event() + release_creation = asyncio.Event() + close_started = asyncio.Event() + release_close = asyncio.Event() + original_create_entry = handler._create_entry + + async def gated_create_entry(invocation: MCPToolInvocation) -> Any: + entry = await original_create_entry(invocation) + entry_created.set() + await release_creation.wait() + return entry + + async def cancelled_close_entry(_entry: Any) -> None: + close_started.set() + await release_close.wait() + raise asyncio.CancelledError + + with ( + _patch_tool(), + patch.object(handler, "_create_entry", gated_create_entry), + patch.object(handler, "_close_entry", cancelled_close_entry), + ): + invocation = asyncio.create_task(handler.invoke_tool(_invocation(headers={"X": "1"}))) + await entry_created.wait() + await handler._cache_lock.acquire() + release_creation.set() + await asyncio.sleep(0) + invocation.cancel() + handler._cache_lock.release() + await close_started.wait() + + shutdown = asyncio.create_task(handler.aclose()) + await asyncio.sleep(0) + assert not shutdown.done() + release_close.set() + + with pytest.raises(asyncio.CancelledError): + _ = await invocation + await asyncio.wait({shutdown}) + with pytest.raises(asyncio.CancelledError): + shutdown.result() + + assert not handler._inflight + assert not handler._inflight_cleanup + # ---------- Result normalisation ------------------------------------------ @@ -1238,38 +2183,43 @@ def boom(**_a: Any) -> Any: class TestCacheKey: def test_key_order_independent(self) -> None: - k1 = DefaultMCPToolHandler._cache_key("https://x/", None, None, {"A": "1", "B": "2"}) - k2 = DefaultMCPToolHandler._cache_key("https://x/", None, None, {"B": "2", "A": "1"}) + k1 = DefaultMCPToolHandler._cache_key("workflow", "https://x/", None, None, {"A": "1", "B": "2"}) + k2 = DefaultMCPToolHandler._cache_key("workflow", "https://x/", None, None, {"B": "2", "A": "1"}) assert k1 == k2 def test_key_distinguishes_values(self) -> None: - k1 = DefaultMCPToolHandler._cache_key("https://x/", None, None, {"A": "1"}) - k2 = DefaultMCPToolHandler._cache_key("https://x/", None, None, {"A": "2"}) + k1 = DefaultMCPToolHandler._cache_key("workflow", "https://x/", None, None, {"A": "1"}) + k2 = DefaultMCPToolHandler._cache_key("workflow", "https://x/", None, None, {"A": "2"}) assert k1 != k2 def test_empty_headers_use_fixed_hash(self) -> None: - k1 = DefaultMCPToolHandler._cache_key("https://x/", None, None, None) - k2 = DefaultMCPToolHandler._cache_key("https://x/", None, None, {}) + k1 = DefaultMCPToolHandler._cache_key("workflow", "https://x/", None, None, None) + k2 = DefaultMCPToolHandler._cache_key("workflow", "https://x/", None, None, {}) assert k1 == k2 + def test_key_distinguishes_workflow_session(self) -> None: + k1 = DefaultMCPToolHandler._cache_key("workflow-a", "https://x/", None, None, None) + k2 = DefaultMCPToolHandler._cache_key("workflow-b", "https://x/", None, None, None) + assert k1 != k2 + def test_key_distinguishes_connection_name(self) -> None: - k1 = DefaultMCPToolHandler._cache_key("https://x/", None, "conn-A", None) - k2 = DefaultMCPToolHandler._cache_key("https://x/", None, "conn-B", None) + k1 = DefaultMCPToolHandler._cache_key("workflow", "https://x/", None, "conn-A", None) + k2 = DefaultMCPToolHandler._cache_key("workflow", "https://x/", None, "conn-B", None) assert k1 != k2 def test_key_distinguishes_server_label(self) -> None: - k1 = DefaultMCPToolHandler._cache_key("https://x/", "Lbl-A", None, None) - k2 = DefaultMCPToolHandler._cache_key("https://x/", "Lbl-B", None, None) + k1 = DefaultMCPToolHandler._cache_key("workflow", "https://x/", "Lbl-A", None, None) + k2 = DefaultMCPToolHandler._cache_key("workflow", "https://x/", "Lbl-B", None, None) assert k1 != k2 def test_key_collapses_header_name_case(self) -> None: - k1 = DefaultMCPToolHandler._cache_key("https://x/", None, None, {"Authorization": "tk"}) - k2 = DefaultMCPToolHandler._cache_key("https://x/", None, None, {"authorization": "tk"}) + k1 = DefaultMCPToolHandler._cache_key("workflow", "https://x/", None, None, {"Authorization": "tk"}) + k2 = DefaultMCPToolHandler._cache_key("workflow", "https://x/", None, None, {"authorization": "tk"}) assert k1 == k2 def test_key_keeps_header_value_case(self) -> None: - k1 = DefaultMCPToolHandler._cache_key("https://x/", None, None, {"X": "Bearer-A"}) - k2 = DefaultMCPToolHandler._cache_key("https://x/", None, None, {"X": "bearer-a"}) + k1 = DefaultMCPToolHandler._cache_key("workflow", "https://x/", None, None, {"X": "Bearer-A"}) + k2 = DefaultMCPToolHandler._cache_key("workflow", "https://x/", None, None, {"X": "bearer-a"}) assert k1 != k2 diff --git a/python/packages/declarative/tests/test_function_tool_executor.py b/python/packages/declarative/tests/test_function_tool_executor.py index cd5c42aa290..207d8ff1839 100644 --- a/python/packages/declarative/tests/test_function_tool_executor.py +++ b/python/packages/declarative/tests/test_function_tool_executor.py @@ -763,11 +763,7 @@ def mock_state() -> MagicMock: mock_state._data = {} def mock_get(key: str, default: Any = None) -> Any: - if key not in mock_state._data: - if default is not None: - return default - raise KeyError(key) - return mock_state._data[key] + return mock_state._data.get(key, default) def mock_set(key: str, value: Any) -> None: mock_state._data[key] = value @@ -1205,6 +1201,7 @@ def my_tool(x: int) -> int: request_id="req-123", function_name="my_tool", arguments={"x": 7}, + workflow_session_id="workflow-a", ) response = ToolApprovalResponse(approved=True) @@ -1241,6 +1238,7 @@ def my_tool(x: int) -> int: request_id="req-456", function_name="my_tool", arguments={"x": 5}, + workflow_session_id="workflow-a", ) response = ToolApprovalResponse(approved=False, reason="Not authorized") diff --git a/python/packages/declarative/tests/test_graph_coverage.py b/python/packages/declarative/tests/test_graph_coverage.py index 7187b317f8c..c3801b327da 100644 --- a/python/packages/declarative/tests/test_graph_coverage.py +++ b/python/packages/declarative/tests/test_graph_coverage.py @@ -2808,6 +2808,7 @@ async def test_handle_external_input_response_no_state(self, mock_context, mock_ agent_name="TestAgent", agent_response="Hello", iteration=1, + workflow_session_id="workflow-a", ) response = AgentExternalInputResponse(user_input="hi there") @@ -2856,6 +2857,7 @@ async def test_handle_external_input_response_agent_not_found(self, mock_context agent_name="NonExistentAgent", agent_response="Hello", iteration=1, + workflow_session_id="workflow-a", ) response = AgentExternalInputResponse(user_input="continue") diff --git a/python/packages/declarative/tests/test_invoke_mcp_tool_executor.py b/python/packages/declarative/tests/test_invoke_mcp_tool_executor.py index b8c2e010c96..25fc8bf62ee 100644 --- a/python/packages/declarative/tests/test_invoke_mcp_tool_executor.py +++ b/python/packages/declarative/tests/test_invoke_mcp_tool_executor.py @@ -9,7 +9,10 @@ """ import sys +from collections.abc import Callable +from pathlib import Path from typing import Any +from unittest.mock import AsyncMock, MagicMock import httpx2 as httpx import pytest @@ -26,9 +29,8 @@ reason="PowerFx engine not available (requires dotnet runtime)", ) -from agent_framework import Content, Message # noqa: E402 +from agent_framework import AgentSession, CheckpointStorage, Content, Message, Workflow # noqa: E402 from agent_framework.exceptions import ToolExecutionException # noqa: E402 - from agent_framework_declarative._workflows import ( # noqa: E402 DECLARATIVE_STATE_KEY, DeclarativeWorkflowError, @@ -115,6 +117,29 @@ def _yaml(action: dict[str, Any]) -> dict[str, Any]: return {"name": "mcp_test", "actions": [action]} +def _build_manual_mcp_workflow( + mcp_handler: MCPToolHandler, *, depth: int, require_approval: bool = False +) -> tuple[Workflow, Workflow]: + from agent_framework import WorkflowBuilder, WorkflowExecutor + from agent_framework_declarative._workflows._executors_mcp import InvokeMcpToolActionExecutor + + leaf = WorkflowBuilder( + name="manual_leaf", + start_executor=InvokeMcpToolActionExecutor( + _action(require_approval=require_approval), mcp_tool_handler=mcp_handler + ), + ).build() + root = leaf + for level in range(depth): + root = WorkflowBuilder( + name=f"manual_parent_{level}", + start_executor=WorkflowExecutor( + root, id=f"child_{level}", propagate_request=True, allow_direct_output=True + ), + ).build() + return root, leaf + + # ---------- Builder enforcement -------------------------------------------- @@ -149,6 +174,102 @@ def test_missing_tool_name_fails_validation(self) -> None: class TestFieldForwarding: + @pytest.mark.parametrize("streaming", [False, True]) + @pytest.mark.parametrize("depth", [0, 1, 2]) + async def test_manual_workflow_scopes_follow_core_run_context(self, streaming: bool, depth: int) -> None: + from agent_framework import InMemoryCheckpointStorage + from agent_framework._workflows._const import WORKFLOW_SESSION_ID_KEY + + mcp_handler = StubMcpHandler(_ok()) + storage = InMemoryCheckpointStorage() + root, leaf = _build_manual_mcp_workflow(mcp_handler, depth=depth) + expected_scopes: list[str] = [] + entry_id: str | None = None + for index in range(2): + if streaming: + await root.run({}, checkpoint_storage=storage, stream=True).get_final_response() + else: + await root.run({}, checkpoint_storage=storage) + expected_scopes.append(leaf._runner.state.get(WORKFLOW_SESSION_ID_KEY)) + if index == 0: + [entry] = [ + checkpoint + for checkpoint in await storage.list_checkpoints(workflow_name=root.name) + if checkpoint.iteration_count == 0 + ] + entry_id = entry.checkpoint_id + + actual_scopes = [invocation.workflow_session_id for invocation in mcp_handler.invocations] + assert len(set(actual_scopes)) == 2 + assert actual_scopes == expected_scopes + + agent = root.as_agent() + for session in (None, AgentSession()): + for _ in range(2): + if streaming: + await agent.run("agent turn", session=session, stream=True).get_final_response() + else: + await agent.run("agent turn", session=session) + expected_scopes.append(leaf._runner.state.get(WORKFLOW_SESSION_ID_KEY)) + assert expected_scopes[-1] == expected_scopes[-2] + assert len(set(expected_scopes)) == 4 + assert [invocation.workflow_session_id for invocation in mcp_handler.invocations] == expected_scopes + + assert entry_id is not None + root, leaf = _build_manual_mcp_workflow(mcp_handler, depth=depth) + if streaming: + await root.run(checkpoint_id=entry_id, checkpoint_storage=storage, stream=True).get_final_response() + else: + await root.run(checkpoint_id=entry_id, checkpoint_storage=storage) + assert mcp_handler.invocations[-1].workflow_session_id == expected_scopes[0] + assert leaf._runner.state.get(WORKFLOW_SESSION_ID_KEY) == expected_scopes[0] + + @pytest.mark.parametrize("streaming", [False, True]) + @pytest.mark.parametrize("restore", [False, True]) + async def test_manual_approval_continuations_restore_core_scopes(self, streaming: bool, restore: bool) -> None: + from agent_framework import InMemoryCheckpointStorage, WorkflowEvent + from agent_framework._workflows._const import WORKFLOW_SESSION_ID_KEY + from agent_framework_declarative._workflows import ToolApprovalResponse + + mcp_handler = StubMcpHandler(_ok()) + storage = InMemoryCheckpointStorage() + root, leaf = _build_manual_mcp_workflow(mcp_handler, depth=1, require_approval=True) + requests: list[WorkflowEvent[Any]] = [] + for _ in range(2): + if streaming: + await root.run({}, checkpoint_storage=storage, stream=True).get_final_response() + else: + await root.run({}, checkpoint_storage=storage) + [request] = [ + request + for request in (await root._runner_context.get_pending_request_info_events()).values() + if request.request_id not in {previous.request_id for previous in requests} + ] + requests.append(request) + scopes = [request.data.workflow_session_id for request in requests] + assert len(set(scopes)) == 2 + assert mcp_handler.call_count == 0 + if restore: + checkpoint_id = root.get_last_checkpoint_id() + assert checkpoint_id is not None + root, leaf = _build_manual_mcp_workflow(mcp_handler, depth=1, require_approval=True) + await root.run(checkpoint_id=checkpoint_id, checkpoint_storage=storage) + + with pytest.raises(ValueError, match="multiple parent workflow scopes"): + await root.run(responses={request.request_id: ToolApprovalResponse(approved=True) for request in requests}) + assert mcp_handler.call_count == 0 + assert len(await root._runner_context.get_pending_request_info_events()) == 2 + for request, scope in reversed(list(zip(requests, scopes, strict=True))): + responses = {request.request_id: ToolApprovalResponse(approved=True)} + if streaming: + await root.run(responses=responses, stream=True).get_final_response() + else: + await root.run(responses=responses) + assert mcp_handler.invocations[-1].workflow_session_id == scope + assert leaf._runner.state.get(WORKFLOW_SESSION_ID_KEY) == scope + assert not await root._runner_context.get_pending_request_info_events() + assert not await leaf._runner_context.get_pending_request_info_events() + @pytest.mark.asyncio async def test_basic_invocation_forwards_required_fields(self) -> None: handler = StubMcpHandler(_ok()) @@ -164,6 +285,969 @@ async def test_basic_invocation_forwards_required_fields(self) -> None: assert inv.headers == {} assert inv.arguments == {} assert inv.connection_name is None + assert inv.workflow_session_id + + @pytest.mark.asyncio + async def test_separate_workflows_receive_separate_session_ids(self) -> None: + handler = StubMcpHandler(_ok()) + factory = WorkflowFactory(mcp_tool_handler=handler) + first = factory.create_workflow_from_definition(_yaml(_action())) + second = factory.create_workflow_from_definition(_yaml(_action())) + + await first.run({}) + await second.run({}) + + assert len(handler.invocations) == 2 + assert handler.invocations[0].workflow_session_id + assert handler.invocations[1].workflow_session_id + assert handler.invocations[0].workflow_session_id != handler.invocations[1].workflow_session_id + + @pytest.mark.parametrize("streaming", [False, True]) + @pytest.mark.parametrize("restore_kind", ["memory", "file", "queued"]) + async def test_same_id_coordinator_relay_preserves_downstream_mcp_scope( + self, + tmp_path: Path, + streaming: bool, + restore_kind: str, + coordinator_relay_workflow: Callable[[Workflow, CheckpointStorage], Workflow], + ) -> None: + mcp_handler = StubMcpHandler(_ok()) + factory = WorkflowFactory(mcp_tool_handler=mcp_handler) + storage = _scope_checkpoint_storage(tmp_path) + + parent = coordinator_relay_workflow(factory.create_workflow_from_definition(_yaml(_action())), storage) + requests = [] + for message in (1, 2): + if streaming: + result = await parent.run(message, stream=True).get_final_response() + else: + result = await parent.run(message) + [request] = [event for event in result.get_request_info_events() if event.request_id not in requests] + requests.append(request.request_id) + scopes = [invocation.workflow_session_id for invocation in mcp_handler.invocations] + assert len(scopes) == 2 + assert scopes[0] and scopes[1] and scopes[0] != scopes[1] + if restore_kind == "file": + checkpoint_id = parent.get_last_checkpoint_id() + assert checkpoint_id is not None + parent = coordinator_relay_workflow(factory.create_workflow_from_definition(_yaml(_action())), storage) + await parent.run(checkpoint_id=checkpoint_id) + + for request_id, scope in reversed(list(zip(requests, scopes, strict=True))): + resume_checkpoint_id: str | None = None + responses: dict[str, str] | None = {request_id: "continue"} + if restore_kind == "queued": + await parent._runner_context.send_request_info_response(request_id, "continue") + await parent._runner.create_checkpoint_if_enabled() + resume_checkpoint_id = parent.get_last_checkpoint_id() + assert resume_checkpoint_id is not None + parent = coordinator_relay_workflow(factory.create_workflow_from_definition(_yaml(_action())), storage) + responses = None + if streaming: + await parent.run( + responses=responses, checkpoint_id=resume_checkpoint_id, stream=True + ).get_final_response() + else: + await parent.run(responses=responses, checkpoint_id=resume_checkpoint_id) + assert mcp_handler.invocations[-1].workflow_session_id == scope + assert len(mcp_handler.invocations) == 4 + assert not await parent._runner_context.get_pending_request_info_events() + + @pytest.mark.parametrize("streaming", [False, True]) + @pytest.mark.parametrize("restore_kind", ["memory", "file", "queued"]) + async def test_core_parent_request_restores_origin_before_child( + self, tmp_path: Path, streaming: bool, restore_kind: str + ) -> None: + from agent_framework import ( + Executor, + Workflow, + WorkflowBuilder, + WorkflowContext, + WorkflowExecutor, + response_handler, + ) + from agent_framework import handler as workflow_handler + from agent_framework._workflows._const import WORKFLOW_SESSION_ID_KEY + + mcp_handler = StubMcpHandler(_ok()) + factory = WorkflowFactory(mcp_tool_handler=mcp_handler) + storage = _scope_checkpoint_storage(tmp_path) + + class Requester(Executor): + @workflow_handler + async def start(self, message: str, ctx: WorkflowContext) -> None: + await ctx.request_info(message, str) + + @response_handler + async def respond(self, original: str, response: str, ctx: WorkflowContext[dict[str, Any]]) -> None: + await ctx.send_message({"original": original, "response": response}) + + def build() -> Workflow: + start = Requester(id="requester") + child = WorkflowExecutor(factory.create_workflow_from_definition(_yaml(_action())), id="child") + return WorkflowBuilder(start_executor=start, checkpoint_storage=storage).add_edge(start, child).build() + + parent = build() + await parent.run("A") + [request_a] = (await parent._runner_context.get_pending_request_info_events()).values() + scope_a = parent._runner.state.get(WORKFLOW_SESSION_ID_KEY) + await parent.run("B") + requests = await parent._runner_context.get_pending_request_info_events() + [request_b] = [request for request in requests.values() if request.request_id != request_a.request_id] + if restore_kind == "file": + checkpoint_id = parent.get_last_checkpoint_id() + assert checkpoint_id is not None + parent = build() + await parent.run(checkpoint_id=checkpoint_id) + + with pytest.raises(ValueError, match="multiple parent workflow scopes"): + await parent.run(responses={request_a.request_id: "continue", request_b.request_id: "continue"}) + assert set(await parent._runner_context.get_pending_request_info_events()) == set(requests) + assert mcp_handler.call_count == 0 + await parent.run(responses={request_b.request_id: "continue"}) + if restore_kind == "queued": + await parent._runner_context.send_request_info_response(request_a.request_id, "continue") + await parent._runner.create_checkpoint_if_enabled() + checkpoint_id = parent.get_last_checkpoint_id() + assert checkpoint_id is not None + parent = build() + if streaming: + await parent.run(checkpoint_id=checkpoint_id, stream=True).get_final_response() + else: + await parent.run(checkpoint_id=checkpoint_id) + elif streaming: + await parent.run(responses={request_a.request_id: "continue"}, stream=True).get_final_response() + else: + await parent.run(responses={request_a.request_id: "continue"}) + + assert parent._runner.state.get(WORKFLOW_SESSION_ID_KEY) == scope_a + assert len(mcp_handler.invocations) == 2 + assert mcp_handler.invocations[0].workflow_session_id != mcp_handler.invocations[1].workflow_session_id + assert not await parent._runner_context.get_pending_request_info_events() + + @pytest.mark.parametrize("message_list", [False, True]) + @pytest.mark.parametrize("streaming", [False, True]) + async def test_fresh_runs_on_same_workflow_receive_separate_session_ids( + self, message_list: bool, streaming: bool + ) -> None: + handler = StubMcpHandler(_ok()) + factory = WorkflowFactory(mcp_tool_handler=handler) + workflow = factory.create_workflow_from_definition(_yaml(_action())) + + for text in ("first run", "second run"): + inputs: list[Message] | dict[str, Any] = ( + [Message(role="user", contents=[Content.from_text(text)])] if message_list else {} + ) + if streaming: + await workflow.run(inputs, stream=True).get_final_response() + else: + await workflow.run(inputs) + + assert len(handler.invocations) == 2 + assert handler.invocations[0].workflow_session_id + assert handler.invocations[0].workflow_session_id != handler.invocations[1].workflow_session_id + + @pytest.mark.parametrize("streaming", [False, True]) + async def test_as_agent_continuation_reuses_workflow_session_id(self, streaming: bool) -> None: + handler = StubMcpHandler(_ok()) + factory = WorkflowFactory(mcp_tool_handler=handler) + workflow = factory.create_workflow_from_definition(_yaml(_action())) + agent = workflow.as_agent(name="mcp-agent") + + for text in ("first turn", "second turn"): + if streaming: + await agent.run(text, stream=True).get_final_response() + else: + await agent.run(text) + + assert len(handler.invocations) == 2 + assert handler.invocations[0].workflow_session_id + assert handler.invocations[0].workflow_session_id == handler.invocations[1].workflow_session_id + + @pytest.mark.parametrize("pause_kind", ["Question", "RequestExternalInput"]) + @pytest.mark.parametrize("streaming", [False, True]) + async def test_response_only_agent_checkpoint_resume_preserves_workflow_session_id( + self, pause_kind: str, streaming: bool + ) -> None: + from agent_framework import InMemoryCheckpointStorage + from agent_framework_declarative._workflows import ExternalInputResponse + + storage = InMemoryCheckpointStorage() + handler = StubMcpHandler(_ok()) + factory = WorkflowFactory(mcp_tool_handler=handler) + definition = { + "name": "mcp_continuation", + "actions": [ + {**_action(tool_name="before"), "id": "before"}, + { + "kind": pause_kind, + "id": "pause", + "question" if pause_kind == "Question" else "prompt": "Continue?", + "variable": "Local.answer", + }, + {**_action(tool_name="after"), "id": "after"}, + ], + } + workflow = factory.create_workflow_from_definition(definition) + agent = workflow.as_agent() + original_session = AgentSession() + if streaming: + response = await agent.run( + "first turn", session=original_session, stream=True, checkpoint_storage=storage + ).get_final_response() + else: + response = await agent.run("first turn", session=original_session, checkpoint_storage=storage) + + [request] = (await workflow._runner_context.get_pending_request_info_events()).values() + [function_call] = [ + content for message in response.messages for content in message.contents if content.type == "function_call" + ] + assert function_call.call_id == request.request_id + checkpoint_id = workflow.get_last_checkpoint_id() + assert checkpoint_id is not None + workflow = factory.create_workflow_from_definition(definition) + agent = workflow.as_agent() + resumed_session = AgentSession() + assert resumed_session.session_id != original_session.session_id + reply = Message( + role="user", + contents=[ + Content( + "function_result", + call_id=request.request_id, + result=ExternalInputResponse(user_input="continue"), + ) + ], + ) + if streaming: + await agent.run( + reply, + session=resumed_session, + stream=True, + checkpoint_id=checkpoint_id, + checkpoint_storage=storage, + ).get_final_response() + else: + await agent.run(reply, session=resumed_session, checkpoint_id=checkpoint_id, checkpoint_storage=storage) + + [original, resumed] = handler.invocations + assert original.tool_name == "before" + assert resumed.tool_name == "after" + assert original.workflow_session_id + assert resumed.workflow_session_id == original.workflow_session_id + assert not await workflow._runner_context.get_pending_request_info_events() + + @pytest.mark.parametrize("pause_kind", ["Question", "RequestExternalInput"]) + @pytest.mark.parametrize("streaming", [False, True]) + @pytest.mark.parametrize("restore", [False, True]) + async def test_pending_external_input_restores_originating_workflow_session_id( + self, pause_kind: str, streaming: bool, restore: bool + ) -> None: + from agent_framework import InMemoryCheckpointStorage, WorkflowEvent + from agent_framework_declarative._workflows import ExternalInputResponse + + storage = InMemoryCheckpointStorage() + handler = StubMcpHandler(_ok()) + factory = WorkflowFactory(mcp_tool_handler=handler) + definition = { + "name": "mcp_pending_input", + "actions": [ + {**_action(tool_name="before"), "id": "before"}, + { + "kind": pause_kind, + "id": "pause", + "question" if pause_kind == "Question" else "prompt": "Continue?", + "variable": "Local.answer", + }, + {**_action(tool_name="after"), "id": "after"}, + ], + } + workflow = factory.create_workflow_from_definition(definition) + requests: list[WorkflowEvent[Any]] = [] + for _ in range(2): + if streaming: + await workflow.run({}, stream=True, checkpoint_storage=storage).get_final_response() + else: + await workflow.run({}, checkpoint_storage=storage) + [request] = [ + event + for event in (await workflow._runner_context.get_pending_request_info_events()).values() + if event.request_id not in {previous.request_id for previous in requests} + ] + requests.append(request) + original_scopes = [invocation.workflow_session_id for invocation in handler.invocations] + assert all(original_scopes) + assert original_scopes[0] != original_scopes[1] + + if restore: + checkpoint_id = workflow.get_last_checkpoint_id() + assert checkpoint_id is not None + workflow = factory.create_workflow_from_definition(definition) + await workflow.run(checkpoint_id=checkpoint_id, checkpoint_storage=storage) + + for request, scope in zip(requests, original_scopes, strict=True): + responses = {request.request_id: ExternalInputResponse(user_input="continue")} + if streaming: + await workflow.run(responses=responses, stream=True).get_final_response() + else: + await workflow.run(responses=responses) + assert handler.invocations[-1].tool_name == "after" + assert handler.invocations[-1].workflow_session_id == scope + assert handler.call_count == 4 + assert not await workflow._runner_context.get_pending_request_info_events() + + @pytest.mark.parametrize("mixed_approval", [False, True]) + @pytest.mark.parametrize("streaming", [False, True]) + @pytest.mark.parametrize("restore", [False, True]) + async def test_external_input_batches_reject_different_workflow_sessions( + self, mixed_approval: bool, streaming: bool, restore: bool + ) -> None: + from agent_framework import InMemoryCheckpointStorage + from agent_framework_declarative._workflows import ExternalInputResponse, ToolApprovalResponse + + storage = InMemoryCheckpointStorage() + handler = StubMcpHandler(_ok()) + factory = WorkflowFactory(mcp_tool_handler=handler) + definition = { + "name": "mcp_input_batches", + "actions": [ + {**_action(tool_name="before"), "id": "before"}, + {"kind": "Question", "id": "pause", "question": "Continue?", "variable": "Local.answer"}, + {**_action(tool_name="after", require_approval=mixed_approval), "id": "after"}, + ], + } + workflow = factory.create_workflow_from_definition(definition) + await workflow.run({}, checkpoint_storage=storage) + [first] = (await workflow._runner_context.get_pending_request_info_events()).values() + await workflow.run({}, checkpoint_storage=storage) + [second] = [ + event + for event in (await workflow._runner_context.get_pending_request_info_events()).values() + if event.request_id != first.request_id + ] + if mixed_approval: + await workflow.run( + responses={second.request_id: ExternalInputResponse(user_input="continue")}, + checkpoint_storage=storage, + ) + [second] = [ + event + for event in (await workflow._runner_context.get_pending_request_info_events()).values() + if event.request_id != first.request_id + ] + if restore: + checkpoint_id = workflow.get_last_checkpoint_id() + assert checkpoint_id is not None + workflow = factory.create_workflow_from_definition(definition) + await workflow.run(checkpoint_id=checkpoint_id, checkpoint_storage=storage) + + responses = { + first.request_id: ExternalInputResponse(user_input="continue"), + second.request_id: ( + ToolApprovalResponse(approved=True) if mixed_approval else ExternalInputResponse(user_input="continue") + ), + } + with pytest.raises(DeclarativeWorkflowError, match="different workflow sessions"): + if streaming: + await workflow.run(responses=responses, stream=True).get_final_response() + else: + await workflow.run(responses=responses) + assert handler.call_count == 2 + assert set(await workflow._runner_context.get_pending_request_info_events()) == set(responses) + + @pytest.mark.parametrize("pause_kind", ["Question", "RequestExternalInput"]) + async def test_external_input_without_originating_scope_is_not_consumed(self, pause_kind: str) -> None: + from agent_framework_declarative._workflows import ExternalInputResponse + + handler = StubMcpHandler(_ok()) + factory = WorkflowFactory(mcp_tool_handler=handler) + workflow = factory.create_workflow_from_definition({ + "name": "mcp_legacy_input", + "actions": [ + {**_action(tool_name="before"), "id": "before"}, + { + "kind": pause_kind, + "id": "pause", + "question" if pause_kind == "Question" else "prompt": "Continue?", + "variable": "Local.answer", + }, + {**_action(tool_name="after"), "id": "after"}, + ], + }) + await workflow.run({}) + [request] = (await workflow._runner_context.get_pending_request_info_events()).values() + request.data.workflow_session_id = None + + with pytest.raises(DeclarativeWorkflowError, match="originating workflow session"): + await workflow.run(responses={request.request_id: ExternalInputResponse(user_input="continue")}) + + assert handler.call_count == 1 + assert set(await workflow._runner_context.get_pending_request_info_events()) == {request.request_id} + + @pytest.mark.parametrize("run_kind", ["direct", "sessionless_agent", "explicit_agent"]) + @pytest.mark.parametrize("streaming", [False, True]) + async def test_entry_checkpoint_replay_preserves_workflow_session_id(self, run_kind: str, streaming: bool) -> None: + from agent_framework import InMemoryCheckpointStorage + + storage = InMemoryCheckpointStorage() + handler = StubMcpHandler(_ok()) + factory = WorkflowFactory(mcp_tool_handler=handler) + definition = _yaml(_action()) + workflow = factory.create_workflow_from_definition(definition) + session = AgentSession() if run_kind == "explicit_agent" else None + + if run_kind == "direct": + if streaming: + await workflow.run({}, stream=True, checkpoint_storage=storage).get_final_response() + else: + await workflow.run({}, checkpoint_storage=storage) + else: + agent = workflow.as_agent() + if streaming: + await agent.run( + "first turn", session=session, stream=True, checkpoint_storage=storage + ).get_final_response() + else: + await agent.run("first turn", session=session, checkpoint_storage=storage) + + [original] = handler.invocations + checkpoints = await storage.list_checkpoints(workflow_name=workflow.name) + [entry_checkpoint] = [checkpoint for checkpoint in checkpoints if checkpoint.iteration_count == 0] + workflow = factory.create_workflow_from_definition(definition) + await workflow.run(checkpoint_id=entry_checkpoint.checkpoint_id, checkpoint_storage=storage) + + assert len(handler.invocations) == 2 + assert original.workflow_session_id + assert handler.invocations[1].workflow_session_id == original.workflow_session_id + + @pytest.mark.parametrize("intervening_run", ["explicit_session", "direct_dict", "direct_message_list"]) + @pytest.mark.parametrize("streaming", [False, True]) + @pytest.mark.parametrize("restore", [False, True]) + async def test_sessionless_agent_scope_survives_intervening_runs( + self, intervening_run: str, streaming: bool, restore: bool + ) -> None: + from agent_framework import InMemoryCheckpointStorage + + storage = InMemoryCheckpointStorage() + handler = StubMcpHandler(_ok()) + factory = WorkflowFactory(mcp_tool_handler=handler) + definition = _yaml(_action()) + workflow = factory.create_workflow_from_definition(definition) + agent = workflow.as_agent(name="mcp-agent") + + async def agent_turn(text: str, session: AgentSession | None = None) -> None: + if streaming: + await agent.run(text, session=session, stream=True, checkpoint_storage=storage).get_final_response() + else: + await agent.run(text, session=session, checkpoint_storage=storage) + + await agent_turn("first sessionless turn") + if intervening_run == "explicit_session": + await agent_turn("explicit session turn", AgentSession()) + else: + inputs: list[Message] | dict[str, Any] = ( + [Message(role="user", contents=[Content.from_text("direct run")])] + if intervening_run == "direct_message_list" + else {} + ) + if streaming: + await workflow.run(inputs, stream=True, checkpoint_storage=storage).get_final_response() + else: + await workflow.run(inputs, checkpoint_storage=storage) + + if restore: + checkpoint_id = workflow.get_last_checkpoint_id() + assert checkpoint_id is not None + workflow = factory.create_workflow_from_definition(definition) + await workflow.run(checkpoint_id=checkpoint_id, checkpoint_storage=storage) + agent = workflow.as_agent(name="mcp-agent") + + await agent_turn("second sessionless turn") + + session_ids = [invocation.workflow_session_id for invocation in handler.invocations] + assert len(session_ids) == 3 + assert all(session_ids) + assert session_ids[0] != session_ids[1] + assert session_ids[0] == session_ids[2] + + @pytest.mark.parametrize("message_list", [False, True]) + @pytest.mark.parametrize("streaming", [False, True]) + @pytest.mark.parametrize("restore", [False, True]) + async def test_sessionless_agent_does_not_inherit_direct_scope( + self, message_list: bool, streaming: bool, restore: bool + ) -> None: + from agent_framework import InMemoryCheckpointStorage + + storage = InMemoryCheckpointStorage() + handler = StubMcpHandler(_ok()) + factory = WorkflowFactory(mcp_tool_handler=handler) + definition = _yaml(_action()) + workflow = factory.create_workflow_from_definition(definition) + inputs: list[Message] | dict[str, Any] = ( + [Message(role="user", contents=[Content.from_text("direct run")])] if message_list else {} + ) + if streaming: + await workflow.run(inputs, stream=True, checkpoint_storage=storage).get_final_response() + else: + await workflow.run(inputs, checkpoint_storage=storage) + + if restore: + checkpoint_id = workflow.get_last_checkpoint_id() + assert checkpoint_id is not None + workflow = factory.create_workflow_from_definition(definition) + await workflow.run(checkpoint_id=checkpoint_id, checkpoint_storage=storage) + + agent = workflow.as_agent(name="mcp-agent") + for text in ("first sessionless turn", "second sessionless turn"): + if streaming: + await agent.run(text, stream=True).get_final_response() + else: + await agent.run(text) + + session_ids = [invocation.workflow_session_id for invocation in handler.invocations] + assert len(session_ids) == 3 + assert all(session_ids) + assert session_ids[0] != session_ids[1] + assert session_ids[1] == session_ids[2] + + @pytest.mark.parametrize("streaming", [False, True]) + @pytest.mark.parametrize("approved", [False, True]) + @pytest.mark.parametrize("restore", [False, True]) + async def test_sessionless_scope_survives_approval_origin_restoration( + self, streaming: bool, approved: bool, restore: bool + ) -> None: + from agent_framework import InMemoryCheckpointStorage + from agent_framework_declarative._workflows import ToolApprovalResponse + + storage = InMemoryCheckpointStorage() + handler = StubMcpHandler(_ok()) + factory = WorkflowFactory(mcp_tool_handler=handler) + downstream = {**_action(tool_name="downstream"), "id": "downstream"} + definition = _yaml(_action(require_approval=True)) + definition["actions"].append(downstream) + workflow = factory.create_workflow_from_definition(definition) + agent = workflow.as_agent(name="mcp-agent") + + async def agent_turn(text: str, session: AgentSession | None = None) -> None: + if streaming: + await agent.run(text, session=session, stream=True, checkpoint_storage=storage).get_final_response() + else: + await agent.run(text, session=session, checkpoint_storage=storage) + + await agent_turn("first sessionless turn") + [sessionless_approval] = (await workflow._runner_context.get_pending_request_info_events()).values() + responses = {sessionless_approval.request_id: ToolApprovalResponse(approved=approved)} + if streaming: + await workflow.run(responses=responses, stream=True).get_final_response() + else: + await workflow.run(responses=responses) + assert [invocation.workflow_session_id for invocation in handler.invocations] == [ + sessionless_approval.data.workflow_session_id + ] * (2 if approved else 1) + + await agent_turn("explicit session turn", AgentSession()) + [explicit_approval] = (await workflow._runner_context.get_pending_request_info_events()).values() + assert sessionless_approval.data.workflow_session_id + assert explicit_approval.data.workflow_session_id + assert sessionless_approval.data.workflow_session_id != explicit_approval.data.workflow_session_id + + if restore: + checkpoint_id = workflow.get_last_checkpoint_id() + assert checkpoint_id is not None + workflow = factory.create_workflow_from_definition(definition) + await workflow.run(checkpoint_id=checkpoint_id, checkpoint_storage=storage) + agent = workflow.as_agent(name="mcp-agent") + + responses = {explicit_approval.request_id: ToolApprovalResponse(approved=approved)} + if streaming: + await workflow.run(responses=responses, stream=True).get_final_response() + else: + await workflow.run(responses=responses) + assert [invocation.workflow_session_id for invocation in handler.invocations] == [ + scope + for scope in (sessionless_approval.data.workflow_session_id, explicit_approval.data.workflow_session_id) + for _ in range(2 if approved else 1) + ] + + await agent_turn("second sessionless turn") + [new_approval] = (await workflow._runner_context.get_pending_request_info_events()).values() + assert new_approval.data.workflow_session_id == sessionless_approval.data.workflow_session_id + responses = {new_approval.request_id: ToolApprovalResponse(approved=approved)} + if streaming: + await workflow.run(responses=responses, stream=True).get_final_response() + else: + await workflow.run(responses=responses) + assert [invocation.workflow_session_id for invocation in handler.invocations] == [ + scope + for scope in ( + sessionless_approval.data.workflow_session_id, + explicit_approval.data.workflow_session_id, + sessionless_approval.data.workflow_session_id, + ) + for _ in range(2 if approved else 1) + ] + assert not await workflow._runner_context.get_pending_request_info_events() + + @pytest.mark.asyncio + async def test_as_agent_sessions_isolate_and_reuse_workflow_session_ids(self) -> None: + handler = StubMcpHandler(_ok()) + factory = WorkflowFactory(mcp_tool_handler=handler) + workflow = factory.create_workflow_from_definition(_yaml(_action())) + agent = workflow.as_agent(name="mcp-agent") + session_a = AgentSession() + session_b = AgentSession() + + await agent.run("session A first turn", session=session_a) + await agent.run("session B first turn", session=session_b) + await agent.run("session A second turn", session=session_a) + + assert len(handler.invocations) == 3 + session_a_id = handler.invocations[0].workflow_session_id + session_b_id = handler.invocations[1].workflow_session_id + assert session_a_id + assert session_b_id + assert session_a_id != session_b_id + assert handler.invocations[2].workflow_session_id == session_a_id + + @pytest.mark.parametrize("streaming", [False, True]) + @pytest.mark.parametrize("run_kind", ["direct", "sessionless_agent", "explicit_agent"]) + @pytest.mark.parametrize("depth", [1, 2]) + async def test_parent_entry_replay_preserves_child_workflow_session_ids( + self, streaming: bool, run_kind: str, depth: int + ) -> None: + from agent_framework import ( + Executor, + InMemoryCheckpointStorage, + WorkflowBuilder, + WorkflowContext, + WorkflowExecutor, + ) + from agent_framework import handler as workflow_handler + + class Start(Executor): + @workflow_handler + async def forward(self, message: list[Message], ctx: WorkflowContext[list[Message]]) -> None: + await ctx.send_message(message) + + mcp_handler = StubMcpHandler(_ok()) + factory = WorkflowFactory(mcp_tool_handler=mcp_handler) + storage = InMemoryCheckpointStorage() + + def build(): + wrappers = [] + for index in range(2): + child = factory.create_workflow_from_definition(_yaml(_action(tool_name=f"tool_{index}"))) + wrapper = WorkflowExecutor(child, id=f"child_{index}_0") + for level in range(1, depth): + child = WorkflowBuilder(start_executor=wrapper).build() + wrapper = WorkflowExecutor(child, id=f"child_{index}_{level}") + wrappers.append(wrapper) + start = Start(id="start") + return WorkflowBuilder(name="parent", start_executor=start).add_fan_out_edges(start, wrappers).build() + + parent = build() + messages = [Message(role="user", contents=[Content.from_text("First turn")])] + if run_kind == "direct": + await parent.run(messages, checkpoint_storage=storage) + else: + await parent.as_agent().run( + messages, + session=AgentSession() if run_kind == "explicit_agent" else None, + checkpoint_storage=storage, + ) + original_scopes = { + invocation.tool_name: invocation.workflow_session_id for invocation in mcp_handler.invocations + } + checkpoints = await storage.list_checkpoints(workflow_name=parent.name) + [entry] = [checkpoint for checkpoint in checkpoints if checkpoint.iteration_count == 0] + parent = build() + if streaming: + await parent.run( + checkpoint_id=entry.checkpoint_id, checkpoint_storage=storage, stream=True + ).get_final_response() + else: + await parent.run(checkpoint_id=entry.checkpoint_id, checkpoint_storage=storage) + + assert len(mcp_handler.invocations) == 4 + assert len(set(original_scopes.values())) == 2 + for invocation in mcp_handler.invocations[2:]: + assert invocation.workflow_session_id == original_scopes[invocation.tool_name] + + @pytest.mark.parametrize("streaming", [False, True]) + @pytest.mark.parametrize("depth", [1, 2]) + @pytest.mark.parametrize("restore", [False, True]) + async def test_resumed_child_restores_parent_scope_before_next_binding( + self, tmp_path: Path, streaming: bool, depth: int, restore: bool + ) -> None: + from agent_framework import Workflow, WorkflowBuilder, WorkflowExecutor + from agent_framework._workflows._const import WORKFLOW_SESSION_ID_KEY + + mcp_handler = StubMcpHandler(_ok()) + factory = WorkflowFactory(mcp_tool_handler=mcp_handler) + storage = _scope_checkpoint_storage(tmp_path) + + def build() -> Workflow: + child = factory.create_workflow_from_definition({ + "name": "first_child", + "actions": [ + _action(tool_name="before"), + _scope_pause_action("Question"), + {"kind": "SendActivity", "activity": {"text": "Completed"}}, + ], + }) + for level in range(1, depth): + nested = WorkflowExecutor(child, id=f"nested_{level}", propagate_request=True, allow_direct_output=True) + child = WorkflowBuilder(start_executor=nested).build() + first = WorkflowExecutor(child, id="first", propagate_request=True) + second = WorkflowExecutor( + factory.create_workflow_from_definition(_yaml(_action(tool_name="after"))), id="second" + ) + return WorkflowBuilder(start_executor=first, checkpoint_storage=storage).add_edge(first, second).build() + + parent = build() + await parent.run({}) + [request_a] = (await parent._runner_context.get_pending_request_info_events()).values() + scope_a = parent._runner.state.get(WORKFLOW_SESSION_ID_KEY) + await parent.run({}) + requests = await parent._runner_context.get_pending_request_info_events() + [request_b] = [request for request in requests.values() if request.request_id != request_a.request_id] + scope_b = parent._runner.state.get(WORKFLOW_SESSION_ID_KEY) + assert scope_a != scope_b + if restore: + checkpoint_id = parent.get_last_checkpoint_id() + assert checkpoint_id is not None + parent = build() + await parent.run(checkpoint_id=checkpoint_id) + + before_rejection = parent._runner.state.get(WORKFLOW_SESSION_ID_KEY) + with pytest.raises(ValueError): + await parent.run(responses={request_a.request_id: 42}) + assert parent._runner.state.get(WORKFLOW_SESSION_ID_KEY) == before_rejection + assert set(await parent._runner_context.get_pending_request_info_events()) == set(requests) + + for request, scope in ((request_b, scope_b), (request_a, scope_a)): + responses = {request.request_id: _scope_pause_response("Question")} + if streaming: + await parent.run(responses=responses, stream=True).get_final_response() + else: + await parent.run(responses=responses) + assert parent._runner.state.get(WORKFLOW_SESSION_ID_KEY) == scope + + scopes = [ + invocation.workflow_session_id for invocation in mcp_handler.invocations if invocation.tool_name == "after" + ] + assert len(scopes) == 4 + assert scopes[0] != scopes[1] + assert scopes[2:] == [scopes[1], scopes[0]] + assert not await parent._runner_context.get_pending_request_info_events() + + @pytest.mark.parametrize("message_list", [False, True]) + @pytest.mark.parametrize("streaming", [False, True]) + @pytest.mark.parametrize("explicit_session", [False, True]) + @pytest.mark.parametrize("depth", [1, 2]) + async def test_nested_agent_turns_reuse_isolated_workflow_session_ids( + self, streaming: bool, explicit_session: bool, depth: int, message_list: bool + ) -> None: + from agent_framework import Executor, WorkflowBuilder, WorkflowContext, WorkflowExecutor, handler + + class Start(Executor): + @handler + async def forward(self, message: list[Message], ctx: WorkflowContext[list[Message]]) -> None: + await ctx.send_message(message) + + mcp_handler = StubMcpHandler(_ok()) + factory = WorkflowFactory(mcp_tool_handler=mcp_handler) + wrappers: list[WorkflowExecutor] = [] + for index in range(2): + child = factory.create_workflow_from_definition(_yaml(_action(tool_name=f"tool_{index}"))) + wrapper = WorkflowExecutor(child, id=f"child_{index}_0") + for level in range(1, depth): + child = WorkflowBuilder(start_executor=wrapper).build() + wrapper = WorkflowExecutor(child, id=f"child_{index}_{level}") + wrappers.append(wrapper) + start = Start(id="start") + parent = WorkflowBuilder(start_executor=start).add_fan_out_edges(start, wrappers).build() + agent = parent.as_agent() + session_a = AgentSession() if explicit_session else None + message = [Message(role="user", contents=[Content.from_text("Agent turn")])] if message_list else "Agent turn" + for session in (session_a, AgentSession(), session_a): + if streaming: + await agent.run(message, session=session, stream=True).get_final_response() + else: + await agent.run(message, session=session) + for _ in range(2): + inputs = [Message(role="user", contents=[Content.from_text("Direct turn")])] + if streaming: + await parent.run(inputs, stream=True).get_final_response() + else: + await parent.run(inputs) + + scopes_by_tool = [ + [invocation.workflow_session_id for invocation in mcp_handler.invocations if invocation.tool_name == tool] + for tool in ("tool_0", "tool_1") + ] + for scopes in scopes_by_tool: + assert len(scopes) == 5 + assert all(scopes) + assert scopes[2] == scopes[0] + assert len({scopes[0], scopes[1], scopes[3], scopes[4]}) == 4 + assert set(scopes_by_tool[0]).isdisjoint(scopes_by_tool[1]) + + @pytest.mark.parametrize("message_list", [False, True]) + @pytest.mark.parametrize("streaming", [False, True]) + @pytest.mark.parametrize("explicit_session", [False, True]) + async def test_direct_fresh_runs_after_agent_turn_receive_separate_session_ids( + self, message_list: bool, streaming: bool, explicit_session: bool + ) -> None: + handler = StubMcpHandler(_ok()) + factory = WorkflowFactory(mcp_tool_handler=handler) + workflow = factory.create_workflow_from_definition(_yaml(_action())) + agent = workflow.as_agent(name="mcp-agent") + + await agent.run("agent turn", session=AgentSession() if explicit_session else None) + for text in ("first direct run", "second direct run"): + inputs: list[Message] | dict[str, Any] = ( + [Message(role="user", contents=[Content.from_text(text)])] if message_list else {} + ) + if streaming: + await workflow.run(inputs, stream=True).get_final_response() + else: + await workflow.run(inputs) + + session_ids = [invocation.workflow_session_id for invocation in handler.invocations] + assert len(session_ids) == 3 + assert all(session_ids) + assert len(set(session_ids)) == 3 + + @pytest.mark.asyncio + async def test_continuation_reuses_workflow_session_id(self) -> None: + from agent_framework_declarative._workflows import ToolApprovalResponse + from agent_framework_declarative._workflows._mcp_handler import get_or_create_workflow_session_id + + handler = StubMcpHandler(_ok()) + factory = WorkflowFactory(mcp_tool_handler=handler) + workflow = factory.create_workflow_from_definition(_yaml(_action(require_approval=True))) + + paused = await workflow.run({}) + [approval] = paused.get_request_info_events() + workflow_session_id = get_or_create_workflow_session_id(workflow._runner.state) # pyright: ignore[reportPrivateUsage] + await workflow.run(responses={approval.request_id: ToolApprovalResponse(approved=True)}) + + assert handler.last_invocation is not None + assert handler.last_invocation.workflow_session_id == workflow_session_id + + @pytest.mark.parametrize("restore", [False, True]) + @pytest.mark.parametrize("agent_sessions", [False, True]) + @pytest.mark.parametrize("streaming", [False, True]) + @pytest.mark.parametrize("approved", [False, True]) + @pytest.mark.parametrize("second_approved", [False, True]) + async def test_pending_approval_uses_originating_workflow_session_id( + self, restore: bool, agent_sessions: bool, streaming: bool, approved: bool, second_approved: bool + ) -> None: + from agent_framework import InMemoryCheckpointStorage + from agent_framework_declarative._workflows import ToolApprovalResponse + + storage = InMemoryCheckpointStorage() + handler = StubMcpHandler(_ok()) + factory = WorkflowFactory(mcp_tool_handler=handler) + downstream = _action(tool_name="downstream") + downstream["id"] = "downstream" + definition = {"name": "mcp_test", "actions": [_action(require_approval=True), downstream]} + workflow = factory.create_workflow_from_definition(definition) + agent = workflow.as_agent(name="mcp-agent") + + if agent_sessions: + await agent.run("first turn", session=AgentSession(), checkpoint_storage=storage) + else: + await workflow.run({}, checkpoint_storage=storage) + [first_approval] = (await workflow._runner_context.get_pending_request_info_events()).values() + first_session_id = first_approval.data.workflow_session_id + + await workflow.run({}, checkpoint_storage=storage) + [second_approval] = [ + event + for event in (await workflow._runner_context.get_pending_request_info_events()).values() + if event.request_id != first_approval.request_id + ] + + assert first_session_id + assert second_approval.data.workflow_session_id + assert first_session_id != second_approval.data.workflow_session_id + + checkpoint_id = None + if restore: + checkpoint_id = workflow.get_last_checkpoint_id() + assert checkpoint_id is not None + workflow = factory.create_workflow_from_definition(definition) + await workflow.run(checkpoint_id=checkpoint_id, checkpoint_storage=storage) + + decisions = {first_approval.request_id: approved, second_approval.request_id: second_approved} + batch = {request_id: ToolApprovalResponse(approved=decision) for request_id, decision in decisions.items()} + with pytest.raises(DeclarativeWorkflowError, match="different workflow sessions"): + if streaming: + await workflow.run( + responses=batch, + stream=True, + checkpoint_id=checkpoint_id, + checkpoint_storage=storage, + ).get_final_response() + else: + await workflow.run(responses=batch, checkpoint_id=checkpoint_id, checkpoint_storage=storage) + assert handler.call_count == 0 + assert set(await workflow._runner_context.get_pending_request_info_events()) == set(batch) + + for approval in (first_approval, second_approval): + responses = {approval.request_id: ToolApprovalResponse(approved=decisions[approval.request_id])} + if streaming: + await workflow.run(responses=responses, stream=True).get_final_response() + else: + await workflow.run(responses=responses) + + expected_tools = [ + tool + for decision in (approved, second_approved) + for tool in (["search", "downstream"] if decision else ["downstream"]) + ] + assert [invocation.tool_name for invocation in handler.invocations] == expected_tools + expected_sessions = [ + approval.data.workflow_session_id + for approval in (first_approval, second_approval) + for _ in range(2 if decisions[approval.request_id] else 1) + ] + assert [invocation.workflow_session_id for invocation in handler.invocations] == expected_sessions + + @pytest.mark.parametrize("streaming", [False, True]) + @pytest.mark.parametrize("approved", [False, True]) + async def test_same_scope_approval_batch_is_supported(self, streaming: bool, approved: bool) -> None: + from agent_framework import WorkflowBuilder + from agent_framework_declarative._workflows import InvokeMcpToolActionExecutor, ToolApprovalResponse + from agent_framework_declarative._workflows._declarative_builder import _validate_response_batch + from agent_framework_declarative._workflows._executors_control_flow import JoinExecutor + + handler = StubMcpHandler(_ok()) + entry = JoinExecutor({"kind": "Entry"}, id="entry") + actions = [ + InvokeMcpToolActionExecutor({**_action(require_approval=True), "id": name}, mcp_tool_handler=handler) + for name in ("first", "second") + ] + workflow = WorkflowBuilder(start_executor=entry).add_fan_out_edges(entry, actions).build() + workflow._response_batch_validator = _validate_response_batch + paused = await workflow.run({}) + approvals = paused.get_request_info_events() + assert len(approvals) == 2 + scopes = {approval.data.workflow_session_id for approval in approvals} + assert len(scopes) == 1 + assert all(scopes) + + responses = {approval.request_id: ToolApprovalResponse(approved=approved) for approval in approvals} + if streaming: + await workflow.run(responses=responses, stream=True).get_final_response() + else: + await workflow.run(responses=responses) + + assert handler.call_count == (2 if approved else 0) + assert all(invocation.workflow_session_id in scopes for invocation in handler.invocations) + assert not await workflow._runner_context.get_pending_request_info_events() @pytest.mark.asyncio async def test_arguments_evaluated_and_preserves_none(self) -> None: @@ -234,6 +1318,1013 @@ async def test_server_label_and_connection_name_forwarded(self) -> None: assert inv.connection_name == "azure-conn" +def _scope_pause_action(kind: str) -> dict[str, Any]: + actions: dict[str, dict[str, Any]] = { + "Question": {"kind": "Question", "question": "Continue?", "variable": "Local.answer"}, + "RequestExternalInput": { + "kind": "RequestExternalInput", + "prompt": "Continue?", + "variable": "Local.answer", + }, + "InvokeFunctionTool": {"kind": "InvokeFunctionTool", "functionName": "echo", "requireApproval": True}, + "InvokeAzureAgent": { + "kind": "InvokeAzureAgent", + "agent": "TestAgent", + "input": {"externalLoop": {"when": '=Local.userInput <> "stop"'}}, + }, + "InvokeMcpTool": _action(require_approval=True), + } + return actions[kind] + + +def _scope_pause_response(kind: str, user_input: str = "stop") -> Any: + from agent_framework_declarative._workflows import ( + AgentExternalInputResponse, + ExternalInputResponse, + ToolApprovalResponse, + ) + + if kind == "InvokeAzureAgent": + return AgentExternalInputResponse(user_input=user_input) + if kind in ("InvokeFunctionTool", "InvokeMcpTool"): + return ToolApprovalResponse(approved=True) + return ExternalInputResponse(user_input=user_input) + + +def _scope_checkpoint_storage(path: Path) -> Any: + from agent_framework import FileCheckpointStorage + + return FileCheckpointStorage( + path, + allowed_checkpoint_types=[ + "agent_framework_declarative._workflows._declarative_base:ActionComplete", + "agent_framework_declarative._workflows._declarative_base:ActionTrigger", + "agent_framework_declarative._workflows._executors_tools:ToolApprovalRequest", + "agent_framework_declarative._workflows._executors_tools:ToolApprovalResponse", + "agent_framework_declarative._workflows._executors_mcp:MCPToolApprovalRequest", + "agent_framework_declarative._workflows._executors_agents:AgentExternalInputRequest", + "agent_framework_declarative._workflows._executors_agents:AgentExternalInputResponse", + "agent_framework_declarative._workflows._executors_agents:ExternalLoopState", + "agent_framework_declarative._workflows._executors_external_input:ExternalInputRequest", + "agent_framework_declarative._workflows._executors_external_input:ExternalInputResponse", + ], + ) + + +class TestRequestScopeContinuations: + @pytest.mark.parametrize("pause_kind", ["InvokeFunctionTool", "InvokeAzureAgent"]) + @pytest.mark.parametrize("streaming", [False, True]) + @pytest.mark.parametrize("restore", [False, True]) + async def test_remaining_request_types_restore_originating_scope( + self, tmp_path: Path, pause_kind: str, streaming: bool, restore: bool + ) -> None: + storage = _scope_checkpoint_storage(tmp_path) + handler = StubMcpHandler(_ok()) + function = MagicMock(return_value="done") + agent = MagicMock(run=AsyncMock(return_value="Need input")) + factory = WorkflowFactory(mcp_tool_handler=handler, agents={"TestAgent": agent}).register_tool("echo", function) + definition = { + "name": "request_continuations", + "actions": [ + {**_action(tool_name="before"), "id": "before"}, + {"kind": "SetValue", "path": "Local.userInput", "value": ""}, + {**_scope_pause_action(pause_kind), "id": "pause"}, + {**_action(tool_name="after"), "id": "after"}, + ], + } + workflow = factory.create_workflow_from_definition(definition) + requests: list[Any] = [] + for _ in range(2): + await workflow.run({}, checkpoint_storage=storage) + [request] = [ + event + for event in (await workflow._runner_context.get_pending_request_info_events()).values() + if event.request_id not in {previous.request_id for previous in requests} + ] + requests.append(request) + scopes = [invocation.workflow_session_id for invocation in handler.invocations] + assert scopes[0] and scopes[1] and scopes[0] != scopes[1] + if restore: + checkpoint_id = workflow.get_last_checkpoint_id() + assert checkpoint_id is not None + workflow = factory.create_workflow_from_definition(definition) + await workflow.run(checkpoint_id=checkpoint_id, checkpoint_storage=storage) + assert set(await workflow._runner_context.get_pending_request_info_events()) == { + request.request_id for request in requests + } + + for request, scope in zip(requests, scopes, strict=True): + responses = {request.request_id: _scope_pause_response(pause_kind)} + if streaming: + await workflow.run(responses=responses, stream=True).get_final_response() + else: + await workflow.run(responses=responses) + assert handler.invocations[-1].tool_name == "after" + assert handler.invocations[-1].workflow_session_id == scope + assert handler.call_count == 4 + assert function.call_count == (2 if pause_kind == "InvokeFunctionTool" else 0) + assert not await workflow._runner_context.get_pending_request_info_events() + + @pytest.mark.parametrize("streaming", [False, True]) + async def test_repeated_agent_input_retains_its_own_loop_state(self, tmp_path: Path, streaming: bool) -> None: + from agent_framework_declarative._workflows._mcp_handler import get_or_create_workflow_session_id + + storage = _scope_checkpoint_storage(tmp_path) + handler = StubMcpHandler(_ok()) + agent = MagicMock(run=AsyncMock(return_value="Need input")) + factory = WorkflowFactory(mcp_tool_handler=handler, agents={"TestAgent": agent}) + definition = { + "name": "repeated_agent_input", + "actions": [ + {**_action(tool_name="before"), "id": "before"}, + {"kind": "SetValue", "path": "Local.userInput", "value": ""}, + {**_scope_pause_action("InvokeAzureAgent"), "id": "pause"}, + {**_action(tool_name="after"), "id": "after"}, + ], + } + workflow = factory.create_workflow_from_definition(definition) + await workflow.run({}) + [first] = (await workflow._runner_context.get_pending_request_info_events()).values() + await workflow.run({}) + [second] = [ + event + for event in (await workflow._runner_context.get_pending_request_info_events()).values() + if event.request_id != first.request_id + ] + first_scope, second_scope = (invocation.workflow_session_id for invocation in handler.invocations) + await workflow.run( + responses={first.request_id: _scope_pause_response("InvokeAzureAgent", "continue")}, + checkpoint_storage=storage, + ) + assert get_or_create_workflow_session_id(workflow._runner.state) == first_scope + [repeated] = [ + event + for event in (await workflow._runner_context.get_pending_request_info_events()).values() + if event.request_id != second.request_id + ] + checkpoint_id = workflow.get_last_checkpoint_id() + assert checkpoint_id is not None + workflow = factory.create_workflow_from_definition(definition) + await workflow.run(checkpoint_id=checkpoint_id, checkpoint_storage=storage) + + for request, scope in ((second, second_scope), (repeated, first_scope)): + responses = {request.request_id: _scope_pause_response("InvokeAzureAgent")} + if streaming: + await workflow.run(responses=responses, stream=True).get_final_response() + else: + await workflow.run(responses=responses) + assert handler.invocations[-1].tool_name == "after" + assert handler.invocations[-1].workflow_session_id == scope + assert handler.call_count == 4 + assert agent.run.await_count == 3 + assert not await workflow._runner_context.get_pending_request_info_events() + + @pytest.mark.parametrize("pause_kind", ["InvokeFunctionTool", "InvokeAzureAgent"]) + @pytest.mark.parametrize("streaming", [False, True]) + async def test_mixed_request_families_reject_different_scopes(self, pause_kind: str, streaming: bool) -> None: + handler = StubMcpHandler(_ok()) + function = MagicMock(return_value="done") + agent = MagicMock(run=AsyncMock(return_value="Need input")) + factory = WorkflowFactory(mcp_tool_handler=handler, agents={"TestAgent": agent}).register_tool("echo", function) + workflow = factory.create_workflow_from_definition({ + "name": "mixed_request_scopes", + "actions": [ + {"kind": "Question", "question": "Continue?", "variable": "Local.answer"}, + {"kind": "SetValue", "path": "Local.userInput", "value": ""}, + {**_scope_pause_action(pause_kind), "id": "pause"}, + _action(), + ], + }) + await workflow.run({}) + [first] = (await workflow._runner_context.get_pending_request_info_events()).values() + await workflow.run({}) + [second] = [ + event + for event in (await workflow._runner_context.get_pending_request_info_events()).values() + if event.request_id != first.request_id + ] + await workflow.run(responses={second.request_id: _scope_pause_response("Question")}) + [second] = [ + event + for event in (await workflow._runner_context.get_pending_request_info_events()).values() + if event.request_id != first.request_id + ] + responses = { + first.request_id: _scope_pause_response("Question"), + second.request_id: _scope_pause_response(pause_kind), + } + with pytest.raises(DeclarativeWorkflowError, match="different workflow sessions"): + if streaming: + await workflow.run(responses=responses, stream=True).get_final_response() + else: + await workflow.run(responses=responses) + assert handler.call_count == 0 + function.assert_not_called() + assert agent.run.await_count == (1 if pause_kind == "InvokeAzureAgent" else 0) + assert set(await workflow._runner_context.get_pending_request_info_events()) == set(responses) + + async def test_legacy_mcp_origin_is_rejected_before_consumption(self) -> None: + handler = StubMcpHandler(_ok()) + factory = WorkflowFactory(mcp_tool_handler=handler) + workflow = factory.create_workflow_from_definition(_yaml(_action(require_approval=True))) + await workflow.run({}) + [legacy] = (await workflow._runner_context.get_pending_request_info_events()).values() + legacy.data.workflow_session_id = None + await workflow.run({}) + + with pytest.raises(DeclarativeWorkflowError, match="migrate_legacy_request_scopes"): + await workflow.run(responses={legacy.request_id: _scope_pause_response("InvokeMcpTool")}) + + assert handler.call_count == 0 + assert legacy.request_id in await workflow._runner_context.get_pending_request_info_events() + + @pytest.mark.parametrize( + "pause_kind", ["Question", "RequestExternalInput", "InvokeFunctionTool", "InvokeAzureAgent", "InvokeMcpTool"] + ) + async def test_opt_in_legacy_migration_survives_file_checkpoints(self, tmp_path: Path, pause_kind: str) -> None: + storage = _scope_checkpoint_storage(tmp_path) + handler = StubMcpHandler(_ok()) + agent = MagicMock(run=AsyncMock(return_value="Need input")) + function = MagicMock(return_value="done") + factory = WorkflowFactory(mcp_tool_handler=handler, agents={"TestAgent": agent}).register_tool("echo", function) + definition = { + "name": "legacy_request_migration", + "actions": [ + {**_action(tool_name="before"), "id": "before"}, + {"kind": "SetValue", "path": "Local.userInput", "value": ""}, + {**_scope_pause_action(pause_kind), "id": "pause"}, + {**_action(tool_name="after"), "id": "after"}, + {**_scope_pause_action("Question"), "id": "followup"}, + {**_action(tool_name="done"), "id": "done"}, + ], + } + workflow = factory.create_workflow_from_definition(definition) + await workflow.run({}, checkpoint_storage=storage) + await workflow.run({}, checkpoint_storage=storage) + original_scopes = {invocation.workflow_session_id for invocation in handler.invocations} + checkpoint_id = workflow.get_last_checkpoint_id() + assert checkpoint_id is not None + checkpoint = await storage.load(checkpoint_id) + assert len(checkpoint.pending_request_info_events) == 2 + for event in checkpoint.pending_request_info_events.values(): + event.data.workflow_session_id = None + if pause_kind == "InvokeAzureAgent": + from agent_framework_declarative._workflows._executors_agents import EXTERNAL_LOOP_STATE_KEY + + checkpoint.state[EXTERNAL_LOOP_STATE_KEY] = event.data.loop_state + del event.data.loop_state + await storage.save(checkpoint) + first_id, second_id = checkpoint.pending_request_info_events + + workflow = factory.create_workflow_from_definition(definition) + await workflow.run(checkpoint_id=checkpoint_id, checkpoint_storage=storage) + with pytest.raises(DeclarativeWorkflowError, match="migrate_legacy_request_scopes"): + await workflow.run(responses={first_id: _scope_pause_response(pause_kind)}) + assert handler.call_count == 2 + assert set(await workflow._runner_context.get_pending_request_info_events()) == {first_id, second_id} + + migration_factory = WorkflowFactory( + mcp_tool_handler=handler, agents={"TestAgent": agent}, migrate_legacy_request_scopes=True + ).register_tool("echo", function) + workflow = migration_factory.create_workflow_from_definition(definition) + await workflow.run(checkpoint_id=checkpoint_id, checkpoint_storage=storage) + with pytest.raises(DeclarativeWorkflowError, match="different workflow sessions"): + await workflow.run( + responses={ + first_id: _scope_pause_response(pause_kind), + second_id: _scope_pause_response(pause_kind), + } + ) + pending = await workflow._runner_context.get_pending_request_info_events() + assert all(event.data.workflow_session_id is None for event in pending.values()) + assert handler.call_count == 2 + function.assert_not_called() + + migrated_scopes: list[str] = [] + for _ in range(2): + workflow = migration_factory.create_workflow_from_definition(definition) + await workflow.run(checkpoint_id=checkpoint_id, checkpoint_storage=storage) + before = handler.call_count + await workflow.run(responses={first_id: _scope_pause_response(pause_kind)}, checkpoint_storage=storage) + scope = handler.invocations[-1].workflow_session_id + assert scope and scope not in original_scopes + assert all(invocation.workflow_session_id == scope for invocation in handler.invocations[before:]) + migrated_scopes.append(scope) + assert migrated_scopes[0] == migrated_scopes[1] + + migrated_checkpoint_id = workflow.get_last_checkpoint_id() + assert migrated_checkpoint_id is not None + workflow = factory.create_workflow_from_definition(definition) + await workflow.run(checkpoint_id=migrated_checkpoint_id, checkpoint_storage=storage) + pending = await workflow._runner_context.get_pending_request_info_events() + [followup] = [event for event in pending.values() if event.request_id != second_id] + assert followup.data.workflow_session_id == migrated_scopes[0] + await workflow.run(responses={followup.request_id: _scope_pause_response("Question")}) + assert handler.invocations[-1].tool_name == "done" + assert handler.invocations[-1].workflow_session_id == migrated_scopes[0] + + migrated_checkpoint = await storage.load(migrated_checkpoint_id) + workflow = migration_factory.create_workflow_from_definition(definition) + await workflow.run(checkpoint_id=migrated_checkpoint.checkpoint_id, checkpoint_storage=storage) + await workflow.run(responses={second_id: _scope_pause_response(pause_kind)}) + assert handler.invocations[-1].tool_name == "after" + assert handler.invocations[-1].workflow_session_id not in original_scopes | set(migrated_scopes) + + @pytest.mark.parametrize("invalid_scope", ["", 42]) + @pytest.mark.parametrize("migrate", [False, True]) + async def test_migration_does_not_replace_invalid_scopes(self, invalid_scope: Any, migrate: bool) -> None: + handler = StubMcpHandler(_ok()) + factory = WorkflowFactory(mcp_tool_handler=handler, migrate_legacy_request_scopes=migrate) + workflow = factory.create_workflow_from_definition(_yaml(_action(require_approval=True))) + await workflow.run({}) + [request] = (await workflow._runner_context.get_pending_request_info_events()).values() + request.data.workflow_session_id = invalid_scope + + with pytest.raises(DeclarativeWorkflowError, match="Invalid originating workflow session"): + await workflow.run(responses={request.request_id: _scope_pause_response("InvokeMcpTool")}) + + assert handler.call_count == 0 + assert set(await workflow._runner_context.get_pending_request_info_events()) == {request.request_id} + + @pytest.mark.parametrize( + "pause_kind", ["Question", "RequestExternalInput", "InvokeFunctionTool", "InvokeAzureAgent", "InvokeMcpTool"] + ) + @pytest.mark.parametrize( + ("restore_kind", "streaming"), + [("memory", False), ("memory", True), ("file", False), ("file", True), ("embedded", False)], + ) + async def test_legacy_queued_response_migration_precedes_replay( + self, tmp_path: Path, pause_kind: str, restore_kind: str, streaming: bool + ) -> None: + from agent_framework import InMemoryCheckpointStorage + from agent_framework._workflows._runner_context import MessageType + + storage = _scope_checkpoint_storage(tmp_path) if restore_kind == "file" else InMemoryCheckpointStorage() + handler = StubMcpHandler(_ok()) + agent = MagicMock(run=AsyncMock(return_value="Need input")) + function = MagicMock(return_value="done") + factory = WorkflowFactory( + checkpoint_storage=storage, mcp_tool_handler=handler, agents={"TestAgent": agent} + ).register_tool("echo", function) + definition = { + "name": "queued_legacy_response", + "actions": [ + {**_action(tool_name="before"), "id": "before"}, + {"kind": "SetValue", "path": "Local.userInput", "value": ""}, + {**_scope_pause_action(pause_kind), "id": "pause"}, + {**_action(tool_name="after"), "id": "after"}, + {**_scope_pause_action("Question"), "id": "followup"}, + ], + } + workflow = factory.create_workflow_from_definition(definition) + await workflow.run({}) + [request] = (await workflow._runner_context.get_pending_request_info_events()).values() + await workflow.run(responses={request.request_id: _scope_pause_response(pause_kind)}) + checkpoints = await storage.list_checkpoints(workflow_name=workflow.name) + [checkpoint] = [ + item + for item in checkpoints + if any(message.type == MessageType.RESPONSE for batch in item.messages.values() for message in batch) + ] + [queued_response] = [message for batch in checkpoint.messages.values() for message in batch] + original = queued_response.original_request_info_event + assert original is not None + assert original.request_id == request.request_id + assert not checkpoint.pending_request_info_events + original_scope = original.data.workflow_session_id + original.data.workflow_session_id = None + await storage.save(checkpoint) + before = handler.call_count + function.reset_mock() + + factory = WorkflowFactory( + checkpoint_storage=storage, + mcp_tool_handler=handler, + agents={"TestAgent": agent}, + migrate_legacy_request_scopes=True, + ).register_tool("echo", function) + workflow = factory.create_workflow_from_definition(definition) + if restore_kind == "embedded": + await workflow._runner.restore_checkpoint(await storage.load(checkpoint.checkpoint_id)) + _ = [event async for event in workflow._runner.run_until_convergence()] + elif streaming: + await workflow.run(checkpoint_id=checkpoint.checkpoint_id, stream=True).get_final_response() + else: + await workflow.run(checkpoint_id=checkpoint.checkpoint_id) + + replayed = handler.invocations[before:] + assert replayed + assert replayed[-1].tool_name == "after" + scope = replayed[-1].workflow_session_id + assert scope == f"legacy-request:{request.request_id}" + assert scope != original_scope + assert all(invocation.workflow_session_id == scope for invocation in replayed) + assert function.call_count == (1 if pause_kind == "InvokeFunctionTool" else 0) + [followup] = (await workflow._runner_context.get_pending_request_info_events()).values() + assert followup.data.workflow_session_id == scope + saved_id = workflow.get_last_checkpoint_id() + assert saved_id is not None + saved = await storage.load(saved_id) + assert saved.pending_request_info_events[followup.request_id].data.workflow_session_id == scope + + @pytest.mark.parametrize("restore_kind", ["file", "embedded"]) + @pytest.mark.parametrize("legacy", [False, True]) + async def test_queued_response_batch_is_validated_before_restoring_state( + self, tmp_path: Path, restore_kind: str, legacy: bool + ) -> None: + from agent_framework import InMemoryCheckpointStorage + from agent_framework.exceptions import WorkflowCheckpointException + + storage = _scope_checkpoint_storage(tmp_path) if restore_kind == "file" else InMemoryCheckpointStorage() + handler = StubMcpHandler(_ok()) + factory = WorkflowFactory( + checkpoint_storage=storage, mcp_tool_handler=handler, migrate_legacy_request_scopes=True + ) + workflow = factory.create_workflow_from_definition(_yaml(_action(require_approval=True))) + await workflow.run({}) + await workflow.run({}) + requests = await workflow._runner_context.get_pending_request_info_events() + assert len(requests) == 2 + for request in requests.values(): + if legacy: + request.data.workflow_session_id = None + await workflow._runner_context.send_request_info_response( + request.request_id, _scope_pause_response("InvokeMcpTool") + ) + await workflow._runner.create_checkpoint_if_enabled() + checkpoint_id = workflow.get_last_checkpoint_id() + assert checkpoint_id is not None + checkpoint = await storage.load(checkpoint_id) + assert not checkpoint.pending_request_info_events + + workflow = factory.create_workflow_from_definition(_yaml(_action(require_approval=True))) + workflow._runner.state.set("preexisting", "keep") + workflow._runner.state.commit() + with pytest.raises(WorkflowCheckpointException) as exc: + if restore_kind == "embedded": + await workflow._runner.restore_checkpoint(checkpoint) + else: + await workflow.run(checkpoint_id=checkpoint_id) + + assert isinstance(exc.value.__cause__, DeclarativeWorkflowError) + assert "different workflow sessions" in str(exc.value.__cause__) + assert workflow._runner.state.get("preexisting") == "keep" + assert handler.call_count == 0 + if legacy: + for batch in checkpoint.messages.values(): + for message in batch: + original = message.original_request_info_event + assert original is not None + assert original.data.workflow_session_id is None + + @pytest.mark.parametrize("restore_kind", ["file", "embedded"]) + async def test_queued_and_new_responses_share_batch_validation(self, tmp_path: Path, restore_kind: str) -> None: + from agent_framework import InMemoryCheckpointStorage + + storage = _scope_checkpoint_storage(tmp_path) if restore_kind == "file" else InMemoryCheckpointStorage() + handler = StubMcpHandler(_ok()) + factory = WorkflowFactory(checkpoint_storage=storage, mcp_tool_handler=handler) + definition = _yaml(_action(require_approval=True)) + workflow = factory.create_workflow_from_definition(definition) + await workflow.run({}) + await workflow.run({}) + first_id, second_id = (await workflow._runner_context.get_pending_request_info_events()).keys() + await workflow._runner_context.send_request_info_response(first_id, _scope_pause_response("InvokeMcpTool")) + await workflow._runner.create_checkpoint_if_enabled() + checkpoint_id = workflow.get_last_checkpoint_id() + assert checkpoint_id is not None + + workflow = factory.create_workflow_from_definition(definition) + with pytest.raises(DeclarativeWorkflowError, match="different workflow sessions"): + if restore_kind == "embedded": + await workflow._runner.restore_checkpoint(await storage.load(checkpoint_id)) + await workflow.run(responses={second_id: _scope_pause_response("InvokeMcpTool")}) + else: + await workflow.run( + checkpoint_id=checkpoint_id, responses={second_id: _scope_pause_response("InvokeMcpTool")} + ) + + assert handler.call_count == 0 + assert set(await workflow._runner_context.get_pending_request_info_events()) == {second_id} + + @pytest.mark.parametrize("restore_kind", ["file", "embedded"]) + async def test_legacy_queued_response_requires_opt_in_before_restore( + self, tmp_path: Path, restore_kind: str + ) -> None: + from agent_framework import InMemoryCheckpointStorage + from agent_framework.exceptions import WorkflowCheckpointException + + storage = _scope_checkpoint_storage(tmp_path) if restore_kind == "file" else InMemoryCheckpointStorage() + handler = StubMcpHandler(_ok()) + factory = WorkflowFactory(checkpoint_storage=storage, mcp_tool_handler=handler) + definition = _yaml(_action(require_approval=True)) + workflow = factory.create_workflow_from_definition(definition) + await workflow.run({}) + [request] = (await workflow._runner_context.get_pending_request_info_events()).values() + request.data.workflow_session_id = None + await workflow._runner_context.send_request_info_response( + request.request_id, _scope_pause_response("InvokeMcpTool") + ) + await workflow._runner.create_checkpoint_if_enabled() + checkpoint_id = workflow.get_last_checkpoint_id() + assert checkpoint_id is not None + checkpoint = await storage.load(checkpoint_id) + + workflow = factory.create_workflow_from_definition(definition) + workflow._runner.state.set("preexisting", "keep") + workflow._runner.state.commit() + with pytest.raises(WorkflowCheckpointException) as exc: + if restore_kind == "embedded": + await workflow._runner.restore_checkpoint(checkpoint) + else: + await workflow.run(checkpoint_id=checkpoint_id) + + assert isinstance(exc.value.__cause__, DeclarativeWorkflowError) + assert "migrate_legacy_request_scopes" in str(exc.value.__cause__) + assert workflow._runner.state.get("preexisting") == "keep" + assert handler.call_count == 0 + + @pytest.mark.parametrize("restore_kind", ["file", "embedded"]) + @pytest.mark.parametrize("depth", [1, 2]) + @pytest.mark.parametrize("mode", ["missing", "invalid", "migrate", "valid"]) + async def test_parent_queued_child_response_preflight_preserves_state( + self, tmp_path: Path, restore_kind: str, depth: int, mode: str + ) -> None: + from copy import deepcopy + + from agent_framework import WorkflowBuilder, WorkflowExecutor + from agent_framework._workflows._checkpoint_encoding import encode_checkpoint_value + from agent_framework._workflows._const import EXECUTOR_STATE_KEY + from agent_framework.exceptions import WorkflowCheckpointException + + storage = _scope_checkpoint_storage(tmp_path) + mcp_handler = StubMcpHandler(_ok()) + child = WorkflowFactory( + mcp_tool_handler=mcp_handler, migrate_legacy_request_scopes=mode == "migrate" + ).create_workflow_from_definition(_yaml(_action(require_approval=True))) + workflows = [child] + parent = child + for level in range(depth): + wrapper = WorkflowExecutor(parent, id=f"nested_{level}", propagate_request=True) + parent = WorkflowBuilder(start_executor=wrapper, checkpoint_storage=storage).build() + workflows.append(parent) + await parent.run({}) + [request] = (await parent._runner_context.get_pending_request_info_events()).values() + await parent._runner_context.send_request_info_response( + request.request_id, _scope_pause_response("InvokeMcpTool") + ) + checkpoint = deepcopy(await parent._runner.build_checkpoint()) + # The child's pending request, not the parent's copy, owns the response policy. + for batch in checkpoint.messages.values(): + for message in batch: + message.original_request_info_event = deepcopy(message.original_request_info_event) + child_checkpoint = checkpoint + for level in reversed(range(depth)): + child_checkpoint = child_checkpoint.state[EXECUTOR_STATE_KEY][f"nested_{level}"]["sub_workflow_checkpoint"] + original_scope = child_checkpoint.pending_request_info_events[request.request_id].data.workflow_session_id + if mode != "valid": + child_checkpoint.pending_request_info_events[request.request_id].data.workflow_session_id = ( + 42 if mode == "invalid" else None + ) + checkpoint_before = encode_checkpoint_value(checkpoint.to_dict()) + await storage.save(checkpoint) + + for workflow in workflows: + workflow._runner.state.set("restore_marker", "live") + workflow._runner.state.commit() + workflow._runner.state.set("uncommitted", "keep") + + async def snapshots() -> list[Any]: + result: list[Any] = [] + for workflow in workflows: + snapshot = await workflow._runner_context.build_checkpoint( + workflow.name, workflow.graph_signature_hash, workflow._runner.state, None, 0 + ) + result.append( + encode_checkpoint_value({ + "state": snapshot.state, + "messages": snapshot.messages, + "requests": snapshot.pending_request_info_events, + }) + ) + return result + + before = await snapshots() + + async def restore() -> None: + if restore_kind == "embedded": + await parent._runner.restore_checkpoint(checkpoint) + else: + await parent._runner.restore_from_checkpoint(checkpoint.checkpoint_id) + + if mode in {"missing", "invalid"}: + with pytest.raises(WorkflowCheckpointException): + await restore() + assert await snapshots() == before + for workflow in workflows: + assert workflow._runner.state.get("restore_marker") == "live" + assert workflow._runner.state.get("uncommitted") == "keep" + else: + await restore() + _ = [event async for event in parent._runner.run_until_convergence()] + assert mcp_handler.call_count == 1 + assert mcp_handler.invocations[0].workflow_session_id == ( + f"legacy-request:{request.request_id}" if mode == "migrate" else original_scope + ) + assert encode_checkpoint_value(checkpoint.to_dict()) == checkpoint_before + assert mcp_handler.call_count == (1 if mode in {"migrate", "valid"} else 0) + + @pytest.mark.parametrize("restore_kind", ["file", "embedded"]) + @pytest.mark.parametrize("depth", [1, 2]) + @pytest.mark.parametrize("targeted", [False, True]) + @pytest.mark.parametrize("mode", ["missing", "invalid", "migrate", "valid"]) + async def test_wrapped_child_response_validation_uses_restored_state( + self, tmp_path: Path, restore_kind: str, depth: int, targeted: bool, mode: str + ) -> None: + from copy import deepcopy + + from agent_framework import ( + Executor, + SubWorkflowRequestMessage, + SubWorkflowResponseMessage, + WorkflowBuilder, + WorkflowContext, + WorkflowExecutor, + handler, + ) + from agent_framework._workflows._checkpoint_encoding import encode_checkpoint_value + from agent_framework._workflows._const import EXECUTOR_STATE_KEY + from agent_framework._workflows._runner_context import WorkflowMessage + + class RestoreRecorder(Executor): + def __init__(self) -> None: + super().__init__(id="coordinator") + self.value = "snapshot" + self.restore_calls = 0 + + @handler + async def forward(self, message: dict[str, Any], ctx: WorkflowContext[dict[str, Any]]) -> None: + await ctx.send_message(message) + + async def on_checkpoint_save(self) -> dict[str, Any]: + return {"value": self.value} + + async def on_checkpoint_restore(self, state: dict[str, Any]) -> None: + self.restore_calls += 1 + self.value = state["value"] + + storage = _scope_checkpoint_storage(tmp_path) + mcp_handler = StubMcpHandler(_ok()) + leaves = [ + WorkflowFactory( + mcp_tool_handler=mcp_handler, migrate_legacy_request_scopes=index == 1 and mode == "migrate" + ).create_workflow_from_definition(_yaml(_action(require_approval=True))) + for index in range(2) + ] + workflows = list(leaves) + wrappers: list[WorkflowExecutor] = [] + for index, leaf in enumerate(leaves): + child = leaf + for level in range(1, depth): + child = WorkflowBuilder( + start_executor=WorkflowExecutor(child, id=f"nested_{index}_{level}", propagate_request=True) + ).build() + workflows.append(child) + wrappers.append(WorkflowExecutor(child, id=f"child_{index}")) + recorder = RestoreRecorder() + parent = ( + WorkflowBuilder(start_executor=recorder, checkpoint_storage=storage) + .add_fan_out_edges(recorder, wrappers) + .build() + ) + workflows.append(parent) + sibling_scope: str | None = None + for wrapper in wrappers: + await wrapper.workflow.run({}) + [request] = (await wrapper.workflow._runner_context.get_pending_request_info_events()).values() + if wrapper is wrappers[0]: + sibling_scope = request.data.workflow_session_id + reply = SubWorkflowRequestMessage(source_event=deepcopy(request), executor_id=wrapper.id).create_response( + _scope_pause_response("InvokeMcpTool") + ) + await parent._runner_context.send_message( + WorkflowMessage(data=reply, source_id=recorder.id, target_id=wrapper.id if targeted else None) + ) + for workflow in workflows: + workflow._runner.state.set("restore_marker", "snapshot") + workflow._runner.state.commit() + checkpoint = deepcopy(await parent._runner.build_checkpoint()) + child_checkpoint = checkpoint.state[EXECUTOR_STATE_KEY]["child_1"]["sub_workflow_checkpoint"] + for level in reversed(range(1, depth)): + child_checkpoint = child_checkpoint.state[EXECUTOR_STATE_KEY][f"nested_1_{level}"][ + "sub_workflow_checkpoint" + ] + [child_request] = child_checkpoint.pending_request_info_events.values() + original_scope = child_request.data.workflow_session_id + if mode != "valid": + child_request.data.workflow_session_id = 42 if mode == "invalid" else None + checkpoint_before = encode_checkpoint_value(checkpoint.to_dict()) + await storage.save(checkpoint) + + for workflow in workflows: + workflow._runner.state.set("restore_marker", "live") + workflow._runner.state.commit() + workflow._runner.state.set("uncommitted", "keep") + recorder.value = "live" + + async def restore() -> None: + if restore_kind == "embedded": + await parent._runner.restore_checkpoint(checkpoint) + else: + await parent._runner.restore_from_checkpoint(checkpoint.checkpoint_id) + + await restore() + if mode in {"missing", "invalid"}: + with pytest.raises(DeclarativeWorkflowError): + _ = [event async for event in parent._runner.run_until_convergence()] + pending = await leaves[1]._runner_context.get_pending_request_info_events() + assert set(pending) == {child_request.request_id} + assert pending[child_request.request_id].data.workflow_session_id == (42 if mode == "invalid" else None) + queued = await parent._runner._get_queued_messages() + assert [ + message.data.source_event.request_id + for messages in queued.values() + for message in messages + if isinstance(message.data, SubWorkflowResponseMessage) + ] == [child_request.request_id] + else: + _ = [event async for event in parent._runner.run_until_convergence()] + assert {invocation.workflow_session_id for invocation in mcp_handler.invocations} == { + sibling_scope, + f"legacy-request:{child_request.request_id}" if mode == "migrate" else original_scope, + } + assert all(workflow._runner.state.get("restore_marker") == "snapshot" for workflow in workflows) + assert all(not workflow._runner.state.has("uncommitted") for workflow in workflows) + assert recorder.value == "snapshot" + assert recorder.restore_calls == 1 + assert encode_checkpoint_value(checkpoint.to_dict()) == checkpoint_before + assert mcp_handler.call_count == (2 if mode in {"migrate", "valid"} else 1) + + @pytest.mark.parametrize("depth", [1, 2]) + @pytest.mark.parametrize("mode", ["missing", "invalid", "queued"]) + async def test_parent_child_response_rejection_preserves_pending_request(self, depth: int, mode: str) -> None: + from copy import deepcopy + + from agent_framework import WorkflowBuilder, WorkflowExecutor + + mcp_handler = StubMcpHandler(_ok()) + child = WorkflowFactory(mcp_tool_handler=mcp_handler).create_workflow_from_definition( + _yaml(_action(require_approval=True)) + ) + parent = child + for level in range(depth): + parent = WorkflowBuilder( + start_executor=WorkflowExecutor(parent, id=f"nested_{level}", propagate_request=True) + ).build() + await parent.run({}) + if mode == "queued": + await parent.run({}) + pending = await parent._runner_context.get_pending_request_info_events() + child_pending = await child._runner_context.get_pending_request_info_events() + request_ids = list(pending) + request_id = request_ids[-1] + original_scope = child_pending[request_id].data.workflow_session_id + if mode == "queued": + await child._runner_context.send_request_info_response( + request_ids[0], _scope_pause_response("InvokeMcpTool") + ) + else: + pending[request_id].data = deepcopy(pending[request_id].data) + child_pending[request_id].data.workflow_session_id = None if mode == "missing" else 42 + + with pytest.raises(DeclarativeWorkflowError): + await parent.run(responses={request_id: _scope_pause_response("InvokeMcpTool")}) + + assert set(await parent._runner_context.get_pending_request_info_events()) == set(request_ids) + assert request_id in await child._runner_context.get_pending_request_info_events() + assert mcp_handler.call_count == 0 + if mode != "queued": + child_pending[request_id].data.workflow_session_id = original_scope + await parent.run(responses={request_id: _scope_pause_response("InvokeMcpTool")}) + assert not await parent._runner_context.get_pending_request_info_events() + assert mcp_handler.call_count == 1 + + @pytest.mark.parametrize("depth", [1, 2]) + @pytest.mark.parametrize("queued", [False, True]) + @pytest.mark.parametrize("origin", [None, 42]) + async def test_sibling_response_rejection_preserves_unmigrated_origins( + self, depth: int, queued: bool, origin: Any + ) -> None: + from agent_framework import Executor, WorkflowBuilder, WorkflowContext, WorkflowExecutor, handler + + class Start(Executor): + @handler + async def forward(self, message: dict[str, Any], ctx: WorkflowContext[dict[str, Any]]) -> None: + await ctx.send_message(message) + + mcp_handler = StubMcpHandler(_ok()) + leaves = [ + WorkflowFactory( + mcp_tool_handler=mcp_handler, migrate_legacy_request_scopes=index == 0 + ).create_workflow_from_definition(_yaml(_action(require_approval=True))) + for index in range(2) + ] + wrappers: list[WorkflowExecutor] = [] + for index, leaf in enumerate(leaves): + wrapper = WorkflowExecutor(leaf, id=f"child_{index}_0", propagate_request=True) + for level in range(1, depth): + child = WorkflowBuilder(start_executor=wrapper).build() + wrapper = WorkflowExecutor(child, id=f"child_{index}_{level}", propagate_request=True) + wrappers.append(wrapper) + start = Start(id="start") + parent = WorkflowBuilder(start_executor=start).add_fan_out_edges(start, wrappers).build() + await parent.run({}) + requests = [ + next(iter((await leaf._runner_context.get_pending_request_info_events()).values())) for leaf in leaves + ] + first, second = requests + second_scope = second.data.workflow_session_id + first.data.workflow_session_id = None + second.data.workflow_session_id = origin + if queued: + await parent._runner_context.send_request_info_response( + first.request_id, _scope_pause_response("InvokeMcpTool") + ) + pending_ids = set(await parent._runner_context.get_pending_request_info_events()) + responses = {request_id: _scope_pause_response("InvokeMcpTool") for request_id in pending_ids} + + with pytest.raises(DeclarativeWorkflowError): + await parent.run(responses=responses) + + assert first.data.workflow_session_id is None + assert second.data.workflow_session_id == origin + assert set(await parent._runner_context.get_pending_request_info_events()) == pending_ids + for leaf, request in zip(leaves, requests, strict=True): + assert set(await leaf._runner_context.get_pending_request_info_events()) == {request.request_id} + assert mcp_handler.call_count == 0 + + second.data.workflow_session_id = second_scope + await parent.run(responses=responses) + assert not await parent._runner_context.get_pending_request_info_events() + assert first.data.workflow_session_id == f"legacy-request:{first.request_id}" + assert {invocation.workflow_session_id for invocation in mcp_handler.invocations} == { + f"legacy-request:{first.request_id}", + second_scope, + } + assert mcp_handler.call_count == 2 + + @pytest.mark.parametrize("restore_kind", ["file", "embedded"]) + @pytest.mark.parametrize("depth", [1, 2]) + @pytest.mark.parametrize("mode", ["missing", "mixed", "valid", "pending"]) + async def test_nested_checkpoint_preflight_preserves_live_state( + self, tmp_path: Path, restore_kind: str, depth: int, mode: str + ) -> None: + from copy import deepcopy + + from agent_framework import ( + Executor, + InProcRunnerContext, + Workflow, + WorkflowBuilder, + WorkflowCheckpoint, + WorkflowContext, + WorkflowExecutor, + handler, + ) + from agent_framework._workflows._checkpoint_encoding import encode_checkpoint_value + from agent_framework._workflows._runner_context import WorkflowMessage + from agent_framework.exceptions import WorkflowCheckpointException + + class RestoreRecorder(Executor): + def __init__(self, executor_id: str) -> None: + super().__init__(id=executor_id) + self.value = "snapshot" + self.restore_calls = 0 + + @handler + async def forward(self, message: str, ctx: WorkflowContext[str]) -> None: + await ctx.send_message(message) + + async def on_checkpoint_save(self) -> dict[str, Any]: + return {"value": self.value} + + async def on_checkpoint_restore(self, state: dict[str, Any]) -> None: + self.restore_calls += 1 + self.value = state["value"] + + storage = _scope_checkpoint_storage(tmp_path) + mcp_handler = StubMcpHandler(_ok()) + leaves: list[Workflow] = [] + for index in range(2): + factory = WorkflowFactory( + mcp_tool_handler=mcp_handler, + migrate_legacy_request_scopes=index == 0 or mode in {"mixed", "valid"}, + ) + leaf = factory.create_workflow_from_definition(_yaml(_action(require_approval=True))) + await leaf.run({}) + if index == 1 and mode == "mixed": + await leaf.run({}) + requests = await leaf._runner_context.get_pending_request_info_events() + for request in requests.values(): + if mode != "mixed" or index == 0: + request.data.workflow_session_id = None + if index != 1 or mode != "pending": + await leaf._runner_context.send_request_info_response( + request.request_id, _scope_pause_response("InvokeMcpTool") + ) + leaves.append(leaf) + + workflows = list(leaves) + recorders: list[RestoreRecorder] = [] + root = leaves[1] + for level in range(depth): + recorder = RestoreRecorder(f"recorder_{level}") + builder = WorkflowBuilder(start_executor=recorder, checkpoint_storage=storage) + nested = WorkflowExecutor(root, id=f"nested_{level}") + if level == 0: + good = WorkflowExecutor(leaves[0], id="good") + builder.add_fan_out_edges(recorder, [good, nested]) + else: + builder.add_edge(recorder, nested) + root = builder.build() + workflows.append(root) + recorders.append(recorder) + for workflow in workflows: + workflow._runner.state.set("restore_marker", "snapshot") + workflow._runner.state.commit() + checkpoint = await root._runner.build_checkpoint() + checkpoint_before = encode_checkpoint_value(checkpoint.to_dict()) + await storage.save(checkpoint) + + for workflow in workflows: + workflow._runner.state.set("restore_marker", "live") + workflow._runner.state.commit() + workflow._runner.state.set("uncommitted", "keep") + await workflow._runner_context.send_message( + WorkflowMessage(data="live", source_id="live", target_id=next(iter(workflow.executors))) + ) + for recorder in recorders: + recorder.value = "live" + + async def snapshots() -> list[WorkflowCheckpoint]: + return [ + deepcopy( + await workflow._runner_context.build_checkpoint( + workflow.name, workflow.graph_signature_hash, workflow._runner.state, None, 0 + ) + ) + for workflow in workflows + ] + + before = await snapshots() + + async def restore() -> None: + if restore_kind == "embedded": + await root._runner.restore_checkpoint(checkpoint) + else: + await root._runner.restore_from_checkpoint(checkpoint.checkpoint_id) + + if mode in {"missing", "mixed"}: + with pytest.raises(WorkflowCheckpointException) as exc: + await restore() + cause: BaseException = exc.value + while cause.__cause__ is not None: + cause = cause.__cause__ + assert isinstance(cause, DeclarativeWorkflowError) + expected = "migrate_legacy_request_scopes" if mode == "missing" else "different workflow sessions" + assert expected in str(cause) + after = await snapshots() + for old, current in zip(before, after, strict=True): + assert encode_checkpoint_value(current.state) == encode_checkpoint_value(old.state) + assert encode_checkpoint_value(current.messages) == encode_checkpoint_value(old.messages) + assert encode_checkpoint_value(current.pending_request_info_events) == encode_checkpoint_value( + old.pending_request_info_events + ) + assert all(workflow._runner.state.get("uncommitted") == "keep" for workflow in workflows) + assert all(recorder.value == "live" and recorder.restore_calls == 0 for recorder in recorders) + else: + await restore() + assert all(workflow._runner.state.get("restore_marker") == "snapshot" for workflow in workflows) + assert all(recorder.value == "snapshot" and recorder.restore_calls == 1 for recorder in recorders) + assert mcp_handler.call_count == 0 + for index, leaf in enumerate(leaves): + context = leaf._runner_context + assert isinstance(context, InProcRunnerContext) + queued = await context.get_queued_response_info_events() + if mode == "pending" and index == 1: + [pending] = (await leaf._runner_context.get_pending_request_info_events()).values() + assert pending.data.workflow_session_id is None + assert not queued + continue + [request] = queued + scope = f"legacy-request:{request.request_id}" + assert request.data.workflow_session_id == scope + _ = [event async for event in leaf._runner.run_until_convergence()] + assert mcp_handler.invocations[-1].workflow_session_id == scope + assert encode_checkpoint_value(checkpoint.to_dict()) == checkpoint_before + assert mcp_handler.call_count == (2 if mode == "valid" else 1 if mode == "pending" else 0) + + # ---------- Output handling ------------------------------------------------ @@ -400,7 +2491,6 @@ class TestApprovalFlow: @pytest.mark.asyncio async def test_auth_context_binding_survives_workflow_resume(self, restore: bool, change_context: bool) -> None: from agent_framework import InMemoryCheckpointStorage, WorkflowBuilder - from agent_framework_declarative._workflows import ( ActionTrigger, ExternalInputResponse, @@ -445,6 +2535,8 @@ def build(): requests = paused.get_request_info_events() approval = next(event for event in requests if isinstance(event.data, MCPToolApprovalRequest)) change_request = next(event for event in requests if event is not approval) + approval_session_id = approval.data.workflow_session_id + assert approval_session_id selector = "second" if change_context else "first" await workflow.run(responses={change_request.request_id: ExternalInputResponse(user_input=selector)}) assert handler.call_count == 0 @@ -454,7 +2546,14 @@ def build(): assert "context-second" not in repr(checkpoints) assert "context-first" not in repr(approval.data) assert "context-second" not in repr(approval.data) - checkpoint = max(checkpoints, key=lambda item: item.timestamp) + # Wall-clock timestamps can tie on Windows. Select the checkpoint after + # the context change, where only the original approval remains pending. + checkpoint = next( + item + for item in checkpoints + if approval.request_id in item.pending_request_info_events + and change_request.request_id not in item.pending_request_info_events + ) assert checkpoint.state[_HEADER_BINDING_KEY] not in repr(approval.data) if restore: workflow = build() @@ -466,6 +2565,7 @@ def build(): assert not resumed.get_request_info_events() assert handler.last_invocation is not None assert handler.last_invocation.headers == {"Authorization": "context-first"} + assert handler.last_invocation.workflow_session_id == approval_session_id return assert handler.call_count == 0 [replacement] = resumed.get_request_info_events() @@ -478,6 +2578,7 @@ def build(): assert handler.last_invocation is not None assert handler.last_invocation.headers == {"Authorization": "context-second"} assert handler.last_invocation.arguments == {"q": "reviewed"} + assert handler.last_invocation.workflow_session_id == approval_session_id @pytest.mark.asyncio async def test_approval_required_emits_request_and_yields(self, mock_state, mock_context) -> None: # type: ignore[no-untyped-def] @@ -578,6 +2679,7 @@ async def test_approval_response_rejected_assigns_error(self, mock_state, mock_c server_url="https://mcp.example/api", server_label=None, arguments={}, + workflow_session_id="workflow-a", ), ToolApprovalResponse(approved=False, reason="not authorized"), mock_context, diff --git a/skills/agent-framework-builder-python/references/declarative.md b/skills/agent-framework-builder-python/references/declarative.md index ce46be674cf..d46ff79e18f 100644 --- a/skills/agent-framework-builder-python/references/declarative.md +++ b/skills/agent-framework-builder-python/references/declarative.md @@ -24,6 +24,17 @@ maturity. Check current package status and public docstrings before making compa - Allowlist external destinations and register only resources the definition may invoke. - Review the documented HTTP body limits and handler ownership rules; use a custom handler when host policy requires different authentication, retries, SSRF controls, or transport behavior. +- Direct `DefaultMCPToolHandler` calls without a non-empty `workflow_session_id` use a fresh session per invocation; + only explicitly scoped calls without a client provider reuse cached MCP sessions. +- Declarative children wrapped in `WorkflowExecutor` retain MCP session continuity across parent agent turns. + Each binding must own a unique child workflow instance; its MCP scope remains isolated from sibling bindings. + Replaying a parent entry checkpoint also retains child scopes when rebuilding the same workflow graph. + Declarative executors manually composed with Core `WorkflowBuilder` also follow Core's active run and inherited scopes. + Child responses restore the originating parent scope before invoking downstream bindings, including after + checkpoint restoration. Submit responses from different parent scopes in separate runs. + The same origin recovery applies to ordinary Core `ctx.request_info()` pauses before invoking a child. + Queued replies participate in response and cancellation preflight. Wrapped child replies are also validated + before live dispatch; do not combine different parent origins in one message or cancellation batch. ## Validation diff --git a/skills/agent-framework-builder-python/references/workflow-builder.md b/skills/agent-framework-builder-python/references/workflow-builder.md index c121e97c66e..12d01fdff1d 100644 --- a/skills/agent-framework-builder-python/references/workflow-builder.md +++ b/skills/agent-framework-builder-python/references/workflow-builder.md @@ -40,6 +40,11 @@ between stages, final output, intermediate output, session ownership, and approv - Sessions are not automatically interchangeable between participants. - Tool approval and workflow human input are different suspension mechanisms. - Workflow streams include framework events in addition to application output. +- Wrapped child replies evaluate routing after preceding handlers and checkpoint restoration, then validate the + selected children before delivery. A later rejection does not roll back earlier execution or restored state; + recover a partially executed failed run from a known-good checkpoint rather than retrying its remaining queue. +- Accepted reply scopes persist across supersteps and in-flight checkpoint replay until messages drain. A later + response-only run may resume a different originating scope. ## Validation